oauth wip

Package.resolved

···
 {
-  "originHash" : "aef8443d2c26c1b290a85bc86c844a258c5dd2c9b4979e9f7b3da92cf56bb581",
+  "originHash" : "76f98c0d52def5e7d5c79516d2e7a4fdfa68c70d7d07bb400dd0095cef1c0b26",
   "pins" : [
     {
       "identity" : "jwt-kit",
···
       "state" : {
         "revision" : "b5f82fb9dc238f2fcac53d721a222513a152613c",
         "version" : "5.3.0"
-      }
-    },
-    {
-      "identity" : "oauthenticator",
-      "kind" : "remoteSourceControl",
-      "location" : "https://github.com/radmakr/OAuthenticator.git",
-      "state" : {
-        "branch" : "CoreAtProtocol",
-        "revision" : "e382a28c7f7dbdb36ec358b1324e0d2320249c70"
       }
     },
     {

Package.swift

···
         // Using fork with fix for WebAuthenticationSession platform guards
         // PR pending at https://github.com/ChimeHQ/OAuthenticator
 //        .package(url: "https://github.com/ChimeHQ/OAuthenticator.git", branch: "main"),
-        .package(url: "https://github.com/radmakr/OAuthenticator.git", branch: "CoreAtProtocol"),
+//        .package(url: "https://github.com/radmakr/OAuthenticator.git", branch: "CoreAtProtocol"),
+        .package(path: "../OAuthenticator"),
         .package(url: "https://github.com/vapor/jwt-kit.git", from: "5.0.0"),
     ],
     targets: [

Sources/CoreATProtocol/APEnvironment.swift

···
     public var accessToken: String?
     public var refreshToken: String?
     public var atProtocoldelegate: CoreATProtocolDelegate?
+    public var tokenRefreshHandler: (@Sendable () async throws -> Bool)?
     public let routerDelegate = APRouterDelegate()
     
     private init() {}
···
 //        self.userAgent = userAgent
 //    }
 }
-

Sources/CoreATProtocol/CoreATProtocol.swift

···
 }
 
 @APActor
+public func setTokenRefreshHandler(_ handler: (@Sendable () async throws -> Bool)?) {
+    APEnvironment.current.tokenRefreshHandler = handler
+}
+
+@APActor
 public func updateTokens(access: String?, refresh: String?) {
     APEnvironment.current.accessToken = access
     APEnvironment.current.refreshToken = refresh

Sources/CoreATProtocol/Networking.swift

···
 @APActor
 public class APRouterDelegate: NetworkRouterDelegate {
     private var shouldRefreshToken = false
+    private var refreshTask: Task<Bool, Error>?
     
     public func intercept(_ request: inout URLRequest) async {
         if let refreshToken = APEnvironment.current.refreshToken, shouldRefreshToken {
···
     }
     
     public func shouldRetry(error: Error, attempts: Int) async throws -> Bool {
-        func getNewToken() async throws -> Bool {
-//            shouldRefreshToken = true
-//            let newSession = try await AtProtoLexicons().refresh(attempts: attempts + 1)
-//            APEnvironment.current.accessToken = newSession.accessJwt
-//            APEnvironment.current.refreshToken = newSession.refreshJwt
-//            await delegate?.sessionUpdated(newSession)
-//
-//            return true
-            false
+        func refreshViaOAuth() async throws -> Bool {
+            guard let handler = APEnvironment.current.tokenRefreshHandler else {
+                return false
+            }
+
+            if let refreshTask {
+                return try await refreshTask.value
+            }
+
+            let task = Task { try await handler() }
+            refreshTask = task
+
+            defer { refreshTask = nil }
+
+            return try await task.value
         }
-        
-        // TODO: verify this works!
-        if case .network(let networkError) = error as? AtError,
+
+        if attempts == 1,
+           case .network(let networkError) = error as? AtError,
            case .statusCode(let statusCode, _) = networkError,
-           let statusCode = statusCode?.rawValue, (400..<500).contains(statusCode),
+           let statusCode = statusCode?.rawValue,
+           statusCode == 401 || statusCode == 403 {
+            return try await refreshViaOAuth()
+        }
+
+        if case .message(let message) = error as? AtError,
+           message.error == AtErrorType.expiredToken.rawValue,
            attempts == 1 {
-            return try await getNewToken()
-        } else if case .message(let message) = error as? AtError,
-                  message.error == AtErrorType.expiredToken.rawValue {
-            return try await getNewToken()
+            return try await refreshViaOAuth()
         }
-        
+
         return false
     }
 }

Sources/CoreATProtocol/OAuth/ATProtoOAuth.swift

···
     private let config: ATProtoOAuthConfig
     private let storage: ATProtoAuthStorage
     private let identityResolver: IdentityResolver
+    private let dpopRequestActor = DPoPRequestActor()
+    private var hasPersistedKey: Bool
 
     // JWT signing keys (pattern from AtProtocol)
     private var keys: JWTKeyCollection
···
         self.identityResolver = IdentityResolver()
 
         // Initialize JWT keys (from AtProto.swift lines 19-23)
-        self.privateKey = ES256PrivateKey()
+        if let storedKeyData = try? await storage.retrievePrivateKey(),
+           let pem = String(data: storedKeyData, encoding: .utf8),
+           let restoredKey = try? ES256PrivateKey(pem: pem) {
+            self.privateKey = restoredKey
+            self.hasPersistedKey = true
+        } else {
+            self.privateKey = ES256PrivateKey()
+            self.hasPersistedKey = false
+        }
         self.keys = JWTKeyCollection()
         await self.keys.add(ecdsa: privateKey)
     }
···
 
         // Restore existing key
         self.privateKey = try ES256PrivateKey(pem: privateKeyPEM)
+        self.hasPersistedKey = true
         self.keys = JWTKeyCollection()
         await self.keys.add(ecdsa: privateKey)
     }
···
         }
 
         // Step 2: Store private key for future sessions
-        let keyPEM = privateKey.pemRepresentation
-        guard let keyData = keyPEM.data(using: .utf8) else {
-            throw ATProtoOAuthError.privateKeyExportFailed
-        }
-        try await storage.storePrivateKey(keyData)
+        try await persistPrivateKey()
 
         // Step 3: Load client metadata
         let provider = URLSession.defaultProvider
···
         do {
             login = try await authenticator.authenticate()
         } catch {
-            throw ATProtoOAuthError.authenticationFailed("OAuth flow failed: \(error.localizedDescription)")
+            if shouldRecoverFromRefreshFailure(error) {
+                try await storage.storeLogin(Login(token: "invalid", validUntilDate: .distantPast))
+                try await resetDPoPKey()
+                do {
+                    login = try await authenticator.authenticate()
+                } catch {
+                    throw ATProtoOAuthError.authenticationFailed("OAuth flow failed: \(error.localizedDescription)")
+                }
+            } else {
+                throw ATProtoOAuthError.authenticationFailed("OAuth flow failed: \(error.localizedDescription)")
+            }
         }
 
         // Step 9: Setup CoreATProtocol environment
···
         )
     }
 
