Added Authorization header caching

Author: dimitribouniol

Sources/WebPush/Helpers/URL+Origin.swift

@@ -0,0 +1,36 @@
+//
+//  URL+Origin.swift
+//  swift-webpush
+//
+//  Created by Dimitri Bouniol on 2024-12-09.
+//  Copyright © 2024 Mochi Development, Inc. All rights reserved.
+//
+
+import Foundation
+
+extension URL {
+    /// Returns the origin for the receiving URL, as defined for use in signing headers for VAPID.
+    ///
+    /// This implementation is similar to the [WHATWG Standard](https://url.spec.whatwg.org/#concept-url-origin), except that it uses the unicode form of the host, and is limited to HTTP and HTTPS schemas.
+    ///
+    /// - SeeAlso: [RFC8292 Voluntary Application Server Identification (VAPID) for Web Push §2. Application Server Self-Identification](https://datatracker.ietf.org/doc/html/rfc8292#section-2)
+    /// - SeeAlso: [RFC6454 The Web Origin Concept §6.1. Unicode Serialization of an Origin](https://datatracker.ietf.org/doc/html/rfc6454#section-6.1)
+    var origin: String {
+        /// Note that we need the unicode variant, which only URLComponents provides.
+        let components = URLComponents(url: self, resolvingAgainstBaseURL: true)
+        guard
+            let scheme = components?.scheme?.lowercased(),
+            let host = components?.host
+        else { return "null" }
+        
+        switch scheme {
+        case "http":
+            let port = components?.port ?? 80
+            return "http://" + host + (port != 80 ? ":\(port)" : "")
+        case "https":
+            let port = components?.port ?? 443
+            return "https://" + host + (port != 443 ? ":\(port)" : "")
+        default: return "null"
+        }
+    }
+}

Sources/WebPush/VAPID/VAPIDConfiguration.swift

@@ -15,9 +15,9 @@ extension VoluntaryApplicationServerIdentification {
         /// This key should be shared by all instances of your push service, and should be kept secure. Rotating this key is not recommended as you'll lose access to subscribers that registered against it.
         ///
         /// Some implementations will choose to use different keys per subscriber. In that case, choose to provide a set of keys instead.
-        public var primaryKey: Key?
-        public var keys: Set<Key>
-        public var deprecatedKeys: Set<Key>?
+        public private(set) var primaryKey: Key?
+        public private(set) var keys: Set<Key>
+        public private(set) var deprecatedKeys: Set<Key>?
         public var contactInformation: ContactInformation
         public var expirationDuration: Duration
         public var validityDuration: Duration
@@ -83,6 +83,25 @@ extension VoluntaryApplicationServerIdentification {
                 validityDuration: validityDuration
             )
         }
+        
+        mutating func updateKeys(
+            primaryKey: Key?,
+            keys: Set<Key>,
+            deprecatedKeys: Set<Key>? = nil
+        ) throws {
+            self.primaryKey = primaryKey
+            var keys = keys
+            if let primaryKey {
+                keys.insert(primaryKey)
+            }
+            guard !keys.isEmpty
+            else { throw CancellationError() } // TODO: No keys error
+            
+            self.keys = keys
+            var deprecatedKeys = deprecatedKeys ?? []
+            deprecatedKeys.subtract(keys)
+            self.deprecatedKeys = deprecatedKeys.isEmpty ? nil : deprecatedKeys
+        }
     }
 }
 
@@ -177,3 +196,9 @@ extension VAPID.Configuration {
         }
     }
 }
+
+extension Date {
+    func adding(_ duration: VAPID.Configuration.Duration) -> Self {
+        addingTimeInterval(TimeInterval(duration.seconds))
+    }
+}

Sources/WebPush/VAPID/VAPIDToken.swift

