Skip to content

Commit daac449

Browse files
charles-zablitcharles-zablit
committed
[darwin] add a check for a matching python architecture
1 parent 279c0fc commit daac449

File tree

2 files changed

+140
-77
lines changed

2 files changed

+140
-77
lines changed

‎Sources/SwiftDriver/Driver/Driver.swift

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -873,12 +873,10 @@ public struct Driver {
873873
throw Error.subcommandPassedToDriver
874874
}
875875

876-
#if os(Windows)
877-
if case .normal(isRepl: true) = invocationMode {
878-
checkIfMatchingPythonArch(
879-
cwd: ProcessEnv.cwd, envBlock: envBlock, diagnosticsEngine: diagnosticsEngine)
880-
}
881-
#endif
876+
if case .normal(isRepl: true) = invocationMode {
877+
checkIfMatchingPythonArch(
878+
cwd: ProcessEnv.cwd, envBlock: envBlock, diagnosticsEngine: diagnosticsEngine)
879+
}
882880

883881
var args = args
884882
if let additional = env["ADDITIONAL_SWIFT_DRIVER_FLAGS"] {

‎Sources/SwiftDriver/Utilities/PythonArchitecture.swift

Lines changed: 136 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -15,45 +15,58 @@ import func TSCBasic.lookupExecutablePath
1515
/// of the Python installation.
1616
///
1717
/// 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
18+
/// install an x86 version of Python, they will get a cryptic error message
1919
/// when running lldb (`0xC000007B`). Calling this function before invoking
20-
/// lldb gives them a warning to help troublshoot the issue.
20+
/// lldb gives them a warning to help troubleshoot the issue.
2121
///
2222
/// - Parameters:
2323
/// - cwd: The current working directory.
24-
/// - env: A dictionary of the environment variables and their values. Usually of the parent shell.
24+
/// - env: The parent shell's ProcessEnvironmentBlock.
2525
/// - diagnosticsEngine: DiagnosticsEngine instance to use for printing the warning.
2626
public func checkIfMatchingPythonArch(
2727
cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, diagnosticsEngine: DiagnosticsEngine
2828
) {
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
29+
#if os(Windows) || os(macOS)
30+
#if arch(arm64)
31+
let toolchainArchitecture = ExecutableArchitecture.arm64
32+
#elseif arch(x86_64)
33+
let toolchainArchitecture = ExecutableArchitecture.x64
34+
#elseif arch(x86)
35+
let toolchainArchitecture = ExecutableArchitecture.x86
36+
#else
37+
return
38+
#endif
39+
40+
#if os(Windows)
41+
let pythonArchitecture = ExecutableArchitecture.readWindowsExecutableArchitecture(
42+
cwd: cwd, envBlock: envBlock, filename: "python.exe")
43+
#elseif os(macOS)
44+
let pythonArchitecture = ExecutableArchitecture.readDarwinExecutableArchitecture(
45+
cwd: cwd, envBlock: envBlock, filename: "python3")
46+
#endif
47+
48+
if pythonArchitecture == .universal {
49+
return
50+
}
51+
52+
if toolchainArchitecture != pythonArchitecture {
53+
diagnosticsEngine.emit(
54+
.warning(
55+
"""
56+
There is an architecture mismatch between the installed toolchain and the resolved Python's architecture:
57+
Toolchain: \(toolchainArchitecture)
58+
Python: \(pythonArchitecture)
59+
"""))
60+
}
3761
#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-
}
5062
}
5163