+    /// Refresh tokens if the stored access token is expired (or if forced).
+    public func refreshLoginIfNeeded(handle: String? = nil, force: Bool = false) async throws -> Login? {
+        guard let login = try await storage.retrieveLogin() else {
+            return nil
+        }
+
+        if !force, login.accessToken.valid {
+            return nil
+        }
+
+        guard hasPersistedKey else {
+            return nil
+        }
+
+        guard login.refreshToken?.valid == true else {
+            return nil
+        }
+
+        let issuer: String
+        if let issuingServer = login.issuingServer {
+            issuer = issuingServer
+        } else if let handle {
+            let identity = try await identityResolver.resolve(handle: handle)
+            issuer = identity.authorizationServer
+        } else {
+            return nil
+        }
+
+        let provider = URLSession.defaultProvider
+        let serverHost = stripScheme(from: issuer)
+        let serverConfig = try await ServerMetadata.load(for: serverHost, provider: provider)
+        let clientConfig = try await ClientMetadata.load(for: config.clientMetadataURL, provider: provider)
+        let jwtGenerator: DPoPSigner.JWTGenerator = { [self] params in
+            try await self.generateJWT(params: params)
+        }
+        let tokenHandling = ATProto.tokenHandling(
+            account: handle,
+            server: serverConfig,
+            jwtGenerator: jwtGenerator
+        )
+
+        guard let refreshProvider = tokenHandling.refreshProvider else {
+            return nil
+        }
+
+        let responseProvider: URLResponseProvider = { request in
+            try await self.dpopRequestActor.response(
+                request: request,
+                jwtGenerator: jwtGenerator,
+                provider: provider,
+                issuingServer: issuer
+            )
+        }
+
+        let refreshedLogin = try await refreshProvider(login, clientConfig.credentials, responseProvider)
+        try await storage.storeLogin(refreshedLogin)
+        return refreshedLogin
+    }
+
     /// Export private key PEM for persistence
     public var privateKeyPEM: String {
         privateKey.pemRepresentation
···
 
         return end == -1 ? url : String(url.prefix(end))
     }
+
+    private func stripScheme(from url: String) -> String {
+        if url.hasPrefix("https://") {
+            return String(url.dropFirst(8))
+        } else if url.hasPrefix("http://") {
+            return String(url.dropFirst(7))
+        }
+        return url
+    }
+
+    private func persistPrivateKey() async throws {
+        let keyPEM = privateKey.pemRepresentation
+        guard let keyData = keyPEM.data(using: .utf8) else {
+            throw ATProtoOAuthError.privateKeyExportFailed
+        }
+        try await storage.storePrivateKey(keyData)
+        hasPersistedKey = true
+    }
+
+    private func resetDPoPKey() async throws {
+        privateKey = ES256PrivateKey()
+        keys = JWTKeyCollection()
+        await keys.add(ecdsa: privateKey)
+        hasPersistedKey = false
+        try await persistPrivateKey()
+    }
+
+    private func shouldRecoverFromRefreshFailure(_ error: Error) -> Bool {
+        guard let authError = error as? AuthenticatorError else {
+            return false
+        }
+
+        switch authError {
+        case .refreshNotPossible, .unauthorizedRefreshFailed, .dpopTokenExpected, .httpResponseExpected:
+            return true
+        default:
+            return false
+        }
+    }
 }
 
 // MARK: - DPoP Payload (from AtProto.swift lines 88-98)
···
 
     func verify(using key: some JWTAlgorithm) throws {
         // No additional verification needed for DPoP
+    }
+}
+
+private actor DPoPRequestActor {
+    private let signer = DPoPSigner()
+
+    func response(
+        request: URLRequest,
+        jwtGenerator: DPoPSigner.JWTGenerator,
+        provider: URLResponseProvider,
+        issuingServer: String?
+    ) async throws -> (Data, URLResponse) {
+        try await signer.response(
+            isolation: self,
+            for: request,
+            using: jwtGenerator,
+            token: nil,
+            tokenHash: nil,
+            issuingServer: issuingServer,
+            provider: provider
+        )
     }
 }