feat: auth improvements (#217)

* feat: improvement to oauth resources resolution

* feat: improvements to refreshToken procedure and updates to unit tests

* feat: improvements to documentation of clien_id

Author: yehorsobko-mac

Sources/MCP/Base/Authorization/OAuthAuthorizer.swift

@@ -217,7 +217,6 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable {
             return nil
         }
         if accessToken.isExpired() {
-            tokenStorage.clear()
             return nil
         }
         return "\(OAuthTokenType.bearer) \(accessToken.value)"
@@ -367,7 +366,6 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable {
     public func prepareAuthorization(for endpoint: URL, session: URLSession) async throws {
         guard configuration.proactiveRefreshWindowSeconds > 0 else { return }
         guard let token = tokenStorage.load() else { return }
-        guard !token.isExpired() else { return }
         guard token.isExpired(skewSeconds: configuration.proactiveRefreshWindowSeconds) else {
             return
         }
@@ -441,8 +439,10 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable {
             candidates.append(fallback)
         }
 
+        let fallbackIssuer = try? discoveryClient.metadataDiscovery
+            .authorizationServerFallbackIssuer(from: endpoint)
         let metadata = try await discoveryClient.fetchProtectedResourceMetadata(
-            candidates: candidates, session: session)
+            candidates: candidates, fallbackIssuer: fallbackIssuer, session: session)
         try validateProtectedResource(metadata: metadata, endpoint: endpoint)
 
         self.protectedResourceMetadata = metadata
@@ -716,10 +716,17 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable {
             expiresAt: expiresAt,
             scopes: scopeSet,
             authorizationServer: selectedAuthorizationServer,
-            refreshToken: decoded.refreshToken
+            refreshToken: decoded.refreshToken,
+            clientID: nonEmptyClientID()
         ))
     }
 
+    /// Returns the configured `client_id` or `nil` if the authorizer has not yet been assigned one.
+    private func nonEmptyClientID() -> String? {
+        let id = configuration.authentication.clientID
+        return id.isEmpty ? nil : id
+    }
+
     private func resolveTokenEndpoint(
         asMetadata: OAuthAuthorizationServerMetadata
     ) throws -> URL {
@@ -823,7 +830,8 @@ public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable {
             expiresAt: nil,
             scopes: requestedScopes ?? [],
             authorizationServer: authorizationServer,
-            refreshToken: nil
+            refreshToken: nil,
+            clientID: nonEmptyClientID()
         ))
     }
 

Sources/MCP/Base/Authorization/OAuthDiscoveryClient.swift

@@ -7,7 +7,7 @@ import Foundation
 /// Internal protocol for fetching OAuth discovery metadata.
 protocol OAuthDiscoveryFetching: Sendable {
     var metadataDiscovery: any OAuthMetadataDiscovering { get }
-    func fetchProtectedResourceMetadata(candidates: [URL], session: URLSession) async throws -> OAuthProtectedResourceMetadata
+    func fetchProtectedResourceMetadata(candidates: [URL], fallbackIssuer: URL?, session: URLSession) async throws -> OAuthProtectedResourceMetadata
     func fetchAuthorizationServerMetadata(candidates: [URL], session: URLSession) async throws -> (server: URL, metadata: OAuthAuthorizationServerMetadata)
 }
 
@@ -29,8 +29,13 @@ struct OAuthDiscoveryClient: Sendable {
     }
 
     /// Fetches Protected Resource Metadata from the first candidate that returns a valid response.
+    ///
+    /// If all candidates fail and `fallbackIssuer` is provided, returns synthetic metadata
+    /// using that issuer as the authorization server — for servers that do not expose a
+    /// PRM document at any well-known path.
     func fetchProtectedResourceMetadata(
         candidates: [URL],
+        fallbackIssuer: URL?,
         session: URLSession
     ) async throws -> OAuthProtectedResourceMetadata {
         let decoder = JSONDecoder()
@@ -56,15 +61,28 @@ struct OAuthDiscoveryClient: Sendable {
                 continue
             }
         }
