Skip to content

Commit 36db423

Browse files
vincentkocvincentkoc
committed
fix(app): retry device tokens on pinned gateways
1 parent ef0eb12 commit 36db423

3 files changed

Lines changed: 173 additions & 4 deletions

File tree

‎apps/macos/Tests/OpenClawIPCTests/GatewayChannelDeviceTokenRetryTests.swift‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
import Foundation
2+
import OpenClawKit
3+
import Testing
4+
5+
private extension NSLock {
6+
func withDeviceRetryLock<T>(_ body: () -> T) -> T {
7+
self.lock()
8+
defer { self.unlock() }
9+
return body()
10+
}
11+
}
12+
13+
private final class ConnectAuthRecorder: @unchecked Sendable {
14+
private let lock = NSLock()
15+
private var auths: [[String: Any]] = []
16+
17+
func append(from message: URLSessionWebSocketTask.Message) {
18+
guard let auth = Self.connectAuth(from: message) else { return }
19+
self.lock.withDeviceRetryLock {
20+
self.auths.append(auth)
21+
}
22+
}
23+
24+
func auth(at index: Int) -> [String: Any]? {
25+
self.lock.withDeviceRetryLock {
26+
guard self.auths.indices.contains(index) else { return nil }
27+
return self.auths[index]
28+
}
29+
}
30+
31+
private static func connectAuth(from message: URLSessionWebSocketTask.Message) -> [String: Any]? {
32+
let data: Data? = switch message {
33+
case let .data(raw):
34+
raw
35+
case let .string(text):
36+
Data(text.utf8)
37+
@unknown default:
38+
nil
39+
}
40+
guard let data,
41+
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
42+
json["type"] as? String == "req",
43+
json["method"] as? String == "connect",
44+
let params = json["params"] as? [String: Any],
45+
let auth = params["auth"] as? [String: Any]
46+
else {
47+
return nil
48+
}
49+
return auth
50+
}
51+
}
52+
53+
private final class TrustedDeviceRetryGatewaySession: WebSocketSessioning, GatewayDeviceTokenRetryTrustProviding, @unchecked Sendable {
54+
let allowsDeviceTokenRetryAuth: Bool
55+
56+
private let lock = NSLock()
57+
private let recorder: ConnectAuthRecorder
58+
private var makeCount = 0
59+
60+
init(recorder: ConnectAuthRecorder, allowsDeviceTokenRetryAuth: Bool) {
61+
self.recorder = recorder
62+
self.allowsDeviceTokenRetryAuth = allowsDeviceTokenRetryAuth
63+
}
64+
65+
func makeWebSocketTask(url: URL) -> WebSocketTaskBox {
66+
_ = url
67+
let attemptIndex = self.lock.withDeviceRetryLock { () -> Int in
68+
let current = self.makeCount
69+
self.makeCount += 1
70+
return current
71+
}
72+
let recorder = self.recorder
73+
let task = GatewayTestWebSocketTask(
74+
sendHook: { _, message, sendIndex in
75+
if sendIndex == 0 {
76+
recorder.append(from: message)
77+
}
78+
},
79+
receiveHook: { task, receiveIndex in
80+
if receiveIndex == 0 {
81+
return .data(GatewayWebSocketTestSupport.connectChallengeData())
82+
}
83+
let id = task.snapshotConnectRequestID() ?? "connect"
84+
if attemptIndex == 0 {
85+
return .data(GatewayWebSocketTestSupport.connectAuthFailureData(
86+
id: id,
87+
detailCode: GatewayConnectAuthDetailCode.authTokenMismatch.rawValue,
88+
canRetryWithDeviceToken: true,
89+
recommendedNextStep: GatewayConnectRecoveryNextStep.retryWithDeviceToken.rawValue))
90+
}
91+
return .data(GatewayWebSocketTestSupport.connectOkData(id: id))
92+
})
93+
return WebSocketTaskBox(task: task)
94+
}
95+
}
96+
97+
@Suite(.serialized)
98+
struct GatewayChannelDeviceTokenRetryTests {
99+
@Test func `remote pinned TLS retries stale shared token with stored device token`() async throws {
100+
let tempDir = FileManager.default.temporaryDirectory
101+
.appendingPathComponent(UUID().uuidString, isDirectory: true)
102+
try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
103+
let previousStateDir = ProcessInfo.processInfo.environment["OPENCLAW_STATE_DIR"]
104+
setenv("OPENCLAW_STATE_DIR", tempDir.path, 1)
105+
defer {
106+
if let previousStateDir {
107+
setenv("OPENCLAW_STATE_DIR", previousStateDir, 1)
108+
} else {
109+
unsetenv("OPENCLAW_STATE_DIR")
110+
}
111+
try? FileManager.default.removeItem(at: tempDir)
112+
}
113+
114+
let identity = DeviceIdentityStore.loadOrCreate()
115+
_ = DeviceAuthStore.storeToken(
116+
deviceId: identity.deviceId,
117+
role: "operator",
118+
token: "stored-device-token")
119+
120+
let recorder = ConnectAuthRecorder()
121+
let session = TrustedDeviceRetryGatewaySession(
122+
recorder: recorder,
123+
allowsDeviceTokenRetryAuth: true)
124+
let options = GatewayConnectOptions(
125+
role: "operator",
126+
scopes: ["operator.read"],
127+
caps: [],
128+
commands: [],
129+
permissions: [:],
130+
clientId: "openclaw-ios-test",
131+
clientMode: "ui",
132+
clientDisplayName: "iOS Test",
133+
includeDeviceIdentity: true)
134+
let channel = try GatewayChannelActor(
135+
url: #require(URL(string: "wss://gateway.example.com")),
136+
token: "stale-shared-token",
137+
session: WebSocketSessionBox(session: session),
138+
connectOptions: options)
139+
140+
do {
141+
try await channel.connect()
142+
Issue.record("expected stale shared-token connect to fail before device-token retry")
143+
} catch let error as GatewayConnectAuthError {
144+
#expect(error.detail == .authTokenMismatch)
145+
}
146+
147+
try await channel.connect()
148+
149+
let firstAuth = try #require(recorder.auth(at: 0))
150+
#expect(firstAuth["token"] as? String == "stale-shared-token")
151+
#expect(firstAuth["deviceToken"] == nil)
152+
153+
let retryAuth = try #require(recorder.auth(at: 1))
154+
#expect(retryAuth["token"] as? String == "stale-shared-token")
155+
#expect(retryAuth["deviceToken"] as? String == "stored-device-token")
156+
157+
await channel.shutdown()
158+
}
159+
}

