From aaa2f32175472a5da42d7712ab1b234e9cfb52d1 Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Fri, 1 May 2026 04:55:59 -0700 Subject: [PATCH] fix(app): retry device tokens on pinned gateways (#75537) --- .../GatewayChannelDeviceTokenRetryTests.swift | 159 ++++++++++++++++++ .../Sources/OpenClawKit/GatewayChannel.swift | 8 +- .../OpenClawKit/GatewayTLSPinning.swift | 10 +- 3 files changed, 173 insertions(+), 4 deletions(-) create mode 100644 apps/macos/Tests/OpenClawIPCTests/GatewayChannelDeviceTokenRetryTests.swift diff --git a/apps/macos/Tests/OpenClawIPCTests/GatewayChannelDeviceTokenRetryTests.swift b/apps/macos/Tests/OpenClawIPCTests/GatewayChannelDeviceTokenRetryTests.swift new file mode 100644 index 00000000000..63710907d39 --- /dev/null +++ b/apps/macos/Tests/OpenClawIPCTests/GatewayChannelDeviceTokenRetryTests.swift @@ -0,0 +1,159 @@ +import Foundation +import OpenClawKit +import Testing + +private extension NSLock { + func withDeviceRetryLock(_ body: () -> T) -> T { + self.lock() + defer { self.unlock() } + return body() + } +} + +private final class ConnectAuthRecorder: @unchecked Sendable { + private let lock = NSLock() + private var auths: [[String: Any]] = [] + + func append(from message: URLSessionWebSocketTask.Message) { + guard let auth = Self.connectAuth(from: message) else { return } + self.lock.withDeviceRetryLock { + self.auths.append(auth) + } + } + + func auth(at index: Int) -> [String: Any]? { + self.lock.withDeviceRetryLock { + guard self.auths.indices.contains(index) else { return nil } + return self.auths[index] + } + } + + private static func connectAuth(from message: URLSessionWebSocketTask.Message) -> [String: Any]? { + let data: Data? = switch message { + case let .data(raw): + raw + case let .string(text): + Data(text.utf8) + @unknown default: + nil + } + guard let data, + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + json["type"] as? String == "req", + json["method"] as? String == "connect", + let params = json["params"] as? [String: Any], + let auth = params["auth"] as? [String: Any] + else { + return nil + } + return auth + } +} + +private final class TrustedDeviceRetryGatewaySession: WebSocketSessioning, GatewayDeviceTokenRetryTrustProviding, @unchecked Sendable { + let allowsDeviceTokenRetryAuth: Bool + + private let lock = NSLock() + private let recorder: ConnectAuthRecorder + private var makeCount = 0 + + init(recorder: ConnectAuthRecorder, allowsDeviceTokenRetryAuth: Bool) { + self.recorder = recorder + self.allowsDeviceTokenRetryAuth = allowsDeviceTokenRetryAuth + } + + func makeWebSocketTask(url: URL) -> WebSocketTaskBox { + _ = url + let attemptIndex = self.lock.withDeviceRetryLock { () -> Int in + let current = self.makeCount + self.makeCount += 1 + return current + } + let recorder = self.recorder + let task = GatewayTestWebSocketTask( + sendHook: { _, message, sendIndex in + if sendIndex == 0 { + recorder.append(from: message) + } + }, + receiveHook: { task, receiveIndex in + if receiveIndex == 0 { + return .data(GatewayWebSocketTestSupport.connectChallengeData()) + } + let id = task.snapshotConnectRequestID() ?? "connect" + if attemptIndex == 0 { + return .data(GatewayWebSocketTestSupport.connectAuthFailureData( + id: id, + detailCode: GatewayConnectAuthDetailCode.authTokenMismatch.rawValue, + canRetryWithDeviceToken: true, + recommendedNextStep: GatewayConnectRecoveryNextStep.retryWithDeviceToken.rawValue)) + } + return .data(GatewayWebSocketTestSupport.connectOkData(id: id)) + }) + return WebSocketTaskBox(task: task) + } +} + +@Suite(.serialized) +struct GatewayChannelDeviceTokenRetryTests { + @Test func `remote pinned TLS retries stale shared token with stored device token`() async throws { + let tempDir = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true) + let previousStateDir = ProcessInfo.processInfo.environment["OPENCLAW_STATE_DIR"] + setenv("OPENCLAW_STATE_DIR", tempDir.path, 1) + defer { + if let previousStateDir { + setenv("OPENCLAW_STATE_DIR", previousStateDir, 1) + } else { + unsetenv("OPENCLAW_STATE_DIR") + } + try? FileManager.default.removeItem(at: tempDir) + } + + let identity = DeviceIdentityStore.loadOrCreate() + _ = DeviceAuthStore.storeToken( + deviceId: identity.deviceId, + role: "operator", + token: "stored-device-token") + + let recorder = ConnectAuthRecorder() + let session = TrustedDeviceRetryGatewaySession( + recorder: recorder, + allowsDeviceTokenRetryAuth: true) + let options = GatewayConnectOptions( + role: "operator", + scopes: ["operator.read"], + caps: [], + commands: [], + permissions: [:], + clientId: "openclaw-ios-test", + clientMode: "ui", + clientDisplayName: "iOS Test", + includeDeviceIdentity: true) + let channel = try GatewayChannelActor( + url: #require(URL(string: "wss://gateway.example.com")), + token: "stale-shared-token", + session: WebSocketSessionBox(session: session), + connectOptions: options) + + do { + try await channel.connect() + Issue.record("expected stale shared-token connect to fail before device-token retry") + } catch let error as GatewayConnectAuthError { + #expect(error.detail == .authTokenMismatch) + } + + try await channel.connect() + + let firstAuth = try #require(recorder.auth(at: 0)) + #expect(firstAuth["token"] as? String == "stale-shared-token") + #expect(firstAuth["deviceToken"] == nil) + + let retryAuth = try #require(recorder.auth(at: 1)) + #expect(retryAuth["token"] as? String == "stale-shared-token") + #expect(retryAuth["deviceToken"] as? String == "stored-device-token") + + await channel.shutdown() + } +} diff --git a/apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayChannel.swift b/apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayChannel.swift index 33722af6dc5..5ec99f5d799 100644 --- a/apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayChannel.swift +++ b/apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayChannel.swift @@ -912,9 +912,6 @@ public actor GatewayChannelActor { } private func isTrustedDeviceRetryEndpoint() -> Bool { - // This client currently treats loopback as the only trusted retry target. - // Unlike the Node gateway client, it does not yet expose a pinned TLS-fingerprint - // trust path for remote retry, so remote fallback remains disabled by default. guard let host = self.url.host?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased(), !host.isEmpty else { @@ -923,6 +920,11 @@ public actor GatewayChannelActor { if host == "localhost" || host == "::1" || host == "127.0.0.1" || host.hasPrefix("127.") { return true } + if self.url.scheme?.lowercased() == "wss", + let trust = self.session as? GatewayDeviceTokenRetryTrustProviding + { + return trust.allowsDeviceTokenRetryAuth + } return false } diff --git a/apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayTLSPinning.swift b/apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayTLSPinning.swift index 8bbbe494f17..c0b45ed27f9 100644 --- a/apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayTLSPinning.swift +++ b/apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayTLSPinning.swift @@ -75,6 +75,10 @@ public protocol GatewayTLSFailureProviding: AnyObject { func consumeLastTLSFailure() -> GatewayTLSValidationFailure? } +public protocol GatewayDeviceTokenRetryTrustProviding: AnyObject { + var allowsDeviceTokenRetryAuth: Bool { get } +} + public enum GatewayTLSStore { private static let keychainService = "ai.openclaw.tls-pinning" @@ -155,7 +159,7 @@ public enum GatewayTLSStore { } } -public final class GatewayTLSPinningSession: NSObject, WebSocketSessioning, URLSessionDelegate, GatewayTLSFailureProviding, @unchecked Sendable { +public final class GatewayTLSPinningSession: NSObject, WebSocketSessioning, URLSessionDelegate, GatewayTLSFailureProviding, GatewayDeviceTokenRetryTrustProviding, @unchecked Sendable { private let params: GatewayTLSParams private let failureLock = NSLock() private var lastTLSFailure: GatewayTLSValidationFailure? @@ -170,6 +174,10 @@ public final class GatewayTLSPinningSession: NSObject, WebSocketSessioning, URLS super.init() } + public var allowsDeviceTokenRetryAuth: Bool { + self.params.expectedFingerprint?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty == false + } + public func consumeLastTLSFailure() -> GatewayTLSValidationFailure? { self.failureLock.lock() defer { self.failureLock.unlock() }