[darwin] add a check for a matching python architecture

charles-zablit · charles-zablit · commit b8a065b0027d · 2025-09-11T12:22:14.000+01:00
diff --git a/Sources/SwiftDriver/Driver/Driver.swift b/Sources/SwiftDriver/Driver/Driver.swift
@@ -875,17 +875,22 @@ public struct Driver {
 
     self.executor = executor
 
-    let invocationMode = try Self.invocationRunMode(forArgs: args).mode
-    if case .subcommand = invocationMode {
-      throw Error.subcommandPassedToDriver
+    if args.count > 1 && args[1] == "-repl" {
+      #if arch(arm64)
+        checkIfMatchingPythonArch(
+          cwd: ProcessEnv.cwd, envBlock: envBlock, toolchainArchitecture: .arm64, diagnosticsEngine: diagnosticsEngine)
+      #elseif arch(x86_64)
+        checkIfMatchingPythonArch(
+          cwd: ProcessEnv.cwd, envBlock: envBlock, toolchainArchitecture: .x64, diagnosticsEngine: diagnosticsEngine)
+      #elseif arch(x86)
+        checkIfMatchingPythonArch(
+          cwd: ProcessEnv.cwd, envBlock: envBlock, toolchainArchitecture: .x86, diagnosticsEngine: diagnosticsEngine)
+      #endif
     }
 
-    #if os(Windows)
-      if case .normal(isRepl: true) = invocationMode {
-        checkIfMatchingPythonArch(
-          cwd: ProcessEnv.cwd, envBlock: envBlock, diagnosticsEngine: diagnosticsEngine)
-      }
-    #endif
+    if case .subcommand = try Self.invocationRunMode(forArgs: args).mode {
+      throw Error.subcommandPassedToDriver
+    }
 
     var args = args
     if let additional = env["ADDITIONAL_SWIFT_DRIVER_FLAGS"] {
diff --git a/Sources/SwiftDriver/Utilities/PythonArchitecture.swift b/Sources/SwiftDriver/Utilities/PythonArchitecture.swift
@@ -1,3 +1,14 @@
+//===--------------- Driver.swift - Swift Driver --------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
 import Foundation
 
 import struct TSCBasic.AbsolutePath
@@ -11,114 +22,187 @@ import typealias TSCBasic.ProcessEnvironmentBlock
 import func TSCBasic.getEnvSearchPaths
 import func TSCBasic.lookupExecutablePath
 
+#if os(Windows)
+  import WinSDK
+#endif
+
 /// Check that the architecture of the toolchain matches the architecture
 /// of the Python installation.
 ///
 /// When installing the x86 toolchain on ARM64 Windows, if the user does not
-/// install an x86 version of Python, they will get acryptic error message
+/// install an x86 version of Python, they will get a cryptic error message
 /// when running lldb (`0xC000007B`). Calling this function before invoking
-/// lldb gives them a warning to help troublshoot the issue.
+/// lldb gives them a warning to help troubleshoot the issue.
 ///
 /// - Parameters:
 ///   - cwd: The current working directory.
-///   - env: A dictionary of the environment variables and their values. Usually of the parent shell.
+///   - env: The parent shell's ProcessEnvironmentBlock.
 ///   - diagnosticsEngine: DiagnosticsEngine instance to use for printing the warning.
 public func checkIfMatchingPythonArch(
-  cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, diagnosticsEngine: DiagnosticsEngine
+  cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock,
+  toolchainArchitecture: ExecutableArchitecture, diagnosticsEngine: DiagnosticsEngine
 ) {
-  #if arch(arm64)
-    let toolchainArchitecture = COFFBinaryExecutableArchitecture.arm64
-  #elseif arch(x86_64)
-    let toolchainArchitecture = COFFBinaryExecutableArchitecture.x64
-  #elseif arch(x86)
-    let toolchainArchitecture = COFFBinaryExecutableArchitecture.x86
-  #else
-    return
+  #if os(Windows) || os(macOS)
+    #if os(Windows)
+      let pythonArchitecture = Process.readWindowsExecutableArchitecture(
+        cwd: cwd, envBlock: envBlock, filename: "python.exe")
+    #elseif os(macOS)
+      let pythonArchitecture = Process.readDarwinExecutableArchitecture(
+        cwd: cwd, envBlock: envBlock, filename: "python3")
+      if pythonArchitecture == .universal {
+        return
+      }
+    #endif
+
+    guard toolchainArchitecture == pythonArchitecture else {
+      diagnosticsEngine.emit(
+        .warning(
+          """
+          There is an architecture mismatch between the installed toolchain and the resolved Python's architecture:
+          Toolchain: \(toolchainArchitecture)
+          Python: \(pythonArchitecture)
+          """))
+      return
+    }
   #endif
-  let pythonArchitecture = COFFBinaryExecutableArchitecture.readWindowsExecutableArchitecture(
-    cwd: cwd, envBlock: envBlock, filename: "python.exe")
-
-  if toolchainArchitecture != pythonArchitecture {
-    diagnosticsEngine.emit(
-      .warning(
-        """
-        There is an architecture mismatch between the installed toolchain and the resolved Python's architecture:
-        Toolchain: \(toolchainArchitecture)
-        Python: \(pythonArchitecture)
-        """))
-  }
 }
 
 /// Some of the architectures that can be stored in a COFF header.
-enum COFFBinaryExecutableArchitecture: String {
+public enum ExecutableArchitecture: String {
   case x86 = "X86"
   case x64 = "X64"
   case arm64 = "ARM64"
+  case universal = "Universal"
   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
-    }
-  }
-
-  /// Resolves the filename from the `Path` environment variable and read its COFF header to determine the architecture
-  /// of the binary.
-  ///
-  /// - Parameters:
-  ///   - cwd: The current working directory.
-  ///   - env: A dictionary of the environment variables and their values. Usually of the parent shell.
-  ///   - filename: The name of the file we are resolving the architecture of.
-  /// - Returns: The architecture of the file which was found in the `Path`.
-  static func readWindowsExecutableArchitecture(
-    cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, filename: String
-  ) -> Self {
-    let searchPaths = getEnvSearchPaths(pathString: envBlock["Path"], currentWorkingDirectory: cwd)
-    guard
-      let filePath = lookupExecutablePath(
-        filename: filename, currentWorkingDirectory: cwd, searchPaths: searchPaths)
-    else {
-      return .unknown
-    }
-    guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else {
-      return .unknown
+  #if os(Windows)
+    static func fromPEMachineByte(machine: Int32) -> Self {
+      // https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#machine-types
+      switch machine {
+      case IMAGE_FILE_MACHINE_I386: return .x86
+      case IMAGE_FILE_MACHINE_AMD64: return .x64
+      case IMAGE_FILE_MACHINE_ARM64: return .arm64
+      default: return .unknown
+      }
     }
+  #endif
 
-    defer { fileHandle.closeFile() }
-
-    // Infering the architecture of a Windows executable from its COFF header involves the following:
-    // 1. Get the COFF header offset from the pointer located at the 0x3C offset (4 bytes long).
-    // 2. Jump to that offset and read the next 6 bytes.
-    // 3. The first 4 are the signature which should be equal to 0x50450000.
-    // 4. The last 2 are the machine architecture which can be infered from the value we get.
-    //
-    // The link below provides a visualization of the COFF header and the process to get to it.
-    // 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
+  #if os(macOS)
+    static func fromMachoCPUType(cpuType: Int32) -> Self {
+      // https://en.wikipedia.org/wiki/Mach-O
+      switch cpuType {
+      case 0x0100_0007: return .x86
+      case 0x0100_000c: return .arm64
+      default: return .unknown
+      }
     }
+  #endif
+}
 
-    let peHeaderOffset = offsetPointer.withUnsafeBytes { $0.load(as: UInt32.self) }
+extension Process {
+  #if os(Windows)
+    /// Resolves the filename from the `Path` environment variable and read its COFF header to determine the architecture
+    /// of the binary.
+    ///
+    /// - Parameters:
+    ///   - cwd: The current working directory.
+    ///   - env: A dictionary of the environment variables and their values. Usually of the parent shell.
+    ///   - filename: The name of the file we are resolving the architecture of.
+    /// - Returns: The architecture of the file which was found in the `Path`.
+    static func readWindowsExecutableArchitecture(
+      cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, filename: String
+    ) -> ExecutableArchitecture {
+      let searchPaths = getEnvSearchPaths(
+        pathString: envBlock["Path"], currentWorkingDirectory: cwd)
+      guard
+        let filePath = lookupExecutablePath(
+          filename: filename, currentWorkingDirectory: cwd, searchPaths: searchPaths)
+      else {
+        return .unknown
+      }
+      guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else {
+        return .unknown
+      }
 
-    fileHandle.seek(toFileOffset: UInt64(peHeaderOffset))
-    guard let coffHeader = try? fileHandle.read(upToCount: 6), coffHeader.count == 6 else {
-      return .unknown
-    }
+      defer { fileHandle.closeFile() }
+
+      // Infering the architecture of a Windows executable from its COFF header involves the following:
+      // 1. Get the COFF header offset from the pointer located at the 0x3C offset (4 bytes long).
+      // 2. Jump to that offset and read the next 6 bytes.
+      // 3. The first 4 are the signature which should be equal to 0x50450000.
+      // 4. The last 2 are the machine architecture which can be infered from the value we get.
+      //
+      // The link below provides a visualization of the COFF header and the process to get to it.
+      // https://upload.wikimedia.org/wikipedia/commons/1/1b/Portable_Executable_32_bit_Structure_in_SVG_fixed.svg
+      guard (try? fileHandle.seek(toOffset: 0x3C)) != nil else {
+        return .unknown
+      }
+      guard let offsetPointer = try? fileHandle.read(upToCount: 4),
+        offsetPointer.count == 4
+      else {
+        return .unknown
+      }
+
+      let peHeaderOffset = offsetPointer.withUnsafeBytes { $0.load(as: UInt32.self) }
 
-    let signature = coffHeader.prefix(4)
-    let machineBytes = coffHeader.suffix(2)
+      guard (try? fileHandle.seek(toOffset: UInt64(peHeaderOffset))) != nil else {
+        return .unknown
+      }
+      guard let coffHeader = try? fileHandle.read(upToCount: 6), coffHeader.count == 6 else {
+        return .unknown
+      }
 
-    guard signature == Data([0x50, 0x45, 0x00, 0x00]) else {
-      return .unknown
+      let signature = coffHeader.prefix(4)
+      let machineBytes = coffHeader.suffix(2)
+
+      guard signature == Data([0x50, 0x45, 0x00, 0x00]) else {
+        return .unknown
+      }
+
+      let machine = machineBytes.withUnsafeBytes { $0.load(as: UInt16.self) }
+      return .fromPEMachineByte(machine: Int32(machine))
     }
+  #endif
+
+  #if os(macOS)
+    static func readDarwinExecutableArchitecture(
+      cwd: AbsolutePath?, envBlock: ProcessEnvironmentBlock, filename: String
+    ) -> ExecutableArchitecture {
+      let magicNumber: UInt32 = 0xcafe_babe
+
+      let searchPaths = getEnvSearchPaths(
+        pathString: envBlock["PATH"], currentWorkingDirectory: cwd)
+      guard
+        let filePath = lookupExecutablePath(
+          filename: filename, currentWorkingDirectory: cwd, searchPaths: searchPaths)
+      else {
+        return .unknown
+      }
+      guard let fileHandle = FileHandle(forReadingAtPath: filePath.pathString) else {
+        return .unknown
+      }
 
-    return .fromPEMachineByte(machine: machineBytes.withUnsafeBytes { $0.load(as: UInt16.self) })
-  }
+      defer {
+        try? fileHandle.close()
+      }
+
+      // The first 4 bytes of a Mach-O header contain the magic number. We use it to determine if the binary is
+      // universal.
+      // https://github.com/apple/darwin-xnu/blob/main/EXTERNAL_HEADERS/mach-o/loader.h
+      let magicData = fileHandle.readData(ofLength: 4)
+      let magic = magicData.withUnsafeBytes { $0.load(as: UInt32.self).bigEndian }
+
+      if magic == magicNumber {
+        return .universal
+      }
+
+      // If the binary is not universal, the next 4 bytes contain the CPU type.
+      guard (try? fileHandle.seek(toOffset: 4)) != nil else {
+        return .unknown
+      }
+      let cpuTypeData = fileHandle.readData(ofLength: 4)
+      let cpuType = cpuTypeData.withUnsafeBytes { $0.load(as: Int32.self) }
+      return .fromMachoCPUType(cpuType: cpuType)
+    }
+  #endif
 }
