|
| 1 | +//===----------------------------------------------------------------------===// |
| 2 | +// |
| 3 | +// This source file is part of the Swift open source project |
| 4 | +// |
| 5 | +// Copyright (c) 2025 Apple Inc. and the Swift project authors |
| 6 | +// Licensed under Apache License v2.0 with Runtime Library Exception |
| 7 | +// |
| 8 | +// See http://swift.org/LICENSE.txt for license information |
| 9 | +// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors |
| 10 | +// |
| 11 | +//===----------------------------------------------------------------------===// |
| 12 | + |
| 13 | +import WinSDK |
| 14 | +import struct Foundation.Data |
| 15 | +import class Foundation.FileHandle |
| 16 | +import struct Foundation.URL |
| 17 | + |
| 18 | +// https://learn.microsoft.com/en-us/windows/win32/api/winternl/nf-winternl-ntqueryinformationprocess |
| 19 | +// Also see winternl.h in the Windows SDK. |
| 20 | + |
| 21 | +fileprivate struct PROCESS_BASIC_INFORMATION { |
| 22 | + var ExitStatus: NTSTATUS = 0 |
| 23 | + var PebBaseAddress: ULONG_PTR = 0 |
| 24 | + var AffinityMask: ULONG_PTR = 0 |
| 25 | + var BasePriority: LONG = 0 |
| 26 | + var UniqueProcessId: ULONG_PTR = 0 |
| 27 | + var InheritedFromUniqueProcessId: ULONG_PTR = 0 |
| 28 | +} |
| 29 | + |
| 30 | +fileprivate typealias NtQueryInformationProcessFunction = @convention(c) (_ ProcessHandle: HANDLE, _ ProcessInformationClass: CInt, _ ProcessInformation: PVOID, _ ProcessInformationLength: ULONG, _ ReturnLength: PULONG) -> NTSTATUS |
| 31 | + |
| 32 | +fileprivate struct _Win32Error: Error { |
| 33 | + let functionName: String |
| 34 | + let error: DWORD |
| 35 | +} |
| 36 | + |
| 37 | +fileprivate nonisolated var KF_FLAG_DEFAULT: DWORD { |
| 38 | + DWORD(WinSDK.KF_FLAG_DEFAULT.rawValue) |
| 39 | +} |
| 40 | + |
| 41 | +fileprivate func SUCCEEDED(_ hr: HRESULT) -> Bool { |
| 42 | + hr >= 0 |
| 43 | +} |
| 44 | + |
| 45 | +fileprivate func _url(for id: KNOWNFOLDERID) throws -> URL { |
| 46 | + var pszPath: PWSTR? |
| 47 | + let hr: HRESULT = withUnsafePointer(to: id) { id in |
| 48 | + SHGetKnownFolderPath(id, KF_FLAG_DEFAULT, nil, &pszPath) |
| 49 | + } |
| 50 | + guard SUCCEEDED(hr) else { throw _Win32Error(functionName: "SHGetKnownFolderPath", error: GetLastError()) } |
| 51 | + defer { CoTaskMemFree(pszPath) } |
| 52 | + return URL(filePath: String(decodingCString: pszPath!, as: UTF16.self), directoryHint: .isDirectory) |
| 53 | +} |
| 54 | + |
| 55 | +extension String { |
| 56 | + fileprivate func withLPWSTR<T>(_ body: (UnsafeMutablePointer<WCHAR>) throws -> T) rethrows -> T { |
| 57 | + try withUnsafeTemporaryAllocation(of: WCHAR.self, capacity: self.utf16.count + 1, { outBuffer in |
| 58 | + try self.withCString(encodedAs: UTF16.self) { inBuffer in |
| 59 | + outBuffer.baseAddress!.initialize(from: inBuffer, count: self.utf16.count) |
| 60 | + outBuffer[outBuffer.count - 1] = 0 |
| 61 | + return try body(outBuffer.baseAddress!) |
| 62 | + } |
| 63 | + }) |
| 64 | + } |
| 65 | +} |
| 66 | + |
| 67 | +func withDebugEventLoop(_ hProcess: HANDLE, _ handle: (_ event: String) throws -> ()) throws { |
| 68 | + let ntdllURL = try _url(for: FOLDERID_System).appendingPathComponent("ntdll.dll") |
| 69 | + guard let ntdll = ntdllURL.withUnsafeFileSystemRepresentation({ s in s.map({ String(cString: $0) }) ?? String() }).withLPWSTR({ LoadLibraryW($0) }) else { |
| 70 | + throw _Win32Error(functionName: "LoadLibraryW", error: GetLastError()) |
| 71 | + } |
| 72 | + |
| 73 | + defer { |
| 74 | + _ = FreeLibrary(ntdll) |
| 75 | + } |
| 76 | + |
| 77 | + guard let ntQueryInformationProc = GetProcAddress(ntdll, "NtQueryInformationProcess") else { |
| 78 | + throw _Win32Error(functionName: "GetProcAddress", error: GetLastError()) |
| 79 | + } |
| 80 | + |
| 81 | + var processBasicInformation = PROCESS_BASIC_INFORMATION() |
| 82 | + var len: ULONG = 0 |
| 83 | + |
| 84 | + let processBasicInformationSize = MemoryLayout.size(ofValue: processBasicInformation) |
| 85 | +#if arch(x86_64) || arch(arm64) |
| 86 | + precondition(processBasicInformationSize == 48) |
| 87 | +#else |
| 88 | + precondition(processBasicInformationSize == 24) |
| 89 | +#endif |
| 90 | + guard unsafeBitCast(ntQueryInformationProc, to: NtQueryInformationProcessFunction.self)(hProcess, 0, &processBasicInformation, ULONG(processBasicInformationSize), &len) == 0 else { |
| 91 | + throw _Win32Error(functionName: "NtQueryInformationProcess", error: GetLastError()) |
| 92 | + } |
| 93 | + |
| 94 | + let peb = Int(processBasicInformation.PebBaseAddress) |
| 95 | + |
| 96 | + var gflags: ULONG = 0 |
| 97 | + var actual: SIZE_T = 0 |
| 98 | + guard ReadProcessMemory(hProcess, UnsafeMutableRawPointer(bitPattern: peb + 0xBC), &gflags, SIZE_T(MemoryLayout.size(ofValue: gflags)), &actual) else { |
| 99 | + throw _Win32Error(functionName: "ReadProcessMemory", error: GetLastError()) |
| 100 | + } |
| 101 | + |
| 102 | + gflags |= 0x2 |
| 103 | + guard WriteProcessMemory(hProcess, UnsafeMutableRawPointer(bitPattern: peb + 0xBC), &gflags, SIZE_T(MemoryLayout.size(ofValue: gflags)), &actual) else { |
| 104 | + throw _Win32Error(functionName: "WriteProcessMemory", error: GetLastError()) |
| 105 | + } |
| 106 | + |
| 107 | + func debugOutputString(_ hProcess: HANDLE, _ dbgEvent: inout DEBUG_EVENT) throws -> String { |
| 108 | + let size = SIZE_T(dbgEvent.u.DebugString.nDebugStringLength) |
| 109 | + return try withUnsafeTemporaryAllocation(of: UInt8.self, capacity: Int(size) + 2) { buffer in |
| 110 | + guard ReadProcessMemory(hProcess, dbgEvent.u.DebugString.lpDebugStringData, buffer.baseAddress, size, nil) else { |
| 111 | + throw _Win32Error(functionName: "ReadProcessMemory", error: GetLastError()) |
| 112 | + } |
| 113 | + |
| 114 | + buffer[Int(size)] = 0 |
| 115 | + buffer[Int(size + 1)] = 0 |
| 116 | + |
| 117 | + if dbgEvent.u.DebugString.fUnicode != 0 { |
| 118 | + return buffer.withMemoryRebound(to: UInt16.self) { String(decoding: $0, as: UTF16.self) } |
| 119 | + } else { |
| 120 | + return try withUnsafeTemporaryAllocation(of: UInt16.self, capacity: Int(size)) { wideBuffer in |
| 121 | + if MultiByteToWideChar(UINT(CP_ACP), 0, buffer.baseAddress, Int32(size), wideBuffer.baseAddress, Int32(size)) == 0 { |
| 122 | + throw _Win32Error(functionName: "MultiByteToWideChar", error: GetLastError()) |
| 123 | + } |
| 124 | + return String(decoding: wideBuffer, as: UTF16.self) |
| 125 | + } |
| 126 | + } |
| 127 | + } |
| 128 | + } |
| 129 | + |
| 130 | + func _WaitForDebugEventEx() throws -> DEBUG_EVENT { |
| 131 | + // WARNING: Only the thread that created the process being debugged can call WaitForDebugEventEx. |
| 132 | + var dbgEvent = DEBUG_EVENT() |
| 133 | + guard WaitForDebugEventEx(&dbgEvent, INFINITE) else { |
| 134 | + // WaitForDebugEventEx will fail if dwCreationFlags did not contain DEBUG_ONLY_THIS_PROCESS |
| 135 | + throw _Win32Error(functionName: "WaitForDebugEventEx", error: GetLastError()) |
| 136 | + } |
| 137 | + return dbgEvent |
| 138 | + } |
| 139 | + |
| 140 | + func runDebugEventLoop() throws { |
| 141 | + do { |
| 142 | + while true { |
| 143 | + var dbgEvent = try _WaitForDebugEventEx() |
| 144 | + if dbgEvent.dwProcessId == GetProcessId(hProcess) { |
| 145 | + switch dbgEvent.dwDebugEventCode { |
| 146 | + case DWORD(OUTPUT_DEBUG_STRING_EVENT): |
| 147 | + try handle(debugOutputString(hProcess, &dbgEvent)) |
| 148 | + case DWORD(EXIT_PROCESS_DEBUG_EVENT): |
| 149 | + return // done! |
| 150 | + default: |
| 151 | + break |
| 152 | + } |
| 153 | + } |
| 154 | + |
| 155 | + guard ContinueDebugEvent(dbgEvent.dwProcessId, dbgEvent.dwThreadId, DBG_EXCEPTION_NOT_HANDLED) else { |
| 156 | + throw _Win32Error(functionName: "WaitForDebugEventEx", error: GetLastError()) |
| 157 | + } |
| 158 | + } |
| 159 | + } catch { |
| 160 | + throw error |
| 161 | + } |
| 162 | + } |
| 163 | + |
| 164 | + try runDebugEventLoop() |
| 165 | +} |
| 166 | + |
| 167 | +func createProcessTrampoline(_ commandLine: String) throws -> Int32 { |
| 168 | + var processInformation = PROCESS_INFORMATION() |
| 169 | + guard commandLine.withLPWSTR({ wCommandLine in |
| 170 | + var startupInfo = STARTUPINFOW() |
| 171 | + startupInfo.cb = DWORD(MemoryLayout.size(ofValue: startupInfo)) |
| 172 | + return CreateProcessW( |
| 173 | + nil, |
| 174 | + wCommandLine, |
| 175 | + nil, |
| 176 | + nil, |
| 177 | + false, |
| 178 | + DWORD(DEBUG_ONLY_THIS_PROCESS), |
| 179 | + nil, |
| 180 | + nil, |
| 181 | + &startupInfo, |
| 182 | + &processInformation, |
| 183 | + ) |
| 184 | + }) else { |
| 185 | + throw _Win32Error(functionName: "CreateProcessW", error: GetLastError()) |
| 186 | + } |
| 187 | + defer { |
| 188 | + _ = CloseHandle(processInformation.hThread) |
| 189 | + _ = CloseHandle(processInformation.hProcess) |
| 190 | + } |
| 191 | + var missingDLLs: [String] = [] |
| 192 | + try withDebugEventLoop(processInformation.hProcess) { message in |
| 193 | + if let match = try #/ ERROR: Unable to load DLL: "(?<moduleName>.*?)",/#.firstMatch(in: message) { |
| 194 | + missingDLLs.append(String(match.output.moduleName)) |
| 195 | + } |
| 196 | + } |
| 197 | + // Don't need to call WaitForSingleObject because the process will have exited after withDebugEventLoop is called |
| 198 | + var exitCode: DWORD = .max |
| 199 | + guard GetExitCodeProcess(processInformation.hProcess, &exitCode) else { |
| 200 | + throw _Win32Error(functionName: "GetExitCodeProcess", error: GetLastError()) |
| 201 | + } |
| 202 | + if exitCode == STATUS_DLL_NOT_FOUND { |
| 203 | + let stderr = FileHandle.standardError |
| 204 | + for missingDLL in missingDLLs { |
| 205 | + try stderr.write(contentsOf: Data("This application has failed to start because \(missingDLL) was not found.\r\n".utf8)) |
| 206 | + } |
| 207 | + } |
| 208 | + return Int32(bitPattern: exitCode) |
| 209 | +} |
| 210 | + |
| 211 | +func main() -> Int32 { |
| 212 | + do { |
| 213 | + var commandLine = String(decodingCString: GetCommandLineW(), as: UTF16.self) |
| 214 | + |
| 215 | + // FIXME: This could probably be more robust |
| 216 | + if commandLine.first == "\"" { |
| 217 | + commandLine = String(commandLine.dropFirst()) |
| 218 | + if let index = commandLine.firstIndex(of: "\"") { |
| 219 | + commandLine = String(commandLine.dropFirst(commandLine.distance(from: commandLine.startIndex, to: index) + 2)) |
| 220 | + } |
| 221 | + } else if let index = commandLine.firstIndex(of: " ") { |
| 222 | + commandLine = String(commandLine.dropFirst(commandLine.distance(from: commandLine.startIndex, to: index) + 1)) |
| 223 | + } else { |
| 224 | + commandLine = "" |
| 225 | + } |
| 226 | + |
| 227 | + return try createProcessTrampoline(commandLine) |
| 228 | + } catch { |
| 229 | + let stderr = FileHandle.standardError |
| 230 | + try? stderr.write(contentsOf: Data("\(error)\r\n".utf8)) |
| 231 | + return EXIT_FAILURE |
| 232 | + } |
| 233 | +} |
| 234 | + |
| 235 | +exit(main()) |
0 commit comments