|
2 | 2 | #include <windows.h> |
3 | 3 | // clang-format on |
4 | 4 | #include <processthreadsapi.h> |
| 5 | +#include <tlhelp32.h> |
5 | 6 | #include <winsock.h> |
6 | 7 |
|
7 | 8 | #include <cassert> |
@@ -141,6 +142,86 @@ int at_fork_in_child(void (*on_fork)()) { |
141 | 142 | return 0; |
142 | 143 | } |
143 | 144 |
|
| 145 | +namespace { |
| 146 | + |
| 147 | +DWORD get_parent_process_id() { |
| 148 | + HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0); |
| 149 | + if (snapshot == INVALID_HANDLE_VALUE) return 0; |
| 150 | + PROCESSENTRY32 pe; |
| 151 | + pe.dwSize = sizeof(pe); |
| 152 | + DWORD own_pid = GetCurrentProcessId(); |
| 153 | + DWORD ppid = 0; |
| 154 | + if (Process32First(snapshot, &pe)) { |
| 155 | + do { |
| 156 | + if (pe.th32ProcessID == own_pid) { |
| 157 | + ppid = pe.th32ParentProcessID; |
| 158 | + break; |
| 159 | + } |
| 160 | + } while (Process32Next(snapshot, &pe)); |
| 161 | + } |
| 162 | + CloseHandle(snapshot); |
| 163 | + return ppid; |
| 164 | +} |
| 165 | + |
| 166 | +} // namespace |
| 167 | + |
| 168 | +SessionCarrier::SessionCarrier(void* handle) : handle_(handle) {} |
| 169 | + |
| 170 | +SessionCarrier::SessionCarrier(SessionCarrier&& rhs) noexcept { |
| 171 | + std::swap(handle_, rhs.handle_); |
| 172 | +} |
| 173 | + |
| 174 | +SessionCarrier& SessionCarrier::operator=(SessionCarrier&& rhs) noexcept { |
| 175 | + std::swap(handle_, rhs.handle_); |
| 176 | + return *this; |
| 177 | +} |
| 178 | + |
| 179 | +SessionCarrier::~SessionCarrier() { |
| 180 | + if (!handle_) return; |
| 181 | + auto* h = static_cast<HANDLE*>(handle_); |
| 182 | + CloseHandle(*h); |
| 183 | + delete h; |
| 184 | +} |
| 185 | + |
| 186 | +Expected<SessionCarrier> SessionCarrier::create(StringView root_session_id) { |
| 187 | + std::string name = |
| 188 | + "Local\\dd-session-ids-" + std::to_string(GetCurrentProcessId()); |
| 189 | + const std::string data{root_session_id}; |
| 190 | + DWORD size = static_cast<DWORD>(data.size()); |
| 191 | + HANDLE h = CreateFileMappingA(INVALID_HANDLE_VALUE, nullptr, PAGE_READWRITE, |
| 192 | + 0, size, name.c_str()); |
| 193 | + if (!h) { |
| 194 | + return Error{Error::Code::OTHER, "CreateFileMapping failed"}; |
| 195 | + } |
| 196 | + void* view = MapViewOfFile(h, FILE_MAP_WRITE, 0, 0, size); |
| 197 | + if (view) { |
| 198 | + std::memcpy(view, data.data(), data.size()); |
| 199 | + UnmapViewOfFile(view); |
| 200 | + } |
| 201 | + return SessionCarrier(new HANDLE{h}); |
| 202 | +} |
| 203 | + |
| 204 | +Optional<std::string> read_inherited_session_id() { |
| 205 | + DWORD ppid = get_parent_process_id(); |
| 206 | + if (!ppid) return nullopt; |
| 207 | + std::string name = "Local\\dd-session-ids-" + std::to_string(ppid); |
| 208 | + HANDLE h = OpenFileMappingA(FILE_MAP_READ, FALSE, name.c_str()); |
| 209 | + if (!h) return nullopt; |
| 210 | + void* view = MapViewOfFile(h, FILE_MAP_READ, 0, 0, 0); |
| 211 | + if (!view) { |
| 212 | + CloseHandle(h); |
| 213 | + return nullopt; |
| 214 | + } |
| 215 | + MEMORY_BASIC_INFORMATION mbi; |
| 216 | + VirtualQuery(view, &mbi, sizeof(mbi)); |
| 217 | + const char* data = static_cast<const char*>(view); |
| 218 | + std::string result(data, strnlen(data, mbi.RegionSize)); |
| 219 | + UnmapViewOfFile(view); |
| 220 | + CloseHandle(h); |
| 221 | + if (result.empty()) return nullopt; |
| 222 | + return result; |
| 223 | +} |
| 224 | + |
144 | 225 | InMemoryFile::InMemoryFile(void* handle) : handle_(handle) {} |
145 | 226 |
|
146 | 227 | InMemoryFile::InMemoryFile(InMemoryFile&& rhs) { |
|
0 commit comments