5264
/// Some of the architectures that can be stored in a COFF header.
53-
enum COFFBinaryExecutableArchitecture: String {
65+
enum ExecutableArchitecture: String {
5466
case x86 = "X86"
5567
case x64 = "X64"
5668
case arm64 = "ARM64"
69+
case universal = "Universal"
5770
case unknown = "Unknown"
5871

5972
static func fromPEMachineByte(machine: UInt16) -> Self {
@@ -66,59 +79,111 @@ enum COFFBinaryExecutableArchitecture: String {
6679
}
6780
}
6881

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
82+
static func fromMachoCPUType(cpuType: Int32) -> Self {
83+
// https://en.wikipedia.org/wiki/Mach-O
84+
switch cpuType {
85+
case 0x0100_0007: return .x86
86+
case 0x0100_000c: return .arm64
87+
default: return .unknown
8988
}
89+
}
9090

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-
}
91+
#if os(Windows)
92+
/// Resolves the filename from the `Path` environment variable and read its COFF header to determine the architecture
93+
/// of the binary.
94+
///
95+
/// - Parameters:
96+
/// - cwd: The current working directory.
97+
/// - env: A dictionary of the environment variables and their values. Usually of the parent shell.
98+
/// - filename: The name of the file we are resolving the architecture of.
99+
/// - Returns: The architecture of the file which was found in the `Path`.
100+
static func readWindowsExecutableArchitecture(
101+
cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, filename: String
102+
) -> Self {
103+
let searchPaths = getEnvSearchPaths(
104+
pathString: envBlock["Path"], currentWorkingDirectory: cwd)
105+
guard
106+
let filePath = lookupExecutablePath(
107+
filename: filename, currentWorkingDirectory: cwd, searchPaths: searchPaths)
108+
else {
109+
return .unknown
110+
}
111+
guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else {
112+
return .unknown
113+
}
107114

108-
let peHeaderOffset = offsetPointer.withUnsafeBytes { $0.load(as: UInt32.self) }
115+
defer { fileHandle.closeFile() }
109116

110-
fileHandle.seek(toFileOffset: UInt64(peHeaderOffset))
111-
guard let coffHeader = try? fileHandle.read(upToCount: 6), coffHeader.count == 6 else {
112-
return .unknown
113-
}
117+
// Infering the architecture of a Windows executable from its COFF header involves the following:
118+
// 1. Get the COFF header offset from the pointer located at the 0x3C offset (4 bytes long).
119+
// 2. Jump to that offset and read the next 6 bytes.
120+
// 3. The first 4 are the signature which should be equal to 0x50450000.
121+
// 4. The last 2 are the machine architecture which can be infered from the value we get.
122+
//
123+
// The link below provides a visualization of the COFF header and the process to get to it.
124+
// https://upload.wikimedia.org/wikipedia/commons/1/1b/Portable_Executable_32_bit_Structure_in_SVG_fixed.svg
125+
fileHandle.seek(toFileOffset: 0x3C)
126+
guard let offsetPointer = try? fileHandle.read(upToCount: 4),
127+
offsetPointer.count == 4
128+
else {
129+
return .unknown
130+
}
131+
132+
let peHeaderOffset = offsetPointer.withUnsafeBytes { $0.load(as: UInt32.self) }
114133

115-
let signature = coffHeader.prefix(4)
116-
let machineBytes = coffHeader.suffix(2)
134+
fileHandle.seek(toFileOffset: UInt64(peHeaderOffset))
135+
guard let coffHeader = try? fileHandle.read(upToCount: 6), coffHeader.count == 6 else {
136+
return .unknown
137+
}
117138

118-
guard signature == Data([0x50, 0x45, 0x00, 0x00]) else {
119-
return .unknown
139+
let signature = coffHeader.prefix(4)
140+
let machineBytes = coffHeader.suffix(2)
141+
142+
guard signature == Data([0x50, 0x45, 0x00, 0x00]) else {
143+
return .unknown
144+
}
145+
146+
return .fromPEMachineByte(machine: machineBytes.withUnsafeBytes { $0.load(as: UInt16.self) })
120147
}
148+
#endif
121149

122-
return .fromPEMachineByte(machine: machineBytes.withUnsafeBytes { $0.load(as: UInt16.self) })
123-
}
150+
#if os(macOS)
151+
static func readDarwinExecutableArchitecture(
152+
cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, filename: String
153+
) -> Self {
154+
let magicNumber: UInt32 = 0xcafe_babe
155+
156+
let searchPaths = getEnvSearchPaths(
157+
pathString: envBlock["PATH"], currentWorkingDirectory: cwd)
158+
guard
159+
let filePath = lookupExecutablePath(
160+
filename: filename, currentWorkingDirectory: cwd, searchPaths: searchPaths)
161+
else {
162+
return .unknown
163+
}
164+
guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else {
165+
return .unknown
166+
}
167+
168+
defer {
169+
try? fileHandle.close()
170+
}
171+
172+
// The first 4 bytes of a Mach-O header contain the magic number. We use it to determine if the binary is
173+
// universal.
174+
// https://github.com/apple/darwin-xnu/blob/main/EXTERNAL_HEADERS/mach-o/loader.h
175+
let magicData = fileHandle.readData(ofLength: 4)
176+
let magic = magicData.withUnsafeBytes { $0.load(as: UInt32.self).bigEndian }
177+
178+
if magic == magicNumber {
179+
return .universal
180+
}
181+
182+
// If the binary is not universal, the next 4 bytes contain the CPU type.
183+
fileHandle.seek(toFileOffset: 4)
184+
let cpuTypeData = fileHandle.readData(ofLength: 4)
185+
let cpuType = cpuTypeData.withUnsafeBytes { $0.load(as: Int32.self) }
186+
return Self.fromMachoCPUType(cpuType: cpuType)
187+
}
188+
#endif
124189
}

0 commit comments

Comments
 (0)