Limit the lifetime of started WinAppDriver.exe processes (#88)

Sources/WinAppDriver/Win32Error.swift

@@ -0,0 +1,10 @@
+import WinSDK
+
+internal struct Win32Error: Error {
+    public var apiName: String
+    public var errorCode: UInt32
+
+    internal static func getLastError(apiName: String) -> Self {
+        Self(apiName: apiName, errorCode: GetLastError())
+    }
+}

Sources/WinAppDriver/Win32ProcessTree.swift

@@ -0,0 +1,83 @@
+import WinSDK
+
+/// Starts and tracks the lifetime of a process tree using Win32 APIs.
+internal class Win32ProcessTree {
+    internal let jobHandle: HANDLE
+    internal let handle: HANDLE
+
+    init(path: String, args: [String]) throws {
+        // Use a job object to ensure that the process tree doesn't outlive us.
+        jobHandle = try Self.createJobObject()
+
+        let commandLine = buildCommandLineArgsString(args: [path] + args)
+        do { handle = try Self.createProcessInJob(commandLine: commandLine, jobHandle: jobHandle) }
+        catch {
+            CloseHandle(jobHandle)
+            throw error
+        }
+    }
+
+    func terminate() throws {
+        if !TerminateJobObject(jobHandle, UINT.max) {
+            throw Win32Error.getLastError(apiName: "TerminateJobObject")
+        }
+    }
+
+    deinit {
+        CloseHandle(handle)
+        CloseHandle(jobHandle)
+    }
+
+    private static func createJobObject() throws -> HANDLE {
+        guard let jobHandle = CreateJobObjectW(nil, nil) else {
+            throw Win32Error.getLastError(apiName: "CreateJobObjectW")
+        }
+
+        var limitInfo = JOBOBJECT_EXTENDED_LIMIT_INFORMATION()
+        limitInfo.BasicLimitInformation.LimitFlags = DWORD(JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE) | DWORD(JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK)
+        guard SetInformationJobObject(jobHandle, JobObjectExtendedLimitInformation,
+                &limitInfo, DWORD(MemoryLayout<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>.size)) else {
+            defer { CloseHandle(jobHandle) }
+            throw Win32Error.getLastError(apiName: "SetInformationJobObject")
+        }
+
+        return jobHandle
+    }
+
+    private static func createProcessInJob(commandLine: String, jobHandle: HANDLE) throws -> HANDLE {
+        try commandLine.withCString(encodedAs: UTF16.self) { commandLine throws in
+            var startupInfo = STARTUPINFOW()
+            startupInfo.cb = DWORD(MemoryLayout<STARTUPINFOW>.size)
+
+            var processInfo = PROCESS_INFORMATION()
+            guard CreateProcessW(
+                nil,
+                UnsafeMutablePointer<WCHAR>(mutating: commandLine),
+                nil,
+                nil,
+                false,
+                DWORD(CREATE_NEW_CONSOLE) | DWORD(CREATE_SUSPENDED) | DWORD(CREATE_NEW_PROCESS_GROUP),
+                nil,
+                nil,
+                &startupInfo,
+                &processInfo
+            ) else {
+                throw Win32Error.getLastError(apiName: "CreateProcessW")
+            }
+
+            defer { CloseHandle(processInfo.hThread) }
+
+            guard AssignProcessToJobObject(jobHandle, processInfo.hProcess) else {
+                defer { CloseHandle(processInfo.hProcess) }
+                throw Win32Error.getLastError(apiName: "AssignProcessToJobObject")
+            }
+
+            guard ResumeThread(processInfo.hThread) != DWORD.max else {
+                defer { CloseHandle(processInfo.hProcess) }
+                throw Win32Error.getLastError(apiName: "ResumeThread")
+            }
+
+            return processInfo.hProcess
+        }
+    }
+}

Sources/WinAppDriver/WinAppDriver.swift

@@ -2,76 +2,49 @@
 import Foundation
 import WinSDK
 
-public enum WinAppDriverError: Error {
-    // Exposes any underlying win32 errors that may surface as a result of process management.
-    case win32Error(lastError: Int)
-}
-
 public class WinAppDriver: WebDriver {
+    /// Raised when the WinAppDriver.exe process fails to start
+    public struct StartError: Error {
+        public var message: String
+    }
+
     public static let defaultIp = "127.0.0.1"
     public static let defaultPort = 4723
+    public static let executableName = "WinAppDriver.exe"
+    public static var defaultExecutablePath: String {
+        let programFilesX86 = ProcessInfo.processInfo.environment["ProgramFiles(x86)"]
+            ?? "\(ProcessInfo.processInfo.environment["SystemDrive"] ?? "C:")\\Program Files (x86)"
+        return "\(programFilesX86)\\Windows Application Driver\\\(executableName)"
+    }
 
-    static let processsName = "WinAppDriver.exe"
-
+    private var processTree: Win32ProcessTree?
     private let httpWebDriver: HTTPWebDriver
 
-    private let port: Int
-    private let ip: String
-
-    private var wadProcessInfo: PROCESS_INFORMATION?
-
-    public init(attachingTo ip: String, port: Int = WinAppDriver.defaultPort) throws {
-        self.ip = ip
-        self.port = port
-
-        httpWebDriver = HTTPWebDriver(endpoint: URL(string: "http://\(ip):\(port)")!)
+    public init(attachingTo ip: String, port: Int = WinAppDriver.defaultPort) {
+        self.httpWebDriver = HTTPWebDriver(endpoint: URL(string: "http://\(ip):\(port)")!)
     }
 
-    public init(_ ip: String = WinAppDriver.defaultIp, port: Int = WinAppDriver.defaultPort) throws {
-        self.ip = ip
-        self.port = port
-
-        httpWebDriver = HTTPWebDriver(endpoint: URL(string: "http://\(ip):\(port)")!)
-
-        let path = "\(ProcessInfo.processInfo.environment["ProgramFiles(x86)"]!)\\Windows Application Driver\\WinAppDriver.exe"
-        let commandLine = ["\"\(path)\"", ip, String(port)].joined(separator: " ")
-        try commandLine.withCString(encodedAs: UTF16.self) { commandLine throws in
-            var startupInfo = STARTUPINFOW()
-            startupInfo.cb = DWORD(MemoryLayout<STARTUPINFOW>.size)
-
-            var processInfo = PROCESS_INFORMATION()
-            guard CreateProcessW(
-                nil,
-                UnsafeMutablePointer<WCHAR>(mutating: commandLine),
-                nil,
-                nil,
-                false,
-                DWORD(CREATE_NEW_CONSOLE),
-                nil,
-                nil,
-                &startupInfo,
-                &processInfo
-            ) else {
-                throw WinAppDriverError.win32Error(lastError: Int(GetLastError()))
-            }
+    public init(startingProcess executablePath: String = defaultExecutablePath, ip: String = WinAppDriver.defaultIp, port: Int = WinAppDriver.defaultPort) throws {
+        do {
+            self.processTree = try Win32ProcessTree(path: executablePath, args: [ ip, String(port) ])
+        } catch let error as Win32Error {
+            throw StartError(message: "Call to Win32 \(error.apiName) failed with error code \(error.errorCode).")
+        }
 
-            wadProcessInfo = processInfo
+        // This gives some time for WinAppDriver to get up and running before
+        // we hammer it with requests, otherwise some requests will timeout.
+        Thread.sleep(forTimeInterval: 1.0)
 
-            // This gives some time for WinAppDriver to get up and running before
-            // we hammer it with requests, otherwise some requests will timeout.
-            Thread.sleep(forTimeInterval: 1.0)
-        }
+        self.httpWebDriver = HTTPWebDriver(endpoint: URL(string: "http://\(ip):\(port)")!)
     }
 
     deinit {
-        if let wadProcessInfo {
-            CloseHandle(wadProcessInfo.hThread)
-
-            if !TerminateProcess(wadProcessInfo.hProcess, 0) {
-                let error = GetLastError()
+        if let processTree {
+            do {
+                try processTree.terminate()
+            } catch {
                 assertionFailure("TerminateProcess failed with error \(error).")
             }
-            CloseHandle(wadProcessInfo.hProcess)
 
             // Add a short delay to let process cleanup happen before we try
             // to launch another instance.