+        if let fallbackIssuer {
+            return OAuthProtectedResourceMetadata(
+                resource: nil,
+                authorizationServers: [fallbackIssuer],
+                scopesSupported: nil
+            )
+        }
         throw OAuthAuthorizationError.metadataDiscoveryFailed
     }
 
-    /// Fetches Authorization Server Metadata from the first candidate that returns a valid response.
+    /// Fetches Authorization Server Metadata from candidates, preferring a response whose
+    /// `issuer` matches the candidate URL (RFC 8414 §3). If no candidate produces a matching
+    /// issuer, the first valid response is accepted and its own `issuer` is used as the
+    /// server identity — accommodating servers that serve metadata at one host but advertise
+    /// a different issuer.
     func fetchAuthorizationServerMetadata(
         candidates: [URL],
         session: URLSession
     ) async throws -> (server: URL, metadata: OAuthAuthorizationServerMetadata) {
         let decoder = JSONDecoder()
+        var firstValid: (server: URL, metadata: OAuthAuthorizationServerMetadata)?
+
         for candidateServer in candidates {
             guard (try? urlValidator.validateAuthorizationServer(
                 candidateServer, context: "Authorization server issuer")) != nil
@@ -95,22 +113,28 @@ struct OAuthDiscoveryClient: Sendable {
                     let asMetadata = try decoder.decode(
                         OAuthAuthorizationServerMetadata.self, from: data)
 
-                    // RFC 8414 §3: issuer field must match the candidate server URL.
-                    // Absent issuer is tolerated (some servers omit it).
-                    if let metadataIssuer = asMetadata.issuer {
-                        guard metadataIssuer.absoluteString.lowercased()
+                    // Prefer metadata whose issuer matches the candidate URL (RFC 8414 §3).
+                    let issuerMatches =
+                        asMetadata.issuer == nil
+                        || asMetadata.issuer?.absoluteString.lowercased()
                             == candidateServer.absoluteString.lowercased()
-                        else {
-                            continue
-                        }
+                    if issuerMatches {
+                        return (server: candidateServer, metadata: asMetadata)
+                    }
+                    // Keep as fallback in case no issuer-matching candidate is found.
+                    if firstValid == nil {
+                        let server = asMetadata.issuer ?? candidateServer
+                        firstValid = (server: server, metadata: asMetadata)
                     }
-
-                    return (server: candidateServer, metadata: asMetadata)
                 } catch {
                     continue
                 }
             }
         }
+
+        if let firstValid {
+            return firstValid
+        }
         throw OAuthAuthorizationError.authorizationServerMetadataDiscoveryFailed
     }
 }

Sources/MCP/Base/Authorization/OAuthModels.swift

