|
| 1 | +import Foundation |
| 2 | + |
| 3 | +import struct TSCBasic.AbsolutePath |
| 4 | +import struct TSCBasic.ByteString |
| 5 | +import struct TSCBasic.Diagnostic |
| 6 | +import protocol TSCBasic.DiagnosticData |
| 7 | +import class TSCBasic.DiagnosticsEngine |
| 8 | +import protocol TSCBasic.FileSystem |
| 9 | +import protocol TSCBasic.OutputByteStream |
| 10 | +import typealias TSCBasic.ProcessEnvironmentBlock |
| 11 | +import func TSCBasic.getEnvSearchPaths |
| 12 | +import func TSCBasic.lookupExecutablePath |
| 13 | + |
| 14 | +/// Check that the architecture of the toolchain matches the architecture |
| 15 | +/// of the Python installation. |
| 16 | +/// |
| 17 | +/// When installing the x86 toolchain on ARM64 Windows, if the user does not |
| 18 | +/// install an x86 version of Python, they will get acryptic error message |
| 19 | +/// when running lldb (`0xC000007B`). Calling this function before invoking |
| 20 | +/// lldb gives them a warning to help troublshoot the issue. |
| 21 | +/// |
| 22 | +/// - Parameters: |
| 23 | +/// - cwd: The current working directory. |
| 24 | +/// - env: A dictionary of the environment variables and their values. Usually of the parent shell. |
| 25 | +/// - diagnosticsEngine: DiagnosticsEngine instance to use for printing the warning. |
| 26 | +public func checkIfMatchingPythonArch( |
| 27 | + cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, diagnosticsEngine: DiagnosticsEngine |
| 28 | +) throws { |
| 29 | + #if arch(arm64) |
| 30 | + let toolchainArchitecture = COFFBinaryExecutableArchitecture.arm64 |
| 31 | + #elseif arch(x86_64) |
| 32 | + let toolchainArchitecture = COFFBinaryExecutableArchitecture.x64 |
| 33 | + #elseif arch(x86) |
| 34 | + let toolchainArchitecture = COFFBinaryExecutableArchitecture.x86 |
| 35 | + #else |
| 36 | + return |
| 37 | + #endif |
| 38 | + let pythonArchitecture = COFFBinaryExecutableArchitecture.readWindowsExecutableArchitecture( |
| 39 | + cwd: cwd, envBlock: envBlock, filename: "python.exe") |
| 40 | + |
| 41 | + if toolchainArchitecture != pythonArchitecture { |
| 42 | + diagnosticsEngine.emit( |
| 43 | + .warning( |
| 44 | + """ |
| 45 | + There is an architecture mismatch between the installed toolchain and the resolved Python's architecture: |
| 46 | + Toolchain: \(toolchainArchitecture) |
| 47 | + Python: \(pythonArchitecture) |
| 48 | + """)) |
| 49 | + } |
| 50 | +} |
| 51 | + |
| 52 | +/// Some of the architectures that can be stored in a COFF header. |
| 53 | +enum COFFBinaryExecutableArchitecture: String { |
| 54 | + case x86 = "X86" |
| 55 | + case x64 = "X64" |
| 56 | + case arm64 = "ARM64" |
| 57 | + case unknown = "Unknown" |
| 58 | + |
| 59 | + static func fromPEMachineByte(machine: UInt16) -> Self { |
| 60 | + // https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#machine-types |
| 61 | + switch machine { |
| 62 | + case 0x014c: return .x86 |
| 63 | + case 0x8664: return .x64 |
| 64 | + case 0xAA64: return .arm64 |
| 65 | + default: return .unknown |
| 66 | + } |
| 67 | + } |
| 68 | + |
| 69 | + /// Resolves the filename from the `Path` environment variable and read its COFF header to determine the architecture |
| 70 | + /// of the binary. |
| 71 | + /// |
| 72 | + /// - Parameters: |
| 73 | + /// - cwd: The current working directory. |
| 74 | + /// - env: A dictionary of the environment variables and their values. Usually of the parent shell. |
| 75 | + /// - filename: The name of the file we are resolving the architecture of. |
| 76 | + /// - Returns: The architecture of the file which was found in the `Path`. |
| 77 | + static func readWindowsExecutableArchitecture( |
| 78 | + cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, filename: String |
| 79 | + ) -> Self { |
| 80 | + let searchPaths = getEnvSearchPaths(pathString: envBlock["Path"], currentWorkingDirectory: cwd) |
| 81 | + guard |
| 82 | + let filePath = lookupExecutablePath( |
| 83 | + filename: filename, currentWorkingDirectory: cwd, searchPaths: searchPaths) |
| 84 | + else { |
| 85 | + return .unknown |
| 86 | + } |
| 87 | + guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else { |
| 88 | + return .unknown |
| 89 | + } |
| 90 | + |
| 91 | + defer { fileHandle.closeFile() } |
| 92 | + |
| 93 | + // Infering the architecture of a Windows executable from its COFF header involves the following: |
| 94 | + // 1. Get the COFF header offset from the pointer located at the 0x3C offset (4 bytes long). |
| 95 | + // 2. Jump to that offset and read the next 6 bytes. |
| 96 | + // 3. The first 4 are the signature which should be equal to 0x50450000. |
| 97 | + // 4. The last 2 are the machine architecture which can be infered from the value we get. |
| 98 | + // |
| 99 | + // The link below provides a visualization of the COFF header and the process to get to it. |
| 100 | + // https://upload.wikimedia.org/wikipedia/commons/1/1b/Portable_Executable_32_bit_Structure_in_SVG_fixed.svg |
| 101 | + fileHandle.seek(toFileOffset: 0x3C) |
| 102 | + guard let offsetPointer = try? fileHandle.read(upToCount: 4), |
| 103 | + offsetPointer.count == 4 |
| 104 | + else { |
| 105 | + return .unknown |
| 106 | + } |
| 107 | + |
| 108 | + let peHeaderOffset = offsetPointer.withUnsafeBytes { $0.load(as: UInt32.self) } |
| 109 | + |
| 110 | + fileHandle.seek(toFileOffset: UInt64(peHeaderOffset)) |
| 111 | + guard let coffHeader = try? fileHandle.read(upToCount: 6), coffHeader.count == 6 else { |
| 112 | + return .unknown |
| 113 | + } |
| 114 | + |
| 115 | + let signature = coffHeader.prefix(4) |
| 116 | + let machineBytes = coffHeader.suffix(2) |
| 117 | + |
| 118 | + guard signature == Data([0x50, 0x45, 0x00, 0x00]) else { |
| 119 | + return .unknown |
| 120 | + } |
| 121 | + |
| 122 | + return .fromPEMachineByte(machine: machineBytes.withUnsafeBytes { $0.load(as: UInt16.self) }) |
| 123 | + } |
| 124 | +} |
0 commit comments