use COFF header to infer python.exe architecture

Author: charles-zablit

Package.resolved

@@ -17,6 +17,7 @@ import class TSCBasic.DiagnosticsEngine
 import class TSCBasic.UnknownLocation
 import enum TSCBasic.ProcessEnv
 import func TSCBasic.withTemporaryDirectory
+import func TSCBasic.lookupExecutablePath
 import protocol TSCBasic.DiagnosticData
 import protocol TSCBasic.FileSystem
 import protocol TSCBasic.OutputByteStream
@@ -869,7 +870,7 @@ public struct Driver {
     #if os(Windows)
       if args[1] == "-repl" {  // a `-` gets automatically added to `swift repl`.
         try Driver.checkIfMatchingPythonArch(
-          executor: executor, env: env, diagnosticsEngine: diagnosticsEngine)
+          cwd: ProcessEnv.cwd, env: env, diagnosticsEngine: diagnosticsEngine)
       }
     #endif
 
@@ -1452,24 +1453,20 @@ extension Driver {
   /// when running lldb (`0xC000007B`). Calling this function before invoking
   /// lldb gives them a warning to help troublshoot the issue.
   public static func checkIfMatchingPythonArch(
-    executor: DriverExecutor, env: [String: String], diagnosticsEngine: DiagnosticsEngine
+    cwd: AbsolutePath?, env: [String: String], diagnosticsEngine: DiagnosticsEngine
   ) throws {
     #if arch(arm64)
-      let toolchainArchitecture = "arm64"
+      let toolchainArchitecture = ExecutableArchitecture.arm64
     #elseif arch(x86_64)
-      let toolchainArchitecture = "amd64"
+      let toolchainArchitecture = ExecutableArchitecture.x64
     #elseif arch(x86)
-      let toolchainArchitecture = "x86"
+      let toolchainArchitecture = ExecutableArchitecture.x86
     #else
       return;
     #endif
-    let pythonArchitecture = try executor.checkNonZeroExit(
-      args: "python.exe", "-c", "import platform; print(platform.machine())",
-      environment: env
-    ).trimmingCharacters(in: .whitespacesAndNewlines)
-      .lowercased()
+    let pythonArchitecture = readWindowsExecutableArchitecture(cwd: cwd, env: env)
 
-    if !pythonArchitecture.contains(toolchainArchitecture) {
+    if toolchainArchitecture != pythonArchitecture {
       diagnosticsEngine.emit(
         .warning(
           """
@@ -1479,6 +1476,68 @@ extension Driver {
           """))
     }
   }
+
+  enum ExecutableArchitecture: String {
+    case x86 = "X86"
+    case x64 = "X64"
+    case arm64 = "ARM64"
+    case unknown = "Unknown"
+
+    static func fromPEMachineByte(machine: UInt16) -> Self {
+      // https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#machine-types
+      switch machine {
+      case 0x014c: return .x86
+      case 0x8664: return .x64
+      case 0xAA64: return .arm64
+      default: return .unknown
+      }
+    }
+  }
+
+  static func readWindowsExecutableArchitecture(cwd: AbsolutePath?, env: [String: String])
+    -> ExecutableArchitecture
+  {
+    guard let pathEnv = env["Path"]?.split(separator: ";") else {
+      return .unknown
+    }
+    let searchPaths = pathEnv.compactMap { path in
+      try? AbsolutePath(validating: String(path))
+    }
+    guard let filePath = lookupExecutablePath(
+        filename: "python.exe", currentWorkingDirectory: cwd, searchPaths: searchPaths)
+    else {
+      return .unknown
+    }
+    guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else {
+      return .unknown
+    }
+
+    defer { fileHandle.closeFile() }
+
+    // https://upload.wikimedia.org/wikipedia/commons/1/1b/Portable_Executable_32_bit_Structure_in_SVG_fixed.svg
+    fileHandle.seek(toFileOffset: 0x3C)
+    guard let offsetPointer = try? fileHandle.read(upToCount: 4),
+      offsetPointer.count == 4
+    else {
+      return .unknown
+    }
+
+    let peHeaderOffset = offsetPointer.withUnsafeBytes { $0.load(as: UInt32.self) }
+
+    fileHandle.seek(toFileOffset: UInt64(peHeaderOffset))
+    guard let coffHeader = try? fileHandle.read(upToCount: 6), coffHeader.count == 6 else {
+      return .unknown
+    }
+
+    let signature = coffHeader.prefix(4)
+    let machineBytes = coffHeader.suffix(2)
+
+    guard signature == Data([0x50, 0x45, 0x00, 0x00]) else {
+      return .unknown
+    }
+
+    return .fromPEMachineByte(machine: machineBytes.withUnsafeBytes { $0.load(as: UInt16.self) })
+  }
 }
 
 extension Driver {