From 599a34d45d99c98bdc575ba42750933c7b51889f Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Tue, 15 Apr 2025 11:20:34 +1200 Subject: [PATCH 1/4] Add PVM command and update dependencies in Tools package --- .gitignore | 2 + PolkaVM/Sources/PolkaVM/Registers.swift | 12 ++++ Tools/Package.swift | 4 +- Tools/Sources/Tools/PVM.swift | 79 +++++++++++++++++++++++++ Tools/Sources/Tools/Tools.swift | 1 + 5 files changed, 97 insertions(+), 1 deletion(-) create mode 100644 Tools/Sources/Tools/PVM.swift diff --git a/.gitignore b/.gitignore index 581b1b7e..4e8a4dab 100644 --- a/.gitignore +++ b/.gitignore @@ -28,3 +28,5 @@ c_cpp_properties.json Tools/openrpc.json tmp + +*.pvm diff --git a/PolkaVM/Sources/PolkaVM/Registers.swift b/PolkaVM/Sources/PolkaVM/Registers.swift index 60f7ae04..5ff69dc1 100644 --- a/PolkaVM/Sources/PolkaVM/Registers.swift +++ b/PolkaVM/Sources/PolkaVM/Registers.swift @@ -152,3 +152,15 @@ extension Registers: Codable { } } } + +extension Registers: CustomStringConvertible { + public var description: String { + var res = "" + for i in 0 ..< 13 { + let value = self[Registers.Index(raw: UInt8(i))] + let formatted = String(format: "%08X", value) + res += "w\(i):\t0x\(formatted) | \(value)\n" + } + return res + } +} diff --git a/Tools/Package.swift b/Tools/Package.swift index a8b9fb7f..d88d75a7 100644 --- a/Tools/Package.swift +++ b/Tools/Package.swift @@ -9,6 +9,7 @@ let package = Package( .macOS(.v15), ], dependencies: [ + .package(path: "../PolkaVM"), .package(path: "../RPC"), .package(path: "../TracingUtils"), .package(path: "../Utils"), @@ -20,9 +21,10 @@ let package = Package( .executableTarget( name: "Tools", dependencies: [ + "PolkaVM", "RPC", - "Utils", "TracingUtils", + "Utils", .product(name: "ArgumentParser", package: "swift-argument-parser"), .product(name: "JSONSchema", package: "swift-json-schema"), .product(name: "JSONSchemaBuilder", package: "swift-json-schema"), diff --git a/Tools/Sources/Tools/PVM.swift b/Tools/Sources/Tools/PVM.swift new file mode 100644 index 00000000..77aa0195 --- /dev/null +++ b/Tools/Sources/Tools/PVM.swift @@ -0,0 +1,79 @@ +import ArgumentParser +import Blockchain +import Codec +import Foundation +import PolkaVM +import TracingUtils +import Utils + +private let logger = Logger(label: "PVM") + +struct PVM: AsyncParsableCommand { + static let configuration = CommandConfiguration( + abstract: "PVM tools", + subcommands: [ + Invoke.self, + ] + ) + + struct Invoke: AsyncParsableCommand { + @Argument(help: "Program blob path") + var programPath: String + + @Option(help: "Program argument data path") + var argumentPath: String? + + @Option(help: "Program argument data in hex") + var argumentHex: String? + + @Option(help: "PC") + var pc: UInt32 = 0 + + @Option(help: "Gas") + var gas: UInt64 = 100_000_000 + + func run() async throws { + let blob: Data = try Data(contentsOf: URL(fileURLWithPath: programPath)) + let argumentData = if let argumentPath { + try Data(contentsOf: URL(fileURLWithPath: argumentPath)) + } else if let argumentHex { + Data(fromHexString: argumentHex).expect("invalid argument hex") + } else { + Data() + } + + logger.info("ArgumentData: \(argumentData.toHexString())") + + // setupTestLogger() + + let config = DefaultPvmConfig() + let gasValue = Gas(gas) + + do { + let state = try VMState(standardProgramBlob: blob, pc: pc, gas: gasValue, argumentData: argumentData) + let engine = Engine(config: config, invocationContext: nil) + let exitReason = await engine.execute(state: state) + let gasUsed = gasValue - Gas(state.getGas()) + + logger.info("Gas used: \(gasUsed)") + logger.info("Registers \n\(state.getRegisters())") + + switch exitReason { + case .halt: + let (addr, len): (UInt32, UInt32) = state.readRegister(Registers.Index(raw: 7), Registers.Index(raw: 8)) + let output = try? state.readMemory(address: addr, length: Int(len)) + if let output { + logger.info("Output: \(output.toHexString())") + } + logger.info("ExitReason: halt") + default: + logger.info("ExitReason: \(exitReason)") + } + } catch let e as StandardProgram.Error { + logger.error("Standard program initialization failed: \(e)") + } catch let e { + logger.error("Unknown error: \(e)") + } + } + } +} diff --git a/Tools/Sources/Tools/Tools.swift b/Tools/Sources/Tools/Tools.swift index 0c230263..51476b38 100644 --- a/Tools/Sources/Tools/Tools.swift +++ b/Tools/Sources/Tools/Tools.swift @@ -10,6 +10,7 @@ struct Boka: AsyncParsableCommand { version: "0.0.1", subcommands: [ OpenRPC.self, + PVM.self, ] ) From f76faafa3b3ecfbd50d2d7943edc2d692007de1a Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Wed, 23 Apr 2025 19:06:56 +1200 Subject: [PATCH 2/4] jit poc working --- Tools/Package.swift | 6 + Tools/Sources/CTools/ctools.c | 121 ++++++++++++ Tools/Sources/CTools/ctools.h | 17 ++ Tools/Sources/Tools/POC.swift | 329 ++++++++++++++++++++++++++++++++ Tools/Sources/Tools/Tools.swift | 1 + 5 files changed, 474 insertions(+) create mode 100644 Tools/Sources/CTools/ctools.c create mode 100644 Tools/Sources/CTools/ctools.h create mode 100644 Tools/Sources/Tools/POC.swift diff --git a/Tools/Package.swift b/Tools/Package.swift index d88d75a7..fcc97c48 100644 --- a/Tools/Package.swift +++ b/Tools/Package.swift @@ -25,12 +25,18 @@ let package = Package( "RPC", "TracingUtils", "Utils", + "CTools", .product(name: "ArgumentParser", package: "swift-argument-parser"), .product(name: "JSONSchema", package: "swift-json-schema"), .product(name: "JSONSchemaBuilder", package: "swift-json-schema"), .product(name: "Runtime", package: "Runtime"), ] ), + .target( + name: "CTools", + sources: ["ctools.h", "ctools.c"], + publicHeadersPath: "." + ), ], swiftLanguageModes: [.version("6")] ) diff --git a/Tools/Sources/CTools/ctools.c b/Tools/Sources/CTools/ctools.c new file mode 100644 index 00000000..30e48488 --- /dev/null +++ b/Tools/Sources/CTools/ctools.c @@ -0,0 +1,121 @@ +#include +#include +#include +#include +#include +#include "ctools.h" + +#if defined(__aarch64__) + +// Fixed encoding for MOVZ: base opcode 0xD2800000 for 64-bit immediate move with zero. +static uint32_t encodeMovZ(unsigned int xReg, unsigned int imm16, unsigned int shift) +{ + uint32_t insn = 0xD2800000; + insn |= ((shift & 3) << 21); + insn |= ((imm16 & 0xFFFF) << 5); + insn |= (xReg & 0x1F); + return insn; +} + +// Fixed encoding for MOVK: base opcode 0xF2A00000 for inserting a 16-bit constant. +static uint32_t encodeMovK(unsigned int xReg, unsigned int imm16, unsigned int shift) +{ + uint32_t insn = 0xF2A00000; + insn |= ((shift & 3) << 21); + insn |= ((imm16 & 0xFFFF) << 5); + insn |= (xReg & 0x1F); + return insn; +} + +// Encode "ldr wReg, [xReg]" with immediate offset zero. +static uint32_t encodeLdrWRegXReg(unsigned int wReg, unsigned int xReg) +{ + uint32_t insn = 0xB9400000; + insn |= ((xReg & 0x1F) << 5); + insn |= (wReg & 0x1F); + return insn; +} + +// Encode "add wDest, wSrc1, wSrc2" (shifted register variant with shift=0). +static uint32_t encodeAddWReg(unsigned int wd, unsigned int wn, unsigned int wm) +{ + uint32_t insn = 0x0B000000; + insn |= ((wm & 0x1F) << 16); + insn |= ((wn & 0x1F) << 5); + insn |= (wd & 0x1F); + return insn; +} + +// Encode "add xDest, xSrc1, xSrc2" (64‑bit, shifted register, shift = 0) +static uint32_t encodeAddXReg(unsigned int xd, unsigned int xn, unsigned int xm) +{ + uint32_t insn = 0x8B000000; /* base opcode for 64‑bit ADD (shifted register) */ + insn |= ((xm & 0x1F) << 16); /* Rm */ + insn |= ((xn & 0x1F) << 5); /* Rn */ + insn |= (xd & 0x1F); /* Rd */ + return insn; +} + +// Encode "str wReg, [xReg]" with immediate offset zero. +static uint32_t encodeStrWRegXReg(unsigned int wReg, unsigned int xReg) +{ + uint32_t insn = 0xB9000000; + insn |= ((xReg & 0x1F) << 5); + insn |= (wReg & 0x1F); + return insn; +} + +// Encode "ret" +static uint32_t encodeRet() +{ + return 0xD65F03C0; +} + +// Build a 64-bit constant into register xReg using MOVZ/MOVK. +static int emitLoadAddress64(unsigned char **p, unsigned int xReg, uint64_t addr) +{ + uint32_t insn; + insn = encodeMovZ(xReg, (uint16_t)(addr & 0xFFFF), 0); + memcpy(*p, &insn, 4); + *p += 4; + insn = encodeMovK(xReg, (uint16_t)((addr >> 16) & 0xFFFF), 1); + memcpy(*p, &insn, 4); + *p += 4; + insn = encodeMovK(xReg, (uint16_t)((addr >> 32) & 0xFFFF), 2); + memcpy(*p, &insn, 4); + *p += 4; + insn = encodeMovK(xReg, (uint16_t)((addr >> 48) & 0xFFFF), 3); + memcpy(*p, &insn, 4); + *p += 4; + return 16; +} + +int emitAddExample(void *codePtr) +{ + unsigned char *p = (unsigned char *)codePtr; + + /* add x0, x0, x1 */ + uint32_t insn = encodeAddXReg(0, 0, 1); + memcpy(p, &insn, 4); + p += 4; + + /* ret */ + insn = encodeRet(); + memcpy(p, &insn, 4); + p += 4; + + return (int)(p - (unsigned char *)codePtr); /* should be 8 bytes */ +} + +#else +// Fallback for unsupported architectures +int emitAddExample(void *codePtr) +{ + (void)codePtr; + return 0; +} +#endif + +int ctools_shm_open(const char *name, int oflag, mode_t mode) { + return shm_open(name, oflag, mode); +} diff --git a/Tools/Sources/CTools/ctools.h b/Tools/Sources/CTools/ctools.h new file mode 100644 index 00000000..44a2806f --- /dev/null +++ b/Tools/Sources/CTools/ctools.h @@ -0,0 +1,17 @@ +#include + +/** + * Emit machine code that: + * - Loads 32-bit values from (dataAddr + reg1Offset) and (dataAddr + reg2Offset) + * - Adds them + * - Stores the 32-bit result at (dataAddr + regDestOffset) + * - Returns (ret) + * + * codePtr: pointer to the code buffer (writable/executable) + * + * Returns the number of bytes of machine code emitted. + */ + int emitAddExample(void *codePtr); + + +int ctools_shm_open(const char *name, int oflag, mode_t mode); diff --git a/Tools/Sources/Tools/POC.swift b/Tools/Sources/Tools/POC.swift new file mode 100644 index 00000000..9ec9d39d --- /dev/null +++ b/Tools/Sources/Tools/POC.swift @@ -0,0 +1,329 @@ +import ArgumentParser +import CTools +import Foundation + +#if os(Linux) + import Glibc +#else + import Darwin +#endif + +// Helper function for logging to standard error +func log(_ message: String) { + fputs(message + "\n", stderr) +} + +// MARK: - IPC Message Protocol + +enum IPCMessage: UInt32 { + case run = 0x0000_0001 + case done = 0x0000_0002 + case exit = 0x0000_0003 + + @discardableResult + func write(to fd: FileHandle) -> Result { + var rawValue = rawValue + return Result { + try fd.write(contentsOf: Data(bytes: &rawValue, count: MemoryLayout.size)) + _ = try? fd.synchronize() + } + } + + static func read(from fd: FileHandle) -> Result { + Result { + let data = try fd.read(upToCount: MemoryLayout.size) + if let data, let decoded = IPCMessage(rawValue: data.withUnsafeBytes { $0.load(as: UInt32.self) }) { + return decoded + } + throw POCError.invalidIPCMessage(data: data) + } + } +} + +// MARK: - Error Handling + +enum POCError: Error, CustomStringConvertible { + case shmOpenFailed(errno: Int32) + case shmTruncateFailed(errno: Int32) + case mmapFailed(errno: Int32) + case childProcessFailed(error: Error) + case ipcReadFailed(errno: Int32) + case invalidIPCMessage(data: Data?) + + var description: String { + switch self { + case let .shmOpenFailed(errno): + "Failed to open shared memory: \(String(cString: strerror(errno)))" + case let .shmTruncateFailed(errno): + "Failed to set shared memory size: \(String(cString: strerror(errno)))" + case let .mmapFailed(errno): + "Failed to map memory: \(String(cString: strerror(errno)))" + case let .childProcessFailed(error): + "Failed to spawn child process: \(error)" + case let .ipcReadFailed(errno): + "Failed to read IPC message: \(String(cString: strerror(errno)))" + case let .invalidIPCMessage(value): + "Received invalid IPC message: \(String(describing: value))" + } + } +} + +struct POC: AsyncParsableCommand { + static let configuration = CommandConfiguration(abstract: "POC tools") + + @Flag var worker: Bool = false + @Option var shmName: String = "/poc_recompiler_shm" + + func run() async throws { + if worker { + try runChildMode(shmName: shmName) + } else { + try runParentMode(shmName: shmName) + } + } +} + +// MARK: - Configuration + +private let SHM_SIZE = 4096 +private let CODE_OFFSET = 0 +private let DATA_OFFSET = 1024 +private let CODE_SIZE = 8 // size of generated stub on AArch64 + +// We'll pass IPC channels on these file descriptor numbers in the child. +private let IPC_READ_FD: Int32 = 3 // Child will read commands from here +private let IPC_WRITE_FD: Int32 = 4 // Child will write responses here + +// MARK: - Child Entry Point + +func runChildMode(shmName: String) throws { + log("[Child] Starting child mode…") + // Use dedicated IPC FDs (3 and 4) rather than standard in/out. + let pipeInFD = FileHandle(fileDescriptor: IPC_READ_FD) + let pipeOutFD = FileHandle(fileDescriptor: IPC_WRITE_FD) + + // 1) Open shared memory. + let shmFD = ctools_shm_open(shmName, O_RDWR, 0o600) + if shmFD < 0 { + throw POCError.shmOpenFailed(errno: errno) + } + log("[Child] Shared memory opened successfully with FD: \(shmFD)") + + // 2) Map shared memory for IPC data only + let childMap = mmap(nil, SHM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, shmFD, 0) + guard childMap != MAP_FAILED else { + throw POCError.mmapFailed(errno: errno) + } + log("[Child] Shared memory mapped successfully at address: \(childMap!)") + + // 3) Allocate a separate JIT region for executing injected code + let codeMap = mmap( + nil, + SHM_SIZE, PROT_READ | PROT_WRITE | PROT_EXEC, + MAP_ANON | MAP_PRIVATE | MAP_JIT, + -1, + 0 + ) + guard codeMap != MAP_FAILED else { + throw POCError.mmapFailed(errno: errno) + } + + log("[Child] JIT region allocated successfully at address: \(codeMap!)") + + // Copy generated machine code from shared memory into JIT buffer + pthread_jit_write_protect_np(0) + memcpy(codeMap, childMap!.advanced(by: CODE_OFFSET), CODE_SIZE) + + log("[Child] Code copied to JIT region") + + sys_icache_invalidate(codeMap, Int(SHM_SIZE)) + pthread_jit_write_protect_np(1) + + log("[Child] JIT region invalidated") + + // 3) Listen for IPC messages on the dedicated FDs. + outer: + while true + { + log("[Child] Waiting for IPC message on fd: \(pipeInFD)") + let messageResult = IPCMessage.read(from: pipeInFD) + log("[Child] Received IPC message: \(messageResult)") + + switch messageResult { + case let .success(message): + switch message { + case .run: + let codePtr = codeMap! + typealias FuncType = @convention(c) (UInt64, UInt64) -> UInt64 + let fn: FuncType = unsafeBitCast(codePtr, to: FuncType.self) + + // Read operands from shared memory + let op1 = childMap!.advanced(by: DATA_OFFSET + 12).loadUnaligned(as: UInt64.self) + let op2 = childMap!.advanced(by: DATA_OFFSET + 20).loadUnaligned(as: UInt64.self) + + log("[Child] Executing code function with operands \(op1) and \(op2)") + let res = fn(op1, op2) + log("[Child] Function returned \(res)") + + // Store the result back to shared memory at offset +4 (keeping the original layout) + childMap!.advanced(by: DATA_OFFSET + 4).storeBytes(of: res, as: UInt64.self) + + log("[Child] Executed code function, sending DONE response") + try IPCMessage.done.write(to: pipeOutFD).get() + log("[Child] DONE response sent") + case .exit: + log("[Child] Received EXIT command, terminating child mode") + break outer + default: + log("[Child] Received unexpected message: \(message)") + } + case let .failure(error): + log("[Child] Error: \(error)") + break outer + } + } + + munmap(codeMap, SHM_SIZE) + munmap(childMap, SHM_SIZE) + close(shmFD) +} + +// MARK: - Parent Entry Point (Using posix_spawn with Extra IPC FDs) + +func runParentMode(shmName: String) throws { + print("[Parent] Starting parent mode…") + + // Create two pipes: one for sending commands to the child and one for receiving responses. + let inputPipe = Pipe() // Parent writes; child reads. + let outputPipe = Pipe() // Child writes; parent reads. + print("[Parent] Pipes created for IPC.") + + // 1) Create and configure shared memory. + shm_unlink(shmName) // Remove any existing shared memory. + let shmFD = ctools_shm_open(shmName, O_CREAT | O_RDWR, 0o600) + if shmFD < 0 { + throw POCError.shmOpenFailed(errno: errno) + } + print("[Parent] Shared memory created with FD: \(shmFD)") + if ftruncate(shmFD, off_t(SHM_SIZE)) != 0 { + throw POCError.shmTruncateFailed(errno: errno) + } + print("[Parent] Shared memory truncated to \(SHM_SIZE) bytes") + + guard let parentMap = mmap(nil, SHM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, shmFD, 0) else { + throw POCError.mmapFailed(errno: errno) + } + if parentMap == MAP_FAILED { + throw POCError.mmapFailed(errno: errno) + } + print("[Parent] Memory mapped at address: \(parentMap)") + + // 2) Spawn the child process using posix_spawn with file actions. + // Instead of modifying STDIN/STDOUT, we'll map the IPC pipes to FDs 3 and 4. + print("[Parent] Launching child process with posix_spawn") + + let childIpcReadFD = inputPipe.fileHandleForReading.fileDescriptor // For child: will be dup'd to FD 3. + let childIpcWriteFD = outputPipe.fileHandleForWriting.fileDescriptor // For child: will be dup'd to FD 4. + + // Keep parent's write end (for inputPipe) and read end (for outputPipe) for IPC. + let parentInputWriteFD = inputPipe.fileHandleForWriting.fileDescriptor + let parentOutputReadFD = outputPipe.fileHandleForReading.fileDescriptor + + var fileActions = posix_spawn_file_actions_t(bitPattern: 0) + posix_spawn_file_actions_init(&fileActions) + + // Duplicate the required descriptors into the child on FDs 3 and 4. + posix_spawn_file_actions_adddup2(&fileActions, childIpcReadFD, IPC_READ_FD) + posix_spawn_file_actions_adddup2(&fileActions, childIpcWriteFD, IPC_WRITE_FD) + // Close parent's ends in the child. + // posix_spawn_file_actions_addclose(&fileActions, parentInputWriteFD) + // posix_spawn_file_actions_addclose(&fileActions, parentOutputReadFD) + + let executable = CommandLine.arguments[0] + let arg0 = executable + let arg1 = "poc" + let arg2 = "--worker" + let arg3 = "--shm-name=\(shmName)" + + var args: [UnsafeMutablePointer?] = [ + strdup(arg0), + strdup(arg1), + strdup(arg2), + strdup(arg3), + nil, + ] + + var pid: pid_t = 0 + let spawnResult = posix_spawn(&pid, executable, &fileActions, nil, &args, environ) + posix_spawn_file_actions_destroy(&fileActions) + + // Free the C strings. + for arg in args { + if let arg { free(arg) } + } + + if spawnResult != 0 { + let errMsg = String(cString: strerror(spawnResult)) + throw POCError.childProcessFailed(error: NSError( + domain: NSPOSIXErrorDomain, + code: Int(spawnResult), + userInfo: [NSLocalizedDescriptionKey: errMsg] + )) + } + print("[Parent] Child process spawned with pid: \(pid)") + + // 3) Initialize shared memory data and emit code. + let codePtr = parentMap + CODE_OFFSET + let dataPtr = parentMap + DATA_OFFSET + (dataPtr + 0).storeBytes(of: UInt32(0), as: UInt32.self) + (dataPtr + 4).storeBytes(of: UInt64(0), as: UInt64.self) + (dataPtr + 12).storeBytes(of: UInt64(10), as: UInt64.self) + (dataPtr + 20).storeBytes(of: UInt64(32), as: UInt64.self) + print("[Parent] Data layout initialized: operands 10 and 32") + + let codeSize = emitAddExample(codePtr) + print("[Parent] Add example code emitted (\(codeSize) bytes) at codePtr: \(codePtr)") + + mprotect(parentMap, SHM_SIZE, PROT_READ | PROT_EXEC) + print("[Parent] Memory protection updated to READ+EXEC") + + // 4) Use the parent's ends of the pipes to communicate. + print("[Parent] Sending RUN command to child") + let runResult = IPCMessage.run.write(to: FileHandle(fileDescriptor: parentInputWriteFD)) + if case let .failure(error) = runResult { + throw error + } + print("[Parent] RUN command sent") + + print("[Parent] Waiting for child's response") + let responseResult = IPCMessage.read(from: FileHandle(fileDescriptor: parentOutputReadFD)) + switch responseResult { + case let .success(message): + if message == .done { + print("[Parent] Child responded with DONE") + } else { + print("[Parent] Child sent unexpected response: \(message)") + } + case let .failure(error): + print("[Parent] Failed to read response: \(error)") + throw error + } + + let result = (dataPtr + 4).load(as: UInt32.self) + print("[Parent] Sum result = \(result) (expected 42)") + + print("[Parent] Sending EXIT command to child") + let exitResult = IPCMessage.exit.write(to: FileHandle(fileDescriptor: parentInputWriteFD)) + if case let .failure(error) = exitResult { + print("[Parent] Failed to send EXIT command: \(error)") + } + print("[Parent] EXIT command sent, proceeding with cleanup") + + munmap(parentMap, SHM_SIZE) + close(shmFD) + print("[Parent] Cleanup complete; waiting for child to exit") + + var status: Int32 = 0 + waitpid(pid, &status, 0) + print("[Parent] Child exit code: \(status)") +} diff --git a/Tools/Sources/Tools/Tools.swift b/Tools/Sources/Tools/Tools.swift index 51476b38..5238cd9e 100644 --- a/Tools/Sources/Tools/Tools.swift +++ b/Tools/Sources/Tools/Tools.swift @@ -11,6 +11,7 @@ struct Boka: AsyncParsableCommand { subcommands: [ OpenRPC.self, PVM.self, + POC.self, ] ) From 5abde646e9c2a50ed6f645a66ef8639b348511a9 Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Thu, 24 Apr 2025 13:20:36 +1200 Subject: [PATCH 3/4] only osx --- Tools/Sources/Tools/POC.swift | 455 +++++++++++++++++----------------- 1 file changed, 230 insertions(+), 225 deletions(-) diff --git a/Tools/Sources/Tools/POC.swift b/Tools/Sources/Tools/POC.swift index 9ec9d39d..466ff66a 100644 --- a/Tools/Sources/Tools/POC.swift +++ b/Tools/Sources/Tools/POC.swift @@ -75,255 +75,260 @@ struct POC: AsyncParsableCommand { @Option var shmName: String = "/poc_recompiler_shm" func run() async throws { - if worker { - try runChildMode(shmName: shmName) - } else { - try runParentMode(shmName: shmName) - } + #if os(macOS) + if worker { + try runChildMode(shmName: shmName) + } else { + try runParentMode(shmName: shmName) + } + #endif } } -// MARK: - Configuration - -private let SHM_SIZE = 4096 -private let CODE_OFFSET = 0 -private let DATA_OFFSET = 1024 -private let CODE_SIZE = 8 // size of generated stub on AArch64 - -// We'll pass IPC channels on these file descriptor numbers in the child. -private let IPC_READ_FD: Int32 = 3 // Child will read commands from here -private let IPC_WRITE_FD: Int32 = 4 // Child will write responses here - -// MARK: - Child Entry Point +#if os(macOS) -func runChildMode(shmName: String) throws { - log("[Child] Starting child mode…") - // Use dedicated IPC FDs (3 and 4) rather than standard in/out. - let pipeInFD = FileHandle(fileDescriptor: IPC_READ_FD) - let pipeOutFD = FileHandle(fileDescriptor: IPC_WRITE_FD) - - // 1) Open shared memory. - let shmFD = ctools_shm_open(shmName, O_RDWR, 0o600) - if shmFD < 0 { - throw POCError.shmOpenFailed(errno: errno) - } - log("[Child] Shared memory opened successfully with FD: \(shmFD)") - - // 2) Map shared memory for IPC data only - let childMap = mmap(nil, SHM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, shmFD, 0) - guard childMap != MAP_FAILED else { - throw POCError.mmapFailed(errno: errno) - } - log("[Child] Shared memory mapped successfully at address: \(childMap!)") - - // 3) Allocate a separate JIT region for executing injected code - let codeMap = mmap( - nil, - SHM_SIZE, PROT_READ | PROT_WRITE | PROT_EXEC, - MAP_ANON | MAP_PRIVATE | MAP_JIT, - -1, - 0 - ) - guard codeMap != MAP_FAILED else { - throw POCError.mmapFailed(errno: errno) - } + // MARK: - Configuration - log("[Child] JIT region allocated successfully at address: \(codeMap!)") + private let SHM_SIZE = 4096 + private let CODE_OFFSET = 0 + private let DATA_OFFSET = 1024 + private let CODE_SIZE = 8 // size of generated stub on AArch64 - // Copy generated machine code from shared memory into JIT buffer - pthread_jit_write_protect_np(0) - memcpy(codeMap, childMap!.advanced(by: CODE_OFFSET), CODE_SIZE) + // We'll pass IPC channels on these file descriptor numbers in the child. + private let IPC_READ_FD: Int32 = 3 // Child will read commands from here + private let IPC_WRITE_FD: Int32 = 4 // Child will write responses here - log("[Child] Code copied to JIT region") + // MARK: - Child Entry Point - sys_icache_invalidate(codeMap, Int(SHM_SIZE)) - pthread_jit_write_protect_np(1) + func runChildMode(shmName: String) throws { + log("[Child] Starting child mode…") + // Use dedicated IPC FDs (3 and 4) rather than standard in/out. + let pipeInFD = FileHandle(fileDescriptor: IPC_READ_FD) + let pipeOutFD = FileHandle(fileDescriptor: IPC_WRITE_FD) - log("[Child] JIT region invalidated") + // 1) Open shared memory. + let shmFD = ctools_shm_open(shmName, O_RDWR, 0o600) + if shmFD < 0 { + throw POCError.shmOpenFailed(errno: errno) + } + log("[Child] Shared memory opened successfully with FD: \(shmFD)") - // 3) Listen for IPC messages on the dedicated FDs. - outer: - while true - { - log("[Child] Waiting for IPC message on fd: \(pipeInFD)") - let messageResult = IPCMessage.read(from: pipeInFD) - log("[Child] Received IPC message: \(messageResult)") + // 2) Map shared memory for IPC data only + let childMap = mmap(nil, SHM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, shmFD, 0) + guard childMap != MAP_FAILED else { + throw POCError.mmapFailed(errno: errno) + } + log("[Child] Shared memory mapped successfully at address: \(childMap!)") + + // 3) Allocate a separate JIT region for executing injected code + let codeMap = mmap( + nil, + SHM_SIZE, PROT_READ | PROT_WRITE | PROT_EXEC, + MAP_ANON | MAP_PRIVATE | MAP_JIT, + -1, + 0 + ) + guard codeMap != MAP_FAILED else { + throw POCError.mmapFailed(errno: errno) + } - switch messageResult { - case let .success(message): - switch message { - case .run: - let codePtr = codeMap! - typealias FuncType = @convention(c) (UInt64, UInt64) -> UInt64 - let fn: FuncType = unsafeBitCast(codePtr, to: FuncType.self) - - // Read operands from shared memory - let op1 = childMap!.advanced(by: DATA_OFFSET + 12).loadUnaligned(as: UInt64.self) - let op2 = childMap!.advanced(by: DATA_OFFSET + 20).loadUnaligned(as: UInt64.self) - - log("[Child] Executing code function with operands \(op1) and \(op2)") - let res = fn(op1, op2) - log("[Child] Function returned \(res)") - - // Store the result back to shared memory at offset +4 (keeping the original layout) - childMap!.advanced(by: DATA_OFFSET + 4).storeBytes(of: res, as: UInt64.self) - - log("[Child] Executed code function, sending DONE response") - try IPCMessage.done.write(to: pipeOutFD).get() - log("[Child] DONE response sent") - case .exit: - log("[Child] Received EXIT command, terminating child mode") + log("[Child] JIT region allocated successfully at address: \(codeMap!)") + + // Copy generated machine code from shared memory into JIT buffer + pthread_jit_write_protect_np(0) + memcpy(codeMap, childMap!.advanced(by: CODE_OFFSET), CODE_SIZE) + + log("[Child] Code copied to JIT region") + + sys_icache_invalidate(codeMap, Int(SHM_SIZE)) + pthread_jit_write_protect_np(1) + + log("[Child] JIT region invalidated") + + // 3) Listen for IPC messages on the dedicated FDs. + outer: + while true + { + log("[Child] Waiting for IPC message on fd: \(pipeInFD)") + let messageResult = IPCMessage.read(from: pipeInFD) + log("[Child] Received IPC message: \(messageResult)") + + switch messageResult { + case let .success(message): + switch message { + case .run: + let codePtr = codeMap! + typealias FuncType = @convention(c) (UInt64, UInt64) -> UInt64 + let fn: FuncType = unsafeBitCast(codePtr, to: FuncType.self) + + // Read operands from shared memory + let op1 = childMap!.advanced(by: DATA_OFFSET + 12).loadUnaligned(as: UInt64.self) + let op2 = childMap!.advanced(by: DATA_OFFSET + 20).loadUnaligned(as: UInt64.self) + + log("[Child] Executing code function with operands \(op1) and \(op2)") + let res = fn(op1, op2) + log("[Child] Function returned \(res)") + + // Store the result back to shared memory at offset +4 (keeping the original layout) + childMap!.advanced(by: DATA_OFFSET + 4).storeBytes(of: res, as: UInt64.self) + + log("[Child] Executed code function, sending DONE response") + try IPCMessage.done.write(to: pipeOutFD).get() + log("[Child] DONE response sent") + case .exit: + log("[Child] Received EXIT command, terminating child mode") + break outer + default: + log("[Child] Received unexpected message: \(message)") + } + case let .failure(error): + log("[Child] Error: \(error)") break outer - default: - log("[Child] Received unexpected message: \(message)") } - case let .failure(error): - log("[Child] Error: \(error)") - break outer } + + munmap(codeMap, SHM_SIZE) + munmap(childMap, SHM_SIZE) + close(shmFD) } - munmap(codeMap, SHM_SIZE) - munmap(childMap, SHM_SIZE) - close(shmFD) -} + // MARK: - Parent Entry Point (Using posix_spawn with Extra IPC FDs) -// MARK: - Parent Entry Point (Using posix_spawn with Extra IPC FDs) + func runParentMode(shmName: String) throws { + print("[Parent] Starting parent mode…") -func runParentMode(shmName: String) throws { - print("[Parent] Starting parent mode…") + // Create two pipes: one for sending commands to the child and one for receiving responses. + let inputPipe = Pipe() // Parent writes; child reads. + let outputPipe = Pipe() // Child writes; parent reads. + print("[Parent] Pipes created for IPC.") - // Create two pipes: one for sending commands to the child and one for receiving responses. - let inputPipe = Pipe() // Parent writes; child reads. - let outputPipe = Pipe() // Child writes; parent reads. - print("[Parent] Pipes created for IPC.") + // 1) Create and configure shared memory. + shm_unlink(shmName) // Remove any existing shared memory. + let shmFD = ctools_shm_open(shmName, O_CREAT | O_RDWR, 0o600) + if shmFD < 0 { + throw POCError.shmOpenFailed(errno: errno) + } + print("[Parent] Shared memory created with FD: \(shmFD)") + if ftruncate(shmFD, off_t(SHM_SIZE)) != 0 { + throw POCError.shmTruncateFailed(errno: errno) + } + print("[Parent] Shared memory truncated to \(SHM_SIZE) bytes") - // 1) Create and configure shared memory. - shm_unlink(shmName) // Remove any existing shared memory. - let shmFD = ctools_shm_open(shmName, O_CREAT | O_RDWR, 0o600) - if shmFD < 0 { - throw POCError.shmOpenFailed(errno: errno) - } - print("[Parent] Shared memory created with FD: \(shmFD)") - if ftruncate(shmFD, off_t(SHM_SIZE)) != 0 { - throw POCError.shmTruncateFailed(errno: errno) - } - print("[Parent] Shared memory truncated to \(SHM_SIZE) bytes") + guard let parentMap = mmap(nil, SHM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, shmFD, 0) else { + throw POCError.mmapFailed(errno: errno) + } + if parentMap == MAP_FAILED { + throw POCError.mmapFailed(errno: errno) + } + print("[Parent] Memory mapped at address: \(parentMap)") + + // 2) Spawn the child process using posix_spawn with file actions. + // Instead of modifying STDIN/STDOUT, we'll map the IPC pipes to FDs 3 and 4. + print("[Parent] Launching child process with posix_spawn") + + let childIpcReadFD = inputPipe.fileHandleForReading.fileDescriptor // For child: will be dup'd to FD 3. + let childIpcWriteFD = outputPipe.fileHandleForWriting.fileDescriptor // For child: will be dup'd to FD 4. + + // Keep parent's write end (for inputPipe) and read end (for outputPipe) for IPC. + let parentInputWriteFD = inputPipe.fileHandleForWriting.fileDescriptor + let parentOutputReadFD = outputPipe.fileHandleForReading.fileDescriptor + + var fileActions = posix_spawn_file_actions_t(bitPattern: 0) + posix_spawn_file_actions_init(&fileActions) + + // Duplicate the required descriptors into the child on FDs 3 and 4. + posix_spawn_file_actions_adddup2(&fileActions, childIpcReadFD, IPC_READ_FD) + posix_spawn_file_actions_adddup2(&fileActions, childIpcWriteFD, IPC_WRITE_FD) + // Close parent's ends in the child. + // posix_spawn_file_actions_addclose(&fileActions, parentInputWriteFD) + // posix_spawn_file_actions_addclose(&fileActions, parentOutputReadFD) + + let executable = CommandLine.arguments[0] + let arg0 = executable + let arg1 = "poc" + let arg2 = "--worker" + let arg3 = "--shm-name=\(shmName)" + + var args: [UnsafeMutablePointer?] = [ + strdup(arg0), + strdup(arg1), + strdup(arg2), + strdup(arg3), + nil, + ] + + var pid: pid_t = 0 + let spawnResult = posix_spawn(&pid, executable, &fileActions, nil, &args, environ) + posix_spawn_file_actions_destroy(&fileActions) + + // Free the C strings. + for arg in args { + if let arg { free(arg) } + } - guard let parentMap = mmap(nil, SHM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, shmFD, 0) else { - throw POCError.mmapFailed(errno: errno) - } - if parentMap == MAP_FAILED { - throw POCError.mmapFailed(errno: errno) - } - print("[Parent] Memory mapped at address: \(parentMap)") - - // 2) Spawn the child process using posix_spawn with file actions. - // Instead of modifying STDIN/STDOUT, we'll map the IPC pipes to FDs 3 and 4. - print("[Parent] Launching child process with posix_spawn") - - let childIpcReadFD = inputPipe.fileHandleForReading.fileDescriptor // For child: will be dup'd to FD 3. - let childIpcWriteFD = outputPipe.fileHandleForWriting.fileDescriptor // For child: will be dup'd to FD 4. - - // Keep parent's write end (for inputPipe) and read end (for outputPipe) for IPC. - let parentInputWriteFD = inputPipe.fileHandleForWriting.fileDescriptor - let parentOutputReadFD = outputPipe.fileHandleForReading.fileDescriptor - - var fileActions = posix_spawn_file_actions_t(bitPattern: 0) - posix_spawn_file_actions_init(&fileActions) - - // Duplicate the required descriptors into the child on FDs 3 and 4. - posix_spawn_file_actions_adddup2(&fileActions, childIpcReadFD, IPC_READ_FD) - posix_spawn_file_actions_adddup2(&fileActions, childIpcWriteFD, IPC_WRITE_FD) - // Close parent's ends in the child. - // posix_spawn_file_actions_addclose(&fileActions, parentInputWriteFD) - // posix_spawn_file_actions_addclose(&fileActions, parentOutputReadFD) - - let executable = CommandLine.arguments[0] - let arg0 = executable - let arg1 = "poc" - let arg2 = "--worker" - let arg3 = "--shm-name=\(shmName)" - - var args: [UnsafeMutablePointer?] = [ - strdup(arg0), - strdup(arg1), - strdup(arg2), - strdup(arg3), - nil, - ] - - var pid: pid_t = 0 - let spawnResult = posix_spawn(&pid, executable, &fileActions, nil, &args, environ) - posix_spawn_file_actions_destroy(&fileActions) - - // Free the C strings. - for arg in args { - if let arg { free(arg) } - } + if spawnResult != 0 { + let errMsg = String(cString: strerror(spawnResult)) + throw POCError.childProcessFailed(error: NSError( + domain: NSPOSIXErrorDomain, + code: Int(spawnResult), + userInfo: [NSLocalizedDescriptionKey: errMsg] + )) + } + print("[Parent] Child process spawned with pid: \(pid)") + + // 3) Initialize shared memory data and emit code. + let codePtr = parentMap + CODE_OFFSET + let dataPtr = parentMap + DATA_OFFSET + (dataPtr + 0).storeBytes(of: UInt32(0), as: UInt32.self) + (dataPtr + 4).storeBytes(of: UInt64(0), as: UInt64.self) + (dataPtr + 12).storeBytes(of: UInt64(10), as: UInt64.self) + (dataPtr + 20).storeBytes(of: UInt64(32), as: UInt64.self) + print("[Parent] Data layout initialized: operands 10 and 32") + + let codeSize = emitAddExample(codePtr) + print("[Parent] Add example code emitted (\(codeSize) bytes) at codePtr: \(codePtr)") + + mprotect(parentMap, SHM_SIZE, PROT_READ | PROT_EXEC) + print("[Parent] Memory protection updated to READ+EXEC") + + // 4) Use the parent's ends of the pipes to communicate. + print("[Parent] Sending RUN command to child") + let runResult = IPCMessage.run.write(to: FileHandle(fileDescriptor: parentInputWriteFD)) + if case let .failure(error) = runResult { + throw error + } + print("[Parent] RUN command sent") - if spawnResult != 0 { - let errMsg = String(cString: strerror(spawnResult)) - throw POCError.childProcessFailed(error: NSError( - domain: NSPOSIXErrorDomain, - code: Int(spawnResult), - userInfo: [NSLocalizedDescriptionKey: errMsg] - )) - } - print("[Parent] Child process spawned with pid: \(pid)") - - // 3) Initialize shared memory data and emit code. - let codePtr = parentMap + CODE_OFFSET - let dataPtr = parentMap + DATA_OFFSET - (dataPtr + 0).storeBytes(of: UInt32(0), as: UInt32.self) - (dataPtr + 4).storeBytes(of: UInt64(0), as: UInt64.self) - (dataPtr + 12).storeBytes(of: UInt64(10), as: UInt64.self) - (dataPtr + 20).storeBytes(of: UInt64(32), as: UInt64.self) - print("[Parent] Data layout initialized: operands 10 and 32") - - let codeSize = emitAddExample(codePtr) - print("[Parent] Add example code emitted (\(codeSize) bytes) at codePtr: \(codePtr)") - - mprotect(parentMap, SHM_SIZE, PROT_READ | PROT_EXEC) - print("[Parent] Memory protection updated to READ+EXEC") - - // 4) Use the parent's ends of the pipes to communicate. - print("[Parent] Sending RUN command to child") - let runResult = IPCMessage.run.write(to: FileHandle(fileDescriptor: parentInputWriteFD)) - if case let .failure(error) = runResult { - throw error - } - print("[Parent] RUN command sent") - - print("[Parent] Waiting for child's response") - let responseResult = IPCMessage.read(from: FileHandle(fileDescriptor: parentOutputReadFD)) - switch responseResult { - case let .success(message): - if message == .done { - print("[Parent] Child responded with DONE") - } else { - print("[Parent] Child sent unexpected response: \(message)") + print("[Parent] Waiting for child's response") + let responseResult = IPCMessage.read(from: FileHandle(fileDescriptor: parentOutputReadFD)) + switch responseResult { + case let .success(message): + if message == .done { + print("[Parent] Child responded with DONE") + } else { + print("[Parent] Child sent unexpected response: \(message)") + } + case let .failure(error): + print("[Parent] Failed to read response: \(error)") + throw error } - case let .failure(error): - print("[Parent] Failed to read response: \(error)") - throw error - } - let result = (dataPtr + 4).load(as: UInt32.self) - print("[Parent] Sum result = \(result) (expected 42)") + let result = (dataPtr + 4).load(as: UInt32.self) + print("[Parent] Sum result = \(result) (expected 42)") - print("[Parent] Sending EXIT command to child") - let exitResult = IPCMessage.exit.write(to: FileHandle(fileDescriptor: parentInputWriteFD)) - if case let .failure(error) = exitResult { - print("[Parent] Failed to send EXIT command: \(error)") - } - print("[Parent] EXIT command sent, proceeding with cleanup") + print("[Parent] Sending EXIT command to child") + let exitResult = IPCMessage.exit.write(to: FileHandle(fileDescriptor: parentInputWriteFD)) + if case let .failure(error) = exitResult { + print("[Parent] Failed to send EXIT command: \(error)") + } + print("[Parent] EXIT command sent, proceeding with cleanup") - munmap(parentMap, SHM_SIZE) - close(shmFD) - print("[Parent] Cleanup complete; waiting for child to exit") + munmap(parentMap, SHM_SIZE) + close(shmFD) + print("[Parent] Cleanup complete; waiting for child to exit") - var status: Int32 = 0 - waitpid(pid, &status, 0) - print("[Parent] Child exit code: \(status)") -} + var status: Int32 = 0 + waitpid(pid, &status, 0) + print("[Parent] Child exit code: \(status)") + } +#endif From 74cda74759a4be58b4a2cade4abb21e23fdf4fb4 Mon Sep 17 00:00:00 2001 From: Bryan Chen Date: Thu, 24 Apr 2025 13:45:38 +1200 Subject: [PATCH 4/4] Replace custom logging function with print statements for child mode IPC messages --- Tools/Sources/Tools/POC.swift | 35 +++++++++++++++-------------------- 1 file changed, 15 insertions(+), 20 deletions(-) diff --git a/Tools/Sources/Tools/POC.swift b/Tools/Sources/Tools/POC.swift index 466ff66a..02f79f73 100644 --- a/Tools/Sources/Tools/POC.swift +++ b/Tools/Sources/Tools/POC.swift @@ -8,11 +8,6 @@ import Foundation import Darwin #endif -// Helper function for logging to standard error -func log(_ message: String) { - fputs(message + "\n", stderr) -} - // MARK: - IPC Message Protocol enum IPCMessage: UInt32 { @@ -101,7 +96,7 @@ struct POC: AsyncParsableCommand { // MARK: - Child Entry Point func runChildMode(shmName: String) throws { - log("[Child] Starting child mode…") + print("[Child] Starting child mode…") // Use dedicated IPC FDs (3 and 4) rather than standard in/out. let pipeInFD = FileHandle(fileDescriptor: IPC_READ_FD) let pipeOutFD = FileHandle(fileDescriptor: IPC_WRITE_FD) @@ -111,14 +106,14 @@ struct POC: AsyncParsableCommand { if shmFD < 0 { throw POCError.shmOpenFailed(errno: errno) } - log("[Child] Shared memory opened successfully with FD: \(shmFD)") + print("[Child] Shared memory opened successfully with FD: \(shmFD)") // 2) Map shared memory for IPC data only let childMap = mmap(nil, SHM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, shmFD, 0) guard childMap != MAP_FAILED else { throw POCError.mmapFailed(errno: errno) } - log("[Child] Shared memory mapped successfully at address: \(childMap!)") + print("[Child] Shared memory mapped successfully at address: \(childMap!)") // 3) Allocate a separate JIT region for executing injected code let codeMap = mmap( @@ -132,26 +127,26 @@ struct POC: AsyncParsableCommand { throw POCError.mmapFailed(errno: errno) } - log("[Child] JIT region allocated successfully at address: \(codeMap!)") + print("[Child] JIT region allocated successfully at address: \(codeMap!)") // Copy generated machine code from shared memory into JIT buffer pthread_jit_write_protect_np(0) memcpy(codeMap, childMap!.advanced(by: CODE_OFFSET), CODE_SIZE) - log("[Child] Code copied to JIT region") + print("[Child] Code copied to JIT region") sys_icache_invalidate(codeMap, Int(SHM_SIZE)) pthread_jit_write_protect_np(1) - log("[Child] JIT region invalidated") + print("[Child] JIT region invalidated") // 3) Listen for IPC messages on the dedicated FDs. outer: while true { - log("[Child] Waiting for IPC message on fd: \(pipeInFD)") + print("[Child] Waiting for IPC message on fd: \(pipeInFD)") let messageResult = IPCMessage.read(from: pipeInFD) - log("[Child] Received IPC message: \(messageResult)") + print("[Child] Received IPC message: \(messageResult)") switch messageResult { case let .success(message): @@ -165,24 +160,24 @@ struct POC: AsyncParsableCommand { let op1 = childMap!.advanced(by: DATA_OFFSET + 12).loadUnaligned(as: UInt64.self) let op2 = childMap!.advanced(by: DATA_OFFSET + 20).loadUnaligned(as: UInt64.self) - log("[Child] Executing code function with operands \(op1) and \(op2)") + print("[Child] Executing code function with operands \(op1) and \(op2)") let res = fn(op1, op2) - log("[Child] Function returned \(res)") + print("[Child] Function returned \(res)") // Store the result back to shared memory at offset +4 (keeping the original layout) childMap!.advanced(by: DATA_OFFSET + 4).storeBytes(of: res, as: UInt64.self) - log("[Child] Executed code function, sending DONE response") + print("[Child] Executed code function, sending DONE response") try IPCMessage.done.write(to: pipeOutFD).get() - log("[Child] DONE response sent") + print("[Child] DONE response sent") case .exit: - log("[Child] Received EXIT command, terminating child mode") + print("[Child] Received EXIT command, terminating child mode") break outer default: - log("[Child] Received unexpected message: \(message)") + print("[Child] Received unexpected message: \(message)") } case let .failure(error): - log("[Child] Error: \(error)") + print("[Child] Error: \(error)") break outer } }