Desktop: harden auth-token storage and local automation bridge

desktop/CHANGELOG.json

@@ -1,5 +1,9 @@
 {
-  "unreleased": [],
+  "unreleased": [
+    "Hardened sign-in by storing auth tokens in the macOS Keychain instead of UserDefaults",
+    "Cleared the browser extension token on sign-out so a new account cannot reuse the previous browser session",
+    "Restricted the local automation bridge to loopback and required a per-launch bearer token"
+  ],
   "releases": [
     {
       "version": "0.11.319",

desktop/Desktop/Sources/AuthKeychain.swift

@@ -0,0 +1,58 @@
+import Foundation
+import Security
+
+/// Keychain-backed storage for sensitive auth tokens (Firebase ID + refresh tokens).
+///
+/// Items are scoped to the current app bundle, marked
+/// `kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly`, and are NOT included in a
+/// keychain access group — so another app running as the same macOS user cannot
+/// read them the way it can read the app's UserDefaults plist.
+enum AuthKeychain {
+    private static var service: String {
+        (Bundle.main.bundleIdentifier ?? "com.omi.desktop") + ".auth"
+    }
+
+    static func set(_ value: String?, forKey key: String) {
+        let baseQuery: [String: Any] = [
+            kSecClass as String: kSecClassGenericPassword,
+            kSecAttrService as String: service,
+            kSecAttrAccount as String: key,
+        ]
+        SecItemDelete(baseQuery as CFDictionary)
+        guard let value, let data = value.data(using: .utf8) else { return }
+        var addQuery = baseQuery
+        addQuery[kSecValueData as String] = data
+        addQuery[kSecAttrAccessible as String] = kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly
+        let status = SecItemAdd(addQuery as CFDictionary, nil)
+        if status != errSecSuccess {
+            NSLog("OMI AUTH: Keychain add failed for %@ (status=%d)", key, Int(status))
+        }
+    }
+
+    static func get(_ key: String) -> String? {
+        let query: [String: Any] = [
+            kSecClass as String: kSecClassGenericPassword,
+            kSecAttrService as String: service,
+            kSecAttrAccount as String: key,
+            kSecReturnData as String: true,
+            kSecMatchLimit as String: kSecMatchLimitOne,
+        ]
+        var out: AnyObject?
+        let status = SecItemCopyMatching(query as CFDictionary, &out)
+        guard status == errSecSuccess,
+              let data = out as? Data,
+              let value = String(data: data, encoding: .utf8)
+        else {
+            return nil
+        }
+        return value
+    }
+
+    static func clearAll() {
+        let query: [String: Any] = [
+            kSecClass as String: kSecClassGenericPassword,
+            kSecAttrService as String: service,
+        ]
+        SecItemDelete(query as CFDictionary)
+    }
+}

desktop/Desktop/Sources/AuthService.swift

@@ -780,30 +780,51 @@ class AuthService {
     // MARK: - Token Storage
 
     private func saveTokens(idToken: String, refreshToken: String, expiresIn: Int, userId: String) {
-        UserDefaults.standard.set(idToken, forKey: kAuthIdToken)
-        UserDefaults.standard.set(refreshToken, forKey: kAuthRefreshToken)
-        // Store expiry time (current time + expiresIn seconds, minus 5 min buffer)
+        AuthKeychain.set(idToken, forKey: kAuthIdToken)
+        AuthKeychain.set(refreshToken, forKey: kAuthRefreshToken)
+        // Non-sensitive metadata stays in UserDefaults so it remains accessible
+        // for lightweight state checks without hitting Keychain.
         let expiryTime = Date().addingTimeInterval(TimeInterval(expiresIn - 300))
         UserDefaults.standard.set(expiryTime.timeIntervalSince1970, forKey: kAuthTokenExpiry)
-        // Store the user ID that owns these tokens (for validation on retrieval)
         UserDefaults.standard.set(userId, forKey: kAuthTokenUserId)
+        // Defensively wipe any legacy plaintext tokens written by earlier versions.
+        UserDefaults.standard.removeObject(forKey: kAuthIdToken)
+        UserDefaults.standard.removeObject(forKey: kAuthRefreshToken)
         NSLog("OMI AUTH: Saved tokens for user %@, expires at %@", userId, expiryTime.description)
     }
 
     private func clearTokens() {
+        AuthKeychain.set(nil, forKey: kAuthIdToken)
+        AuthKeychain.set(nil, forKey: kAuthRefreshToken)
+        // Purge legacy plaintext copies too.
         UserDefaults.standard.removeObject(forKey: kAuthIdToken)
         UserDefaults.standard.removeObject(forKey: kAuthRefreshToken)
         UserDefaults.standard.removeObject(forKey: kAuthTokenExpiry)
         UserDefaults.standard.removeObject(forKey: kAuthTokenUserId)
         NSLog("OMI AUTH: Cleared all tokens")
     }
 
+    /// Read a token from Keychain. On first access after upgrade, transparently
+    /// migrate any legacy plaintext value out of UserDefaults.
+    private func readStoredToken(forKey key: String) -> String? {
+        if let value = AuthKeychain.get(key), !value.isEmpty {
+            return value
+        }
+        if let legacy = UserDefaults.standard.string(forKey: key), !legacy.isEmpty {
+            AuthKeychain.set(legacy, forKey: key)
+            UserDefaults.standard.removeObject(forKey: key)
+            NSLog("OMI AUTH: Migrated %@ from UserDefaults to Keychain", key)
+            return legacy
+        }
+        return nil
+    }
+
     private var storedIdToken: String? {
-        UserDefaults.standard.string(forKey: kAuthIdToken)
+        readStoredToken(forKey: kAuthIdToken)
     }
 
     private var storedRefreshToken: String? {
-        UserDefaults.standard.string(forKey: kAuthRefreshToken)
+        readStoredToken(forKey: kAuthRefreshToken)
     }
 
     private var storedTokenUserId: String? {
@@ -1107,6 +1128,11 @@ class AuthService {
         // transcriptionEnabled: removeObject works since nothing writes it back.
         UserDefaults.standard.removeObject(forKey: "transcriptionEnabled")
 
+        // Browser automation token binds the app to the user's Chrome session.
+        // Leaving it in place after sign-out would let the next user of this Mac
+        // drive the previous user's still-logged-in browser.
+        UserDefaults.standard.removeObject(forKey: "playwrightExtensionToken")
+
         NSLog("OMI AUTH: Signed out and cleared saved state + onboarding")
     }
 

desktop/Desktop/Sources/DesktopAutomationBridge.swift

@@ -1,6 +1,7 @@
 import AppKit
 import Foundation
 import Network
+import Security
 
 enum DesktopAutomationLaunchOptions {
   static let enableFlag = "--automation-bridge"
@@ -101,16 +102,34 @@ final class DesktopAutomationBridge {
 
   private let queue = DispatchQueue(label: "com.omi.desktop.automation-bridge")
   private var listener: NWListener?
+  private var sessionToken: String?
 
   private init() {}
 
+  /// Path where the per-launch bearer token is written so automation clients
+  /// (agent-swift, test scripts) can read it. File is created with mode 0600
+  /// and rewritten on every app launch.
+  static var tokenFileURL: URL {
+    let fm = FileManager.default
+    let base = fm.urls(for: .applicationSupportDirectory, in: .userDomainMask).first
+      ?? URL(fileURLWithPath: NSHomeDirectory()).appendingPathComponent("Library/Application Support")
+    return base.appendingPathComponent("Omi").appendingPathComponent("automation-bridge.token")
+  }
+
   func startIfNeeded() {
     guard DesktopAutomationLaunchOptions.isEnabled else { return }
     guard listener == nil else { return }
 
     do {
+      let token = try Self.generateSessionToken()
+      try Self.writeTokenFile(token)
+      self.sessionToken = token
+
       let parameters = NWParameters.tcp
       parameters.allowLocalEndpointReuse = true
+      // Bind to loopback only. Without this, NWListener binds to all interfaces,
+      // exposing unauthenticated automation to the local network.
+      parameters.requiredInterfaceType = .loopback
       guard let port = NWEndpoint.Port(rawValue: DesktopAutomationLaunchOptions.port) else {
         log("DesktopAutomationBridge: invalid port \(DesktopAutomationLaunchOptions.port)")
         return
@@ -125,18 +144,75 @@ final class DesktopAutomationBridge {
       listener.start(queue: queue)
       self.listener = listener
       log(
-        "DesktopAutomationBridge: listening on http://127.0.0.1:\(DesktopAutomationLaunchOptions.port)"
+        "DesktopAutomationBridge: listening on http://127.0.0.1:\(DesktopAutomationLaunchOptions.port) (token at \(Self.tokenFileURL.path))"
       )
     } catch {
       logError("DesktopAutomationBridge: failed to start listener", error: error)
     }
   }
 
   private func handleConnection(_ connection: NWConnection) {
+    // Defense in depth: even with requiredInterfaceType = .loopback, verify the
+    // remote endpoint is a loopback address before processing any data.
+    if !Self.isLoopbackEndpoint(connection.endpoint) {
+      log("DesktopAutomationBridge: rejected non-loopback connection from \(connection.endpoint)")
+      connection.cancel()
+      return
+    }
     connection.start(queue: queue)
     receiveRequest(on: connection, buffer: Data())
   }
 
+  private static func isLoopbackEndpoint(_ endpoint: NWEndpoint) -> Bool {
+    switch endpoint {
+    case .hostPort(let host, _):
+      switch host {
+      case .ipv4(let addr):
+        return addr.isLoopback
+      case .ipv6(let addr):
+        return addr.isLoopback
+      case .name(let name, _):
+        return name == "localhost" || name == "127.0.0.1" || name == "::1"
+      @unknown default:
+        return false
+      }
+    default:
+      return false
+    }
+  }
+
+  private static func generateSessionToken() throws -> String {
+    var bytes = [UInt8](repeating: 0, count: 32)
+    let status = SecRandomCopyBytes(kSecRandomDefault, bytes.count, &bytes)
+    guard status == errSecSuccess else {
+      throw NSError(
+        domain: "DesktopAutomationBridge", code: Int(status),
+        userInfo: [NSLocalizedDescriptionKey: "SecRandomCopyBytes failed"])
+    }
+    return Data(bytes).base64EncodedString()
+      .replacingOccurrences(of: "+", with: "-")
+      .replacingOccurrences(of: "/", with: "_")
+      .replacingOccurrences(of: "=", with: "")
+  }
+
+  private static func writeTokenFile(_ token: String) throws {
+    let url = tokenFileURL
+    let fm = FileManager.default
+    try fm.createDirectory(
+      at: url.deletingLastPathComponent(), withIntermediateDirectories: true,
+      attributes: [.posixPermissions: 0o700])
+    try token.write(to: url, atomically: true, encoding: .utf8)
+    try fm.setAttributes([.posixPermissions: 0o600], ofItemAtPath: url.path)
+  }
+
+  private func extractBearerToken(from headers: [String: String]) -> String? {
+    guard let raw = headers["authorization"] else { return nil }
+    let trimmed = raw.trimmingCharacters(in: .whitespaces)
+    let prefix = "bearer "
+    guard trimmed.lowercased().hasPrefix(prefix) else { return nil }
+    return String(trimmed.dropFirst(prefix.count)).trimmingCharacters(in: .whitespaces)
+  }
+
   private func receiveRequest(on connection: NWConnection, buffer: Data) {
     connection.receive(minimumIncompleteLength: 1, maximumLength: 64 * 1024) {
       [weak self] data, _, isComplete, error in
@@ -189,11 +265,15 @@ final class DesktopAutomationBridge {
     let path = String(requestParts[1])
 
     var contentLength = 0
+    var headers: [String: String] = [:]
     for line in lines.dropFirst() {
       let pieces = line.split(separator: ":", maxSplits: 1)
       guard pieces.count == 2 else { continue }
-      if pieces[0].lowercased() == "content-length" {
-        contentLength = Int(pieces[1].trimmingCharacters(in: .whitespaces)) ?? 0
+      let name = pieces[0].lowercased()
+      let value = pieces[1].trimmingCharacters(in: .whitespaces)
+      headers[name] = value
+      if name == "content-length" {
+        contentLength = Int(value) ?? 0
       }
     }
 
@@ -204,10 +284,38 @@ final class DesktopAutomationBridge {
     }
 
     let body = Data(data[bodyStart..<data.index(bodyStart, offsetBy: contentLength)])
-    return HTTPRequest(method: method, path: path, body: body)
+    return HTTPRequest(method: method, path: path, headers: headers, body: body)
+  }
+
+  private static func secureCompare(_ a: String, _ b: String) -> Bool {
+    let ab = Array(a.utf8)
+    let bb = Array(b.utf8)
+    if ab.count != bb.count { return false }
+    var diff: UInt8 = 0
+    for i in 0..<ab.count { diff |= ab[i] ^ bb[i] }
+    return diff == 0
   }
 
   private func route(request: HTTPRequest) async -> HTTPResponse {
+    // Every route (including /health) requires the per-launch bearer token.
+    // Without this any local process — or any LAN host if the listener ever
+    // falls back to all-interfaces — could trigger /gmail-read on the
+    // currently signed-in user and walk away with their inbox.
+    guard let expected = sessionToken else {
+      return jsonResponse(
+        DesktopAutomationResponse<DesktopAutomationSnapshot>(
+          ok: false, result: nil, error: "bridge_not_ready"),
+        statusCode: 503)
+    }
+    guard let provided = extractBearerToken(from: request.headers),
+      Self.secureCompare(provided, expected)
+    else {
+      return jsonResponse(
+        DesktopAutomationResponse<DesktopAutomationSnapshot>(
+          ok: false, result: nil, error: "unauthorized"),
+        statusCode: 401)
+    }
+
     switch (request.method, request.path) {
     case ("GET", "/health"):
       let snapshot = await DesktopAutomationStateStore.shared.current()
@@ -370,7 +478,9 @@ final class DesktopAutomationBridge {
     switch response.statusCode {
     case 200: statusText = "OK"
     case 400: statusText = "Bad Request"
+    case 401: statusText = "Unauthorized"
     case 404: statusText = "Not Found"
+    case 503: statusText = "Service Unavailable"
     default: statusText = "Internal Server Error"
     }
 
@@ -399,6 +509,7 @@ final class DesktopAutomationBridge {
 private struct HTTPRequest {
   let method: String
   let path: String
+  let headers: [String: String]
   let body: Data
 }