@@ -113,7 +113,7 @@ struct OAuthTokenErrorResponse: Decodable {
 ///
 /// Stored by ``TokenStorage`` and produced by ``OAuthAuthorizer`` after a successful
 /// token request. Use ``isExpired(now:skewSeconds:)`` to check validity before use.
-public struct OAuthAccessToken: Sendable {
+public struct OAuthAccessToken: Sendable, Codable {
     /// The raw bearer token string for use in the `Authorization` header.
     public let value: String
 
@@ -135,21 +135,31 @@ public struct OAuthAccessToken: Sendable {
     /// The refresh token, if the authorization server issued one alongside the access token.
     public let refreshToken: String?
 
+    /// The OAuth `client_id` this token was issued to, if the authorizer had one at storage time.
+    ///
+    /// Captures the `client_id` held by the authorizer's ``OAuthConfiguration/authentication`` when
+    /// the token was saved — including identifiers assigned by Dynamic Client Registration
+    /// (RFC 7591). `nil` when no `client_id` was configured (for example, the placeholder used
+    /// before registration completes).
+    public let clientID: String?
+
     /// Creates a new access token record.
     public init(
         value: String,
         tokenType: String,
         expiresAt: Date?,
         scopes: Set<String>,
         authorizationServer: URL?,
-        refreshToken: String?
+        refreshToken: String?,
+        clientID: String? = nil
     ) {
         self.value = value
         self.tokenType = tokenType
         self.expiresAt = expiresAt
         self.scopes = scopes
         self.authorizationServer = authorizationServer
         self.refreshToken = refreshToken
+        self.clientID = clientID
     }
 
     /// Returns `true` if the token has expired or will expire within the skew window.

Tests/MCPTests/OAuthAuthorizerTests.swift

@@ -64,7 +64,7 @@ final class MockDiscoveryClient: OAuthDiscoveryFetching, @unchecked Sendable {
         )
     }
 
-    func fetchProtectedResourceMetadata(candidates: [URL], session: URLSession) async throws -> OAuthProtectedResourceMetadata {
+    func fetchProtectedResourceMetadata(candidates: [URL], fallbackIssuer: URL?, session: URLSession) async throws -> OAuthProtectedResourceMetadata {
         fetchProtectedResourceMetadataCallCount += 1
         return protectedResourceMetadataResult
     }
@@ -100,6 +100,10 @@ final class MockTokenClient: OAuthTokenRequesting, @unchecked Sendable {
 
 final class MockClientRegistrar: OAuthClientRegistering, @unchecked Sendable {
     var registerCallCount = 0
+    var registrationResult: (
+        response: OAuthClientRegistrationResponse,
+        updatedAuthentication: OAuthConfiguration.TokenEndpointAuthentication
+    )?
 
     func register(
         configuration: OAuthConfiguration,
@@ -110,7 +114,7 @@ final class MockClientRegistrar: OAuthClientRegistering, @unchecked Sendable {
         updatedAuthentication: OAuthConfiguration.TokenEndpointAuthentication
     )? {
         registerCallCount += 1
-        return nil
+        return registrationResult
     }
 }
 
@@ -381,4 +385,45 @@ struct OAuthAuthorizerTests {
 
         #expect(registrar.registerCallCount == 0)
     }
+
+    @Test("handleChallenge persists the DCR-assigned clientID on the saved token")
+    func testHandleChallengePersistsDCRClientIDOnToken() async throws {
+        let assignedClientID = "dcr-assigned-client-id"
+        let tokenStorage = InMemoryTokenStorage()
+        let registrar = MockClientRegistrar()
+        registrar.registrationResult = (
+            response: OAuthClientRegistrationResponse(
+                clientID: assignedClientID,
+                clientSecret: nil,
+                tokenEndpointAuthMethod: nil,
+                clientSecretExpiresAt: nil
+            ),
+            updatedAuthentication: .none(clientID: assignedClientID)
+        )
+
+        let config = OAuthConfiguration(
+            authentication: .none(clientID: "")
+        )
+        let authorizer = OAuthAuthorizer(
+            configuration: config,
+            tokenStorage: tokenStorage,
+            urlValidator: MockURLValidator(),
+            discoveryClient: MockDiscoveryClient(),
+            tokenEndpointClient: MockTokenClient(),
+            clientRegistrar: registrar,
+            authCodeFlow: MockAuthCodeFlow()
+        )
+
+        let handled = try await authorizer.handleChallenge(
+            statusCode: 401,
+            headers: headers401,
+            endpoint: endpoint,
+            operationKey: nil,
+            session: .shared
+        )
+
+        #expect(handled == true)
+        #expect(registrar.registerCallCount == 1)
+        #expect(tokenStorage.load()?.clientID == assignedClientID)
+    }
 }

Tests/MCPTests/OAuthDiscoveryClientTests.swift

@@ -49,6 +49,7 @@ import Testing
 
             let metadata = try await makeClient().fetchProtectedResourceMetadata(
                 candidates: [URL(string: "https://example.com/.well-known/oauth-protected-resource")!],
+                fallbackIssuer: nil,
                 session: session
             )
             let expected = OAuthProtectedResourceMetadata(
@@ -76,6 +77,7 @@ import Testing
                     URL(string: "https://example.com/.well-known/oauth-protected-resource/mcp")!,
                     URL(string: "https://example.com/.well-known/oauth-protected-resource")!,
                 ],
+                fallbackIssuer: nil,
                 session: session
             )
             let expected = OAuthProtectedResourceMetadata(
@@ -103,6 +105,7 @@ import Testing
                     URL(string: "https://example.com/.well-known/oauth-protected-resource/mcp")!,
                     URL(string: "https://example.com/.well-known/oauth-protected-resource")!,
                 ],
+                fallbackIssuer: nil,
                 session: session
             )
             let expected = OAuthProtectedResourceMetadata(
@@ -112,7 +115,7 @@ import Testing
             #expect(metadata == expected)
         }
 
-        @Test("Throws metadataDiscoveryFailed when all candidates fail")
+        @Test("Throws metadataDiscoveryFailed when all candidates fail and no fallback issuer")
         func testFetchProtectedResourceMetadataThrowsWhenAllFail() async throws {
             let (session, key) = makeIsolatedSession()
             await IsolatedMockURLProtocol.setHandler(key: key) { request in
@@ -126,11 +129,34 @@ import Testing
                     candidates: [
                         URL(string: "https://example.com/.well-known/oauth-protected-resource")!
                     ],
+                    fallbackIssuer: nil,
                     session: session
                 )
             }
         }
 
+        @Test("Returns synthetic metadata with fallback issuer when all candidates fail")
+        func testFetchProtectedResourceMetadataUsesFallbackIssuer() async throws {
+            let fallback = URL(string: "https://example.com")!
+            let (session, key) = makeIsolatedSession()
+            await IsolatedMockURLProtocol.setHandler(key: key) { request in
+                let response = HTTPURLResponse(
+                    url: request.url!, statusCode: 404, httpVersion: nil, headerFields: nil)!
+                return (response, Data())
+            }
+
+            let metadata = try await makeClient().fetchProtectedResourceMetadata(
+                candidates: [URL(string: "https://example.com/.well-known/oauth-protected-resource")!],
+                fallbackIssuer: fallback,
+                session: session
+            )
+            let expected = OAuthProtectedResourceMetadata(
+                resource: nil,
+                authorizationServers: [fallback],
+                scopesSupported: nil)
+            #expect(metadata == expected)
+        }
+
         // MARK: - fetchAuthorizationServerMetadata
 
         @Test("Returns server and metadata when issuer matches")
@@ -162,23 +188,23 @@ import Testing
             #expect(metadata == expectedMetadata)
         }
 
-        @Test("Skips candidate when issuer field does not match")
-        func testFetchAuthorizationServerMetadataSkipsIssuerMismatch() async throws {
-            let wrongIssuerBody = try makeASMetadataBody(issuer: "https://other.example.com")
+        @Test("Uses metadata issuer as server identity when it differs from candidate URL")
+        func testFetchAuthorizationServerMetadataUsesMetadataIssuer() async throws {
+            let metadataIssuer = "https://other.example.com"
+            let body = try makeASMetadataBody(issuer: metadataIssuer)
             let (session, key) = makeIsolatedSession()
             await IsolatedMockURLProtocol.setHandler(key: key) { _ in
                 let response = HTTPURLResponse(
                     url: URL(string: "https://auth.example.com")!,
                     statusCode: 200, httpVersion: nil, headerFields: nil)!
-                return (response, wrongIssuerBody)
+                return (response, body)
             }
 
-            await #expect(throws: OAuthAuthorizationError.self) {
-                try await makeClient().fetchAuthorizationServerMetadata(
-                    candidates: [URL(string: "https://auth.example.com")!],
-                    session: session
-                )
-            }
+            let (server, _) = try await makeClient().fetchAuthorizationServerMetadata(
+                candidates: [URL(string: "https://auth.example.com")!],
+                session: session
+            )
+            #expect(server == URL(string: metadataIssuer)!)
         }
 
         @Test("Skips private IP candidates without making HTTP calls")