Handle Ctrl + C properly (#73)

* Terminate when Ctrl+C is encountered when entering credentials * Handle Ctrl+C by catching SIGINT and converting it to task cancellation * maxCharacters instead of (buf.count - 2) * readStdinCredential: fix maxCharacters to be 255 * "user" variable should be named "credential"
cirruslabs · May 16, 2022 · 60b7054 · 60b7054
1 parent fa9e314
commit 60b7054
Show file tree

Hide file tree

Showing 4 changed files with 102 additions and 43 deletions.
diff --git a/Sources/tart/Commands/Run.swift b/Sources/tart/Commands/Run.swift
@@ -18,50 +18,55 @@ struct Run: AsyncParsableCommand {
     let vmDir = try VMStorageLocal().open(name)
     vm = try VM(vmDir: vmDir)
 
-    Task {
-      do {
-        try await vm!.run()
+    await withThrowingTaskGroup(of: Void.self) { group in
+      group.addTask {
+        do {
+          try await vm!.run()
 
-        Foundation.exit(0)
-      } catch {
-        print(error)
+          Foundation.exit(0)
+        } catch {
+          print(error)
 
-        Foundation.exit(1)
+          Foundation.exit(1)
+        }
       }
-    }
 
-    if noGraphics {
-      dispatchMain()
-    } else {
-      // UI mumbo-jumbo
-      let nsApp = NSApplication.shared
-      nsApp.setActivationPolicy(.regular)
-      nsApp.activate(ignoringOtherApps: true)
-
-      nsApp.applicationIconImage = NSImage(data: AppIconData)
-
-      struct MainApp: App {
-        var body: some Scene {
-          WindowGroup(vm!.name) {
-            Group {
-              VMView(vm: vm!).onAppear {
-                NSWindow.allowsAutomaticWindowTabbing = false
-              }
-            }.frame(width: CGFloat(vm!.config.display.width), height: CGFloat(vm!.config.display.height))
-          }.commands {
-            // Remove some standard menu options
-            CommandGroup(replacing: .help, addition: {})
-            CommandGroup(replacing: .newItem, addition: {})
-            CommandGroup(replacing: .pasteboard, addition: {})
-            CommandGroup(replacing: .textEditing, addition: {})
-            CommandGroup(replacing: .undoRedo, addition: {})
-            CommandGroup(replacing: .windowSize, addition: {})
-          }
-        }
+      if noGraphics {
+        dispatchMain()
+      } else {
+        runUI()
       }
+    }
+  }
 
-      MainApp.main()
+  private func runUI() {
+    let nsApp = NSApplication.shared
+    nsApp.setActivationPolicy(.regular)
+    nsApp.activate(ignoringOtherApps: true)
+
+    nsApp.applicationIconImage = NSImage(data: AppIconData)
+
+    struct MainApp: App {
+      var body: some Scene {
+        WindowGroup(vm!.name) {
+          Group {
+            VMView(vm: vm!).onAppear {
+              NSWindow.allowsAutomaticWindowTabbing = false
+            }
+          }.frame(width: CGFloat(vm!.config.display.width), height: CGFloat(vm!.config.display.height))
+        }.commands {
+                  // Remove some standard menu options
+                  CommandGroup(replacing: .help, addition: {})
+                  CommandGroup(replacing: .newItem, addition: {})
+                  CommandGroup(replacing: .pasteboard, addition: {})
+                  CommandGroup(replacing: .textEditing, addition: {})
+                  CommandGroup(replacing: .undoRedo, addition: {})
+                  CommandGroup(replacing: .windowSize, addition: {})
+                }
+      }
     }
+
+    MainApp.main()
   }
 }
 

diff --git a/Sources/tart/Credentials.swift b/Sources/tart/Credentials.swift
@@ -1,5 +1,10 @@
 import Foundation
 
+enum CredentialsError: Error {
+  case CredentialRequired(which: String)
+  case CredentialTooLong(message: String)
+}
+
 class Credentials {
   static func retrieveKeychain(host: String) throws -> (String, String)? {
     let query: [String: Any] = [kSecClass as String: kSecClassInternetPassword,
@@ -34,13 +39,26 @@ class Credentials {
   }
 
   static func retrieveStdin() throws -> (String, String) {
-    print("User: ", terminator: "")
-    let user = readLine() ?? ""
+    let user = try readStdinCredential(name: "username", prompt: "User: ", isSensitive: false)
+    let password = try readStdinCredential(name: "password", prompt: "Password: ", isSensitive: true)
+
+    return (user, password)
+  }
 
-    let rawPass = getpass("Password: ")
-    let pass = String(cString: rawPass!, encoding: .utf8)!
+  private static func readStdinCredential(name: String, prompt: String, maxCharacters: Int = 255, isSensitive: Bool) throws -> String {
+    var buf = [CChar](repeating: 0, count: maxCharacters + 1 /* sentinel */ + 1 /* NUL */)
+    guard let rawCredential = readpassphrase(prompt, &buf, buf.count, isSensitive ? RPP_ECHO_OFF : RPP_ECHO_ON) else {
+      throw CredentialsError.CredentialRequired(which: name)
+    }
+
+    let credential = String(cString: rawCredential).trimmingCharacters(in: .newlines)
+
+    if credential.count > maxCharacters {
+      throw CredentialsError.CredentialTooLong(
+        message: "\(name) should contain no more than \(maxCharacters) characters")
+    }
 
-    return (user, pass)
+    return credential
   }
 
   static func store(host: String, user: String, password: String) throws {

diff --git a/Sources/tart/Root.swift b/Sources/tart/Root.swift
@@ -1,4 +1,5 @@
 import ArgumentParser
+import Foundation
 
 @main
 struct Root: AsyncParsableCommand {
@@ -16,4 +17,27 @@ struct Root: AsyncParsableCommand {
       Push.self,
       Delete.self,
     ])
+
+  public static func main() async throws {
+    // Handle cancellation by Ctrl+C
+    let task = withUnsafeCurrentTask { $0 }!
+    let sigintSrc = DispatchSource.makeSignalSource(signal: SIGINT)
+    sigintSrc.setEventHandler {
+      task.cancel()
+    }
+    sigintSrc.activate()
+
+    // Parse and run command
+    do {
+      var command = try parseAsRoot()
+
+      if var asyncCommand = command as? AsyncParsableCommand {
+        try await asyncCommand.run()
+      } else {
+        try command.run()
+      }
+    } catch {
+      exit(withError: error)
+    }
+  }
 }
diff --git a/Sources/tart/VM.swift b/Sources/tart/VM.swift
@@ -140,7 +140,19 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
       }
     }
 
-    sema.wait()
+    await withTaskCancellationHandler(operation: {
+      sema.wait()
+    }, onCancel: {
+      sema.signal()
+    })
+
+    if Task.isCancelled {
+      DispatchQueue.main.sync {
+        Task {
+          try await self.virtualMachine.stop()
+        }
+      }
+    }
   }
 
   static func craftConfiguration(diskURL: URL, auxStorage: VZMacAuxiliaryStorage, vmConfig: VMConfig) throws -> VZVirtualMachineConfiguration {