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
This commit is contained in:
yehorsobko 2026-04-17 16:34:59 +02:00 • committed by GitHub
commit cb6a62f7c2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 144 additions and 31 deletions

View file

@ -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()
))
}

View file

@ -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
}
}

View file

@ -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,6 +135,14 @@ 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,
@ -142,7 +150,8 @@ public struct OAuthAccessToken: Sendable {
expiresAt: Date?,
scopes: Set<String>,
authorizationServer: URL?,
refreshToken: String?
refreshToken: String?,
clientID: String? = nil
) {
self.value = value
self.tokenType = tokenType
@ -150,6 +159,7 @@ public struct OAuthAccessToken: Sendable {
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.

View file

@ -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)
}
}

View file

@ -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")