‎apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayChannel.swift‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -912,9 +912,6 @@ public actor GatewayChannelActor {
912912
}
913913

914914
private func isTrustedDeviceRetryEndpoint() -> Bool {
915-
// This client currently treats loopback as the only trusted retry target.
916-
// Unlike the Node gateway client, it does not yet expose a pinned TLS-fingerprint
917-
// trust path for remote retry, so remote fallback remains disabled by default.
918915
guard let host = self.url.host?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased(),
919916
!host.isEmpty
920917
else {
@@ -923,6 +920,11 @@ public actor GatewayChannelActor {
923920
if host == "localhost" || host == "::1" || host == "127.0.0.1" || host.hasPrefix("127.") {
924921
return true
925922
}
923+
if self.url.scheme?.lowercased() == "wss",
924+
let trust = self.session as? GatewayDeviceTokenRetryTrustProviding
925+
{
926+
return trust.allowsDeviceTokenRetryAuth
927+
}
926928
return false
927929
}
928930

‎apps/shared/OpenClawKit/Sources/OpenClawKit/GatewayTLSPinning.swift‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,10 @@ public protocol GatewayTLSFailureProviding: AnyObject {
7575
func consumeLastTLSFailure() -> GatewayTLSValidationFailure?
7676
}
7777

78+
public protocol GatewayDeviceTokenRetryTrustProviding: AnyObject {
79+
var allowsDeviceTokenRetryAuth: Bool { get }
80+
}
81+
7882
public enum GatewayTLSStore {
7983
private static let keychainService = "ai.openclaw.tls-pinning"
8084

@@ -155,7 +159,7 @@ public enum GatewayTLSStore {
155159
}
156160
}
157161

158-
public final class GatewayTLSPinningSession: NSObject, WebSocketSessioning, URLSessionDelegate, GatewayTLSFailureProviding, @unchecked Sendable {
162+
public final class GatewayTLSPinningSession: NSObject, WebSocketSessioning, URLSessionDelegate, GatewayTLSFailureProviding, GatewayDeviceTokenRetryTrustProviding, @unchecked Sendable {
159163
private let params: GatewayTLSParams
160164
private let failureLock = NSLock()
161165
private var lastTLSFailure: GatewayTLSValidationFailure?
@@ -170,6 +174,10 @@ public final class GatewayTLSPinningSession: NSObject, WebSocketSessioning, URLS
170174
super.init()
171175
}
172176

177+
public var allowsDeviceTokenRetryAuth: Bool {
178+
self.params.expectedFingerprint?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty == false
179+
}
180+
173181
public func consumeLastTLSFailure() -> GatewayTLSValidationFailure? {
174182
self.failureLock.lock()
175183
defer { self.failureLock.unlock() }

0 commit comments

Comments
 (0)