From 9ab374c65c58757ba5bb11a074bc9106310535ad Mon Sep 17 00:00:00 2001 From: Thomas Rademaker Date: Thu, 12 Mar 2026 15:38:08 -0400 Subject: [PATCH] Add auth proxy support to CoreATProtocol for confidential client OAuth --- .../CoreATProtocol/OAuth/ATProtoOAuth.swift | 74 ++++++- Sources/CoreATProtocol/OAuth/AuthProxy.swift | 168 +++++++++++++++ Tests/CoreATProtocolTests/OAuthTests.swift | 198 ++++++++++++++++++ 3 files changed, 433 insertions(+), 7 deletions(-) create mode 100644 Sources/CoreATProtocol/OAuth/AuthProxy.swift diff --git a/Sources/CoreATProtocol/OAuth/ATProtoOAuth.swift b/Sources/CoreATProtocol/OAuth/ATProtoOAuth.swift index e6726e0..ad7b4cb 100644 --- a/Sources/CoreATProtocol/OAuth/ATProtoOAuth.swift +++ b/Sources/CoreATProtocol/OAuth/ATProtoOAuth.swift @@ -25,15 +25,18 @@ public struct ATProtoOAuthConfig: Sendable { public let clientMetadataURL: String public let redirectURI: String public let scopes: [String] + public let authProxyBaseURL: String? public init( clientMetadataURL: String, redirectURI: String, - scopes: [String] = ["atproto", "transition:generic"] + scopes: [String] = ["atproto", "transition:generic"], + authProxyBaseURL: String? = nil ) { self.clientMetadataURL = clientMetadataURL self.redirectURI = redirectURI self.scopes = scopes + self.authProxyBaseURL = authProxyBaseURL } } @@ -43,17 +46,23 @@ public struct ATProtoAuthStorage: Sendable { public let storeLogin: @Sendable (Login) async throws -> Void public let retrievePrivateKey: @Sendable () async throws -> Data? public let storePrivateKey: @Sendable (Data) async throws -> Void + public let retrieveAuthProxyKeyID: (@Sendable () async throws -> String?)? + public let storeAuthProxyKeyID: (@Sendable (String) async throws -> Void)? public init( retrieveLogin: @escaping @Sendable () async throws -> Login?, storeLogin: @escaping @Sendable (Login) async throws -> Void, retrievePrivateKey: @escaping @Sendable () async throws -> Data?, - storePrivateKey: @escaping @Sendable (Data) async throws -> Void + storePrivateKey: @escaping @Sendable (Data) async throws -> Void, + retrieveAuthProxyKeyID: (@Sendable () async throws -> String?)? = nil, + storeAuthProxyKeyID: (@Sendable (String) async throws -> Void)? = nil ) { self.retrieveLogin = retrieveLogin self.storeLogin = storeLogin self.retrievePrivateKey = retrievePrivateKey self.storePrivateKey = storePrivateKey + self.retrieveAuthProxyKeyID = retrieveAuthProxyKeyID + self.storeAuthProxyKeyID = storeAuthProxyKeyID } } @@ -221,7 +230,27 @@ public final class ATProtoOAuth: Sendable { try await self.generateJWT(params: params) } - // Step 7: Create authenticator + // Step 7: Create proxy URL loader if auth proxy is configured + let proxyKeyIDStorage: AuthProxyKeyIDStorage? + let proxyURLLoader: URLResponseProvider? + + if let proxyBaseURL = config.authProxyBaseURL { + let initialKeyID = try? await storage.retrieveAuthProxyKeyID?() + let keyIDStorage = AuthProxyKeyIDStorage(keyID: initialKeyID) + proxyKeyIDStorage = keyIDStorage + proxyURLLoader = makeProxyURLLoader( + proxyBaseURL: proxyBaseURL, + tokenEndpoint: serverConfig.tokenEndpoint, + parEndpoint: serverConfig.pushedAuthorizationRequestEndpoint, + issuer: serverConfig.issuer, + keyIDStorage: keyIDStorage + ) + } else { + proxyKeyIDStorage = nil + proxyURLLoader = nil + } + + // Step 8: Create authenticator let tokenHandling = buildTokenHandling( accountHint: identifier, server: serverConfig, @@ -238,9 +267,9 @@ public final class ATProtoOAuth: Sendable { userAuthenticator: userAuthenticator ) - let authenticator = Authenticator(config: authenticatorConfig) + let authenticator = Authenticator(config: authenticatorConfig, urlLoader: proxyURLLoader) - // Step 8: Trigger authentication with user interaction + // Step 9: Trigger authentication with user interaction let login: Login do { login = try await authenticator.authenticate() @@ -258,7 +287,12 @@ public final class ATProtoOAuth: Sendable { } } - // Step 9: Setup CoreATProtocol environment + // Step 10: Persist auth proxy key ID + if let proxyKeyIDStorage, let keyID = await proxyKeyIDStorage.keyID { + try? await storage.storeAuthProxyKeyID?(keyID) + } + + // Step 11: Setup CoreATProtocol environment setup( hostURL: identity.pdsEndpoint, accessJWT: login.accessToken.value, @@ -345,17 +379,43 @@ public final class ATProtoOAuth: Sendable { return nil } + // Use proxy-aware provider when auth proxy is configured + let baseProvider: URLResponseProvider + let proxyKeyIDStorage: AuthProxyKeyIDStorage? + + if let proxyBaseURL = config.authProxyBaseURL { + let initialKeyID = try? await storage.retrieveAuthProxyKeyID?() + let keyIDStorage = AuthProxyKeyIDStorage(keyID: initialKeyID) + proxyKeyIDStorage = keyIDStorage + baseProvider = makeProxyURLLoader( + proxyBaseURL: proxyBaseURL, + tokenEndpoint: serverConfig.tokenEndpoint, + parEndpoint: serverConfig.pushedAuthorizationRequestEndpoint, + issuer: serverConfig.issuer, + keyIDStorage: keyIDStorage + ) + } else { + baseProvider = provider + proxyKeyIDStorage = nil + } + let responseProvider: URLResponseProvider = { request in try await self.dpopRequestActor.response( request: request, jwtGenerator: jwtGenerator, - provider: provider, + provider: baseProvider, issuingServer: issuer ) } let refreshedLogin = try await refreshProvider(login, clientConfig.credentials, responseProvider) try await storage.storeLogin(refreshedLogin) + + // Persist updated auth proxy key ID + if let proxyKeyIDStorage, let keyID = await proxyKeyIDStorage.keyID { + try? await storage.storeAuthProxyKeyID?(keyID) + } + return refreshedLogin } diff --git a/Sources/CoreATProtocol/OAuth/AuthProxy.swift b/Sources/CoreATProtocol/OAuth/AuthProxy.swift new file mode 100644 index 0000000..22b759d --- /dev/null +++ b/Sources/CoreATProtocol/OAuth/AuthProxy.swift @@ -0,0 +1,168 @@ +import Foundation +import OAuthenticator + +// MARK: - Auth Proxy Request Models + +/// Request body for proxying token exchange/refresh through the auth proxy. +struct AuthProxyTokenRequest: Encodable { + let tokenEndpoint: String + let issuer: String + let grantType: String + let code: String? + let redirectURI: String? + let codeVerifier: String? + let refreshToken: String? + let keyID: String? + + enum CodingKeys: String, CodingKey { + case tokenEndpoint = "token_endpoint" + case issuer + case grantType = "grant_type" + case code + case redirectURI = "redirect_uri" + case codeVerifier = "code_verifier" + case refreshToken = "refresh_token" + case keyID = "key_id" + } +} + +/// Request body for proxying PAR requests through the auth proxy. +struct AuthProxyPARRequest: Encodable { + let parEndpoint: String + let issuer: String + let keyID: String? + let loginHint: String? + let scope: String + let codeChallenge: String + let codeChallengeMethod: String + let state: String + let redirectURI: String + + enum CodingKeys: String, CodingKey { + case parEndpoint = "par_endpoint" + case issuer + case keyID = "key_id" + case loginHint = "login_hint" + case scope + case codeChallenge = "code_challenge" + case codeChallengeMethod = "code_challenge_method" + case state + case redirectURI = "redirect_uri" + } +} + +// MARK: - Key ID Storage + +/// Thread-safe storage for the auth proxy key ID, used during a single auth flow. +actor AuthProxyKeyIDStorage { + var keyID: String? + + init(keyID: String? = nil) { + self.keyID = keyID + } + + func update(_ newKeyID: String) { + keyID = newKeyID + } +} + +// MARK: - Proxy URL Loader + +/// Creates a `URLResponseProvider` that intercepts requests to the auth server's PAR and token +/// endpoints, redirecting them through the auth proxy. DPoP proofs are generated by the caller +/// for the real endpoint URLs and forwarded unchanged — the proxy is transparent to DPoP. +func makeProxyURLLoader( + proxyBaseURL: String, + tokenEndpoint: String, + parEndpoint: String, + issuer: String, + keyIDStorage: AuthProxyKeyIDStorage +) -> URLResponseProvider { + { request in + guard let requestURL = request.url?.absoluteString else { + return try await URLSession.shared.data(for: request) + } + + let isTokenRequest = requestURL == tokenEndpoint + let isPARRequest = requestURL == parEndpoint + + guard isTokenRequest || isPARRequest else { + return try await URLSession.shared.data(for: request) + } + + let proxyPath = isTokenRequest ? "/oauth/token" : "/oauth/par" + guard let proxyURL = URL(string: proxyBaseURL + proxyPath) else { + return try await URLSession.shared.data(for: request) + } + + var proxyRequest = URLRequest(url: proxyURL) + proxyRequest.httpMethod = "POST" + proxyRequest.setValue("application/json", forHTTPHeaderField: "Content-Type") + proxyRequest.setValue("application/json", forHTTPHeaderField: "Accept") + + // Forward DPoP header (already generated for the real endpoint URL) + if let dpop = request.value(forHTTPHeaderField: "DPoP") { + proxyRequest.setValue(dpop, forHTTPHeaderField: "DPoP") + } + + let currentKeyID = await keyIDStorage.keyID + + if isPARRequest { + // Convert form-encoded PAR body to JSON for the proxy + let formParams = parseFormEncoded(request.httpBody) + let parBody = AuthProxyPARRequest( + parEndpoint: parEndpoint, + issuer: issuer, + keyID: currentKeyID, + loginHint: formParams["login_hint"], + scope: formParams["scope"] ?? "", + codeChallenge: formParams["code_challenge"] ?? "", + codeChallengeMethod: formParams["code_challenge_method"] ?? "", + state: formParams["state"] ?? "", + redirectURI: formParams["redirect_uri"] ?? "" + ) + proxyRequest.httpBody = try JSONEncoder().encode(parBody) + } else { + // Token request: add proxy-specific fields to the existing JSON body + var bodyDict: [String: Any] = [:] + if let body = request.httpBody { + bodyDict = (try? JSONSerialization.jsonObject(with: body) as? [String: Any]) ?? [:] + } + bodyDict["token_endpoint"] = tokenEndpoint + bodyDict["issuer"] = issuer + if let keyID = currentKeyID { + bodyDict["key_id"] = keyID + } + proxyRequest.httpBody = try JSONSerialization.data(withJSONObject: bodyDict) + } + + let (data, response) = try await URLSession.shared.data(for: proxyRequest) + + // Capture Auth-Proxy-Key-ID from every proxy response + if let httpResponse = response as? HTTPURLResponse, + let newKeyID = httpResponse.value(forHTTPHeaderField: "Auth-Proxy-Key-ID") { + await keyIDStorage.update(newKeyID) + } + + return (data, response) + } +} + +// MARK: - Helpers + +/// Parses a `application/x-www-form-urlencoded` body into a dictionary. +func parseFormEncoded(_ data: Data?) -> [String: String] { + guard let data, let string = String(data: data, encoding: .utf8) else { + return [:] + } + + var result: [String: String] = [:] + for pair in string.split(separator: "&") { + let parts = pair.split(separator: "=", maxSplits: 1) + guard parts.count == 2, + let key = parts[0].removingPercentEncoding, + let value = parts[1].removingPercentEncoding else { continue } + result[key] = value + } + return result +} diff --git a/Tests/CoreATProtocolTests/OAuthTests.swift b/Tests/CoreATProtocolTests/OAuthTests.swift index d0a62aa..8454dea 100644 --- a/Tests/CoreATProtocolTests/OAuthTests.swift +++ b/Tests/CoreATProtocolTests/OAuthTests.swift @@ -74,6 +74,204 @@ struct IdentityResolverTests { } } +@Suite("Auth Proxy Models") +struct AuthProxyTests { + + @Test("ATProtoOAuthConfig with auth proxy base URL") + func testConfigWithProxy() { + let config = ATProtoOAuthConfig( + clientMetadataURL: "https://example.com/client-metadata.json", + redirectURI: "example://callback", + authProxyBaseURL: "https://auth.example.com" + ) + + #expect(config.authProxyBaseURL == "https://auth.example.com") + } + + @Test("ATProtoOAuthConfig defaults to nil auth proxy base URL") + func testConfigWithoutProxy() { + let config = ATProtoOAuthConfig( + clientMetadataURL: "https://example.com/client-metadata.json", + redirectURI: "example://callback" + ) + + #expect(config.authProxyBaseURL == nil) + } + + @Test("AuthProxyTokenRequest encodes correctly for authorization code") + func testTokenRequestEncoding() throws { + let request = AuthProxyTokenRequest( + tokenEndpoint: "https://bsky.social/oauth/token", + issuer: "https://bsky.social", + grantType: "authorization_code", + code: "test_code", + redirectURI: "example://callback", + codeVerifier: "test_verifier", + refreshToken: nil, + keyID: "atproto-auth-1" + ) + + let data = try JSONEncoder().encode(request) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["token_endpoint"] as? String == "https://bsky.social/oauth/token") + #expect(json?["issuer"] as? String == "https://bsky.social") + #expect(json?["grant_type"] as? String == "authorization_code") + #expect(json?["code"] as? String == "test_code") + #expect(json?["redirect_uri"] as? String == "example://callback") + #expect(json?["code_verifier"] as? String == "test_verifier") + #expect(json?["key_id"] as? String == "atproto-auth-1") + #expect(json?["refresh_token"] == nil) + } + + @Test("AuthProxyTokenRequest encodes correctly for refresh") + func testRefreshTokenRequestEncoding() throws { + let request = AuthProxyTokenRequest( + tokenEndpoint: "https://bsky.social/oauth/token", + issuer: "https://bsky.social", + grantType: "refresh_token", + code: nil, + redirectURI: nil, + codeVerifier: nil, + refreshToken: "test_refresh_token", + keyID: "atproto-auth-2" + ) + + let data = try JSONEncoder().encode(request) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["grant_type"] as? String == "refresh_token") + #expect(json?["refresh_token"] as? String == "test_refresh_token") + #expect(json?["key_id"] as? String == "atproto-auth-2") + #expect(json?["code"] == nil) + } + + @Test("AuthProxyPARRequest encodes correctly") + func testPARRequestEncoding() throws { + let request = AuthProxyPARRequest( + parEndpoint: "https://bsky.social/oauth/par", + issuer: "https://bsky.social", + keyID: "atproto-auth-1", + loginHint: "alice.bsky.social", + scope: "atproto transition:generic", + codeChallenge: "test_challenge", + codeChallengeMethod: "S256", + state: "test_state", + redirectURI: "example://callback" + ) + + let data = try JSONEncoder().encode(request) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["par_endpoint"] as? String == "https://bsky.social/oauth/par") + #expect(json?["issuer"] as? String == "https://bsky.social") + #expect(json?["key_id"] as? String == "atproto-auth-1") + #expect(json?["login_hint"] as? String == "alice.bsky.social") + #expect(json?["scope"] as? String == "atproto transition:generic") + #expect(json?["code_challenge"] as? String == "test_challenge") + #expect(json?["code_challenge_method"] as? String == "S256") + #expect(json?["state"] as? String == "test_state") + #expect(json?["redirect_uri"] as? String == "example://callback") + } + + @Test("AuthProxyPARRequest encodes nil key ID and login hint") + func testPARRequestWithNils() throws { + let request = AuthProxyPARRequest( + parEndpoint: "https://bsky.social/oauth/par", + issuer: "https://bsky.social", + keyID: nil, + loginHint: nil, + scope: "atproto", + codeChallenge: "challenge", + codeChallengeMethod: "S256", + state: "state", + redirectURI: "example://callback" + ) + + let data = try JSONEncoder().encode(request) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["key_id"] == nil) + #expect(json?["login_hint"] == nil) + } + + @Test("parseFormEncoded handles standard form body") + func testParseFormEncoded() { + let body = "client_id=test&scope=atproto&redirect_uri=example%3A%2F%2Fcallback&state=abc123" + let result = parseFormEncoded(body.data(using: .utf8)) + + #expect(result["client_id"] == "test") + #expect(result["scope"] == "atproto") + #expect(result["redirect_uri"] == "example://callback") + #expect(result["state"] == "abc123") + } + + @Test("parseFormEncoded handles empty and nil data") + func testParseFormEncodedEdgeCases() { + #expect(parseFormEncoded(nil).isEmpty) + #expect(parseFormEncoded(Data()).isEmpty) + } + + @Test("AuthProxyKeyIDStorage stores and retrieves key ID") + func testKeyIDStorage() async { + let storage = AuthProxyKeyIDStorage() + let initial = await storage.keyID + #expect(initial == nil) + + await storage.update("key-1") + let updated = await storage.keyID + #expect(updated == "key-1") + + await storage.update("key-2") + let rotated = await storage.keyID + #expect(rotated == "key-2") + } + + @Test("AuthProxyKeyIDStorage initializes with existing key ID") + func testKeyIDStorageWithInitialValue() async { + let storage = AuthProxyKeyIDStorage(keyID: "existing-key") + let value = await storage.keyID + #expect(value == "existing-key") + } + + @Test("ATProtoAuthStorage accepts auth proxy key ID closures") + func testStorageWithProxyKeyID() async throws { + let keyIDHolder = AuthProxyKeyIDStorage() + + let storage = ATProtoAuthStorage( + retrieveLogin: { nil }, + storeLogin: { _ in }, + retrievePrivateKey: { nil }, + storePrivateKey: { _ in }, + retrieveAuthProxyKeyID: { await keyIDHolder.keyID }, + storeAuthProxyKeyID: { await keyIDHolder.update($0) } + ) + + #expect(storage.retrieveAuthProxyKeyID != nil) + #expect(storage.storeAuthProxyKeyID != nil) + + let initial = try await storage.retrieveAuthProxyKeyID?() + #expect(initial == nil) + + try await storage.storeAuthProxyKeyID?("test-key") + let retrieved = try await storage.retrieveAuthProxyKeyID?() + #expect(retrieved == "test-key") + } + + @Test("ATProtoAuthStorage defaults proxy key ID closures to nil") + func testStorageWithoutProxyKeyID() { + let storage = ATProtoAuthStorage( + retrieveLogin: { nil }, + storeLogin: { _ in }, + retrievePrivateKey: { nil }, + storePrivateKey: { _ in } + ) + + #expect(storage.retrieveAuthProxyKeyID == nil) + #expect(storage.storeAuthProxyKeyID == nil) + } +} + @Suite("DPoP JWT") struct DPoPTests { -- 2.51.2