@@ -26,6 +26,16 @@ extension VAPID {
 
         static let jwtHeader = Array(#"{"typ":"JWT","alg":"ES256"}"#.utf8).base64URLEncodedString()
 
+        init(
+            origin: String,
+            contactInformation: VAPID.Configuration.ContactInformation,
+            expiration: Date
+        ) {
+            self.audience = origin
+            self.subject = contactInformation
+            self.expiration = Int(expiration.timeIntervalSince1970)
+        }
+        
         init(
             origin: String,
             contactInformation: VAPID.Configuration.ContactInformation,

Sources/WebPush/WebPushManager.swift

@@ -12,21 +12,23 @@ import Logging
 import NIOCore
 import ServiceLifecycle
 
-actor WebPushManager: Service, Sendable {
+actor WebPushManager: Sendable {
     public let vapidConfiguration: VAPID.Configuration
 
     nonisolated let logger: Logger
     let httpClient: HTTPClient
 
     let vapidKeyLookup: [VAPID.Key.ID : VAPID.Key]
-    var vapidAuthorizationCache: [String : (authorization: String, expiration: Date)] = [:]
+    var vapidAuthorizationCache: [String : (authorization: String, validUntil: Date)] = [:]
 
     public init(
         vapidConfiguration: VAPID.Configuration,
         // TODO: Add networkConfiguration for proxy, number of simultaneous pushes, etc…
         logger: Logger? = nil,
         eventLoopGroupProvider: NIOEventLoopGroupProvider = .shared(.singletonMultiThreadedEventLoopGroup)
     ) {
+        assert(vapidConfiguration.validityDuration <= vapidConfiguration.expirationDuration, "The validity duration must be earlier than the expiration duration since it represents when the VAPID Authorization token will be refreshed ahead of it expiring.");
+        assert(vapidConfiguration.expirationDuration <= .hours(24), "The expiration duration must be less than 24 hours or else push endpoints will reject messages sent to them.");
         self.vapidConfiguration = vapidConfiguration
         let allKeys = vapidConfiguration.keys + Array(vapidConfiguration.deprecatedKeys ?? [])
         self.vapidKeyLookup = Dictionary(
@@ -54,6 +56,93 @@ actor WebPushManager: Service, Sendable {
         }
     }
 
+    func loadCurrentVAPIDAuthorizationHeader(
+        endpoint: URL,
+        signingKey: VAPID.Key
+    ) throws -> String {
+        let origin = endpoint.origin
+        let cacheKey = "\(signingKey.id)|\(origin)"
+        
+        let now = Date()
+        let expirationDate = min(now.adding(vapidConfiguration.expirationDuration), now.adding(.hours(24)))
+        let renewalDate = min(now.adding(vapidConfiguration.validityDuration), expirationDate)
+        
+        if let cachedHeader = vapidAuthorizationCache[cacheKey],
+           now < cachedHeader.validUntil
+        { return cachedHeader.authorization }
+        
+        let token = VAPID.Token(
+            origin: origin,
+            contactInformation: vapidConfiguration.contactInformation,
+            expiration: expirationDate
+        )
+        
+        let authorization = try token.generateAuthorization(signedBy: signingKey)
+        vapidAuthorizationCache[cacheKey] = (authorization, validUntil: renewalDate)
+        
+        return authorization
+    }
+    
+    /// Request a VAPID key to supply to the client when requesting a new subscription.
+    ///
+    /// The ID returned is already in a format that browsers expect `applicationServerKey` to be:
+    /// ```js
+    /// const serviceRegistration = await navigator.serviceWorker?.register("/serviceWorker.mjs", { type: "module" });
+    /// const applicationServerKey = await loadVAPIDKey();
+    /// const subscription = await serviceRegistration.pushManager.subscribe({
+    ///     userVisibleOnly: true,
+    ///     applicationServerKey,
+    /// });
+    ///
+    /// ...
+    ///
+    /// async function loadVAPIDKey() {
+    ///     const httpResponse = await fetch(`/vapidKey`);
+    ///
+    ///     const webPushOptions = await httpResponse.json();
+    ///     if (httpResponse.status != 200) throw new Error(webPushOptions.reason);
+    ///
+    ///     return webPushOptions.vapid;
+    /// }
+    /// ```
+    ///
+    /// Simply provide a route to supply the key, as shown for Vapor below:
+    /// ```swift
+    /// app.get("vapidKey", use: loadVapidKey)
+    ///
+    /// ...
+    ///
+    /// struct WebPushOptions: Codable, Content, Hashable, Sendable {
+    ///     static let defaultContentType = HTTPMediaType(type: "application", subType: "webpush-options+json")
+    ///
+    ///     var vapid: VAPID.Key.ID
+    /// }
+    ///
+    /// @Sendable func loadVapidKey(request: Request) async throws -> WebPushOptions {
+    ///     WebPushOptions(vapid: manager.nextVAPIDKeyID)
+    /// }
+    /// ```
+    ///
+    /// - Note: If you supplied multiple keys in your VAPID configuration, you must specify the key ID along with the subscription you received from the browser. This can be easily done client side:
+    /// ```js
+    /// export async function registerSubscription(subscription, applicationServerKey) {
+    ///     const subscriptionStatusResponse = await fetch(`/registerSubscription`, {
+    ///         method: "POST",
+    ///         body: {
+    ///             ...subscription.toJSON(),
+    ///             applicationServerKey
+    ///         }
+    ///     });
+    ///
+    ///     ...
+    /// }
+    /// ```
+    public nonisolated var nextVAPIDKeyID: VAPID.Key.ID {
+        vapidConfiguration.primaryKey?.id ?? vapidConfiguration.keys.randomElement()!.id
+    }
+}
+
+extension WebPushManager: Service {
     public func run() async throws {
         logger.info("Starting up WebPushManager")
         try await withTaskCancellationOrGracefulShutdownHandler {