APNs Push Notifications

// Package.swift (abridged)
dependencies: [
    .package(url: "https://github.com/vapor/vapor.git", from: "4.122.2"),
    .package(url: "https://github.com/vapor/apns.git", from: "5.1.0"),
    // Only needed because the code below names APNSClientConfiguration and APNSAlertNotification directly:
    .package(url: "https://github.com/kylebrowning/APNSwift.git", from: "7.1.0"),
],
targets: [
    .executableTarget(name: "App", dependencies: [
        .product(name: "Vapor", package: "vapor"),
        .product(name: "VaporAPNS", package: "apns"),
        .product(name: "APNS", package: "apnswift"),
        .product(name: "APNSCore", package: "apnswift"),
    ]),
    .testTarget(name: "AppTests", dependencies: [
        .target(name: "App"),
        .product(name: "VaporTesting", package: "vapor"),
        .product(name: "APNSTestServer", package: "apnswift"), // the mock APNs server used below
        .product(name: "Crypto", package: "swift-crypto"),     // to generate a throwaway key in tests
    ]),
]

import APNS
import APNSCore
import Foundation
import Vapor
import VaporAPNS

/// Non-secret push settings, kept in `app.storage` next to the APNs containers.
struct PushSettings: Sendable {
    /// The app's bundle identifier: the `apns-topic` of every notification.
    let topic: String
}

extension Application {
    private struct PushSettingsKey: StorageKey {
        typealias Value = PushSettings
    }

    var pushSettings: PushSettings? {
        get { storage[PushSettingsKey.self] }
        set { storage[PushSettingsKey.self] = newValue }
    }
}

enum APNSSetupError: Error, CustomStringConvertible {
    case missing(String)
    case unreadableKey(String)

    var description: String {
        switch self {
        case .missing(let name): "\(name) is not set"
        case .unreadableKey(let why): "APNS signing key unreadable: \(why)"
        }
    }
}

/// Reads the PEM text of the `.p8` key from the environment, never from the repository.
///
/// - `APNS_KEY_PATH`: a file (a Docker/Kubernetes secret mounted read-only, or a Fly.io secret file)
/// - `APNS_KEY_P8`: the PEM text itself; literal `\n` sequences are turned into newlines
/// - `APNS_KEY_P8_BASE64`: the whole PEM base64-encoded, for platforms that dislike newlines
func loadAPNSKeyPEM() throws -> String {
    if let path = Environment.get("APNS_KEY_PATH") {
        guard let pem = try? String(contentsOfFile: path, encoding: .utf8) else {
            throw APNSSetupError.unreadableKey("cannot read \(path)")
        }
        return pem
    }
    if let inline = Environment.get("APNS_KEY_P8") {
        return inline.replacingOccurrences(of: "\\n", with: "\n")
    }
    if let encoded = Environment.get("APNS_KEY_P8_BASE64") {
        guard let data = Data(base64Encoded: encoded) else {
            throw APNSSetupError.unreadableKey("APNS_KEY_P8_BASE64 is not valid base64")
        }
        return String(decoding: data, as: UTF8.self)
    }
    throw APNSSetupError.missing("APNS_KEY_PATH, APNS_KEY_P8 or APNS_KEY_P8_BASE64")
}

/// Registers a development and a production APNs client that share one signing key.
func configureAPNS(_ app: Application) async throws {
    guard let keyID = Environment.get("APNS_KEY_ID") else { throw APNSSetupError.missing("APNS_KEY_ID") }
    guard let teamID = Environment.get("APNS_TEAM_ID") else { throw APNSSetupError.missing("APNS_TEAM_ID") }
    guard let topic = Environment.get("APNS_TOPIC") else { throw APNSSetupError.missing("APNS_TOPIC") }

    // `.init(pemRepresentation:)` resolves to swift-crypto's `P256.Signing.PrivateKey`; a malformed key throws here, at boot.
    await app.apns.configure(.jwt(
        privateKey: try .init(pemRepresentation: try loadAPNSKeyPEM()),
        keyIdentifier: keyID,
        teamIdentifier: teamID
    ))
    app.pushSettings = PushSettings(topic: topic)
}

import Fluent
import Foundation
import Vapor

/// Which APNs endpoint a token was issued for. An Xcode build registers a development token,
/// TestFlight and App Store builds register a production token.
enum PushEnvironment: String, Codable, Sendable {
    case development, production
}

/// What the token addresses: the app on a device, one running Live Activity, or "start a Live Activity".
enum PushTokenKind: String, Codable, Sendable {
    case device, liveActivity, pushToStart
}

final class PushToken: Model, @unchecked Sendable {
    static let schema = "push_tokens"

    @ID(key: .id) var id: UUID?
    @Parent(key: "user_id") var user: User
    @Field(key: "token") var token: String
    @Enum(key: "environment") var environment: PushEnvironment
    @Enum(key: "kind") var kind: PushTokenKind
    /// Set for `.liveActivity` tokens: the loan the activity shows.
    @OptionalField(key: "loan_id") var loanID: UUID?
    /// When the app last told us about this token. A 410 response older than this is stale.
    @Field(key: "registered_at") var registeredAt: Date

    init() {}

    init(userID: User.IDValue, token: String, environment: PushEnvironment,
         kind: PushTokenKind = .device, loanID: UUID? = nil, registeredAt: Date = .now) {
        self.$user.id = userID
        self.token = token
        self.environment = environment
        self.kind = kind
        self.loanID = loanID
        self.registeredAt = registeredAt
    }
}

struct CreatePushToken: AsyncMigration {
    func prepare(on db: any Database) async throws {
        let environment = try await db.enum("push_environment")
            .case("development").case("production").create()
        let kind = try await db.enum("push_token_kind")
            .case("device").case("liveActivity").case("pushToStart").create()
        try await db.schema(PushToken.schema)
            .id()
            .field("user_id", .uuid, .required, .references("users", "id", onDelete: .cascade))
            .field("token", .string, .required)
            .field("environment", environment, .required)
            .field("kind", kind, .required)
            .field("loan_id", .uuid)
            .field("registered_at", .datetime, .required)
            .unique(on: "token")
            .create()
    }

    func revert(on db: any Database) async throws {
        try await db.schema(PushToken.schema).delete()
        try await db.enum("push_token_kind").delete()
        try await db.enum("push_environment").delete()
    }
}

import Fluent
import Foundation
import Vapor

struct RegisterPushTokenRequest: Content {
    var token: String
    var environment: PushEnvironment
    var kind: PushTokenKind?
    var loanID: UUID?
}

/// APNs device tokens are hex strings. Apple does not promise a length, so check the alphabet and a sane range.
func normalizedPushToken(_ raw: String) throws -> String {
    let token = raw.lowercased()
    guard (32...200).contains(token.count), token.count.isMultiple(of: 2), token.allSatisfy(\.isHexDigit) else {
        throw Abort(.unprocessableEntity, reason: "token must be a hexadecimal APNs token")
    }
    return token
}

struct PushController: RouteCollection {
    func boot(routes: any RoutesBuilder) throws {
        let devices = routes.grouped("v1", "devices")
        devices.put(use: register)
        devices.delete(":token", use: unregister)
        routes.post("v1", "loans", ":loanID", "remind", use: remind)
    }

    /// PUT /v1/devices: idempotent, called every time the app receives a token.
    @Sendable
    func register(req: Request) async throws -> HTTPStatus {
        let user = try req.auth.require(User.self)
        let body = try req.content.decode(RegisterPushTokenRequest.self)
        let token = try normalizedPushToken(body.token)
        let kind = body.kind ?? .device
        if kind == .liveActivity {
            guard let loanID = body.loanID,
                  let loan = try await Loan.find(loanID, on: req.db),
                  loan.$user.id == user.id
            else { throw Abort(.unprocessableEntity, reason: "a Live Activity token needs one of your loans") }
            _ = loan
        }

        if let existing = try await PushToken.query(on: req.db).filter(\.$token == token).first() {
            // The same device after another sign-in: the token moves to the new user.
            existing.$user.id = try user.requireID()
            existing.environment = body.environment
            existing.kind = kind
            existing.loanID = body.loanID
            existing.registeredAt = .now
            try await existing.save(on: req.db)
            return .ok
        }
        let created = PushToken(userID: try user.requireID(), token: token, environment: body.environment,
                                kind: kind, loanID: body.loanID)
        try await created.save(on: req.db)
        return .created
    }

    /// DELETE /v1/devices/:token: call it from the sign-out flow.
    @Sendable
    func unregister(req: Request) async throws -> HTTPStatus {
        let user = try req.auth.require(User.self)
        let token = try normalizedPushToken(req.parameters.require("token"))
        try await PushToken.query(on: req.db)
            .filter(\.$token == token)
            .filter(\.$user.$id == user.requireID())
            .delete()
        return .noContent
    }

    /// POST /v1/loans/:loanID/remind: what a due-date job calls for each loan; here it is a librarian action.
    @Sendable
    func remind(req: Request) async throws -> RemindResponse {
        let user = try req.auth.require(User.self)
        guard user.role != .reader else { throw Abort(.forbidden) }
        guard let loan = try await Loan.query(on: req.db)
            .filter(\.$id == req.parameters.require("loanID", as: UUID.self))
            .with(\.$book)
            .first()
        else { throw Abort(.notFound) }
        let push = try PushService(app: req.application, logger: req.logger)
        let outcome = try await push.sendLoanDue(loan, bookTitle: loan.book.title, on: req.db)
        return RemindResponse(delivered: outcome.delivered, pruned: outcome.pruned, retryable: outcome.retryable, dropped: outcome.dropped)
    }
}

struct RemindResponse: Content {
    let delivered: Int
    let pruned: Int
    let retryable: Int
    let dropped: Int
}

curl -X PUT http://localhost:8080/v1/devices \
  -H 'Authorization: Bearer <BOOKSHELF_ACCESS_TOKEN>' -H 'Content-Type: application/json' \
  -d '{"token":"abababababababababababababababababababababababababababababababab","environment":"development"}'

import APNS
import APNSCore
import Fluent
import Foundation
import Logging
import Vapor
import VaporAPNS

/// Custom data that rides next to the `aps` dictionary. Never use the key `aps` here.
struct LoanDuePayload: Encodable, Sendable {
    let kind = "loan-due"
    let loanID: UUID
    let bookID: UUID
}

/// A hint for a silent push: "something changed, go and sync".
struct SyncHint: Encodable, Sendable {
    let resource: String
}

/// The Live Activity `ContentState`; it must match the app's `ActivityAttributes.ContentState` key for key.
struct LoanActivityState: Encodable, Sendable {
    let title: String
    let daysLeft: Int
}

/// The static half of the Live Activity (`ActivityAttributes`); its type name is sent as `attributes-type`.
struct LoanAttributes: Encodable, Sendable {
    let loanID: UUID
}

struct PushOutcome: Sendable, Equatable {
    var delivered = 0
    var pruned = 0
    var retryable = 0
    var dropped = 0
}

struct PushService: Sendable {
    let app: Application
    let topic: String
    let logger: Logger

    init(app: Application, logger: Logger) throws {
        guard let settings = app.pushSettings else { throw Abort(.serviceUnavailable, reason: "push is not configured") }
        self.app = app
        self.topic = settings.topic
        self.logger = logger
    }

    private func client(for environment: PushEnvironment) async -> APNSGenericClient {
        await app.apns.client(environment == .production ? .production : .development)
    }

    // MARK: - Alert

    func sendLoanDue(_ loan: Loan, bookTitle: String, on db: any Database) async throws -> PushOutcome {
        let tokens = try await PushToken.query(on: db)
            .filter(\.$user.$id == loan.$user.id)
            .filter(\.$kind == .device)
            .all()
        let loanID = try loan.requireID()
        var notification = APNSAlertNotification(
            alert: .init(
                title: .raw("Loan due soon"),
                body: .raw("\"\(bookTitle)\" is due back on \(loan.dueAt.formatted(date: .abbreviated, time: .omitted)).")
            ),
            expiration: .timeIntervalSince1970InSeconds(Int(loan.dueAt.timeIntervalSince1970)),
            priority: .immediately,
            topic: topic,
            payload: LoanDuePayload(loanID: loanID, bookID: loan.$book.id),
            sound: .default,
            threadID: "loan-\(loanID)",
            category: "LOAN_DUE"
        )
        notification.collapseID = "loan-\(loanID)" // a newer reminder replaces the older one on the lock screen
        var outcome = PushOutcome()
        for token in tokens {
            let apns = await client(for: token.environment)
            await record(&outcome, for: token, on: db) {
                try await apns.sendAlertNotification(notification, deviceToken: token.token)
            }
        }
        return outcome
    }

    // MARK: - Background

    func sendSyncHint(to userID: User.IDValue, resource: String, on db: any Database) async throws -> PushOutcome {
        let tokens = try await PushToken.query(on: db)
            .filter(\.$user.$id == userID)
            .filter(\.$kind == .device)
            .all()
        let notification = APNSBackgroundNotification(
            expiration: .immediately,
            topic: topic,
            payload: SyncHint(resource: resource)
        )
        var outcome = PushOutcome()
        for token in tokens {
            let apns = await client(for: token.environment)
            await record(&outcome, for: token, on: db) {
                try await apns.sendBackgroundNotification(notification, deviceToken: token.token)
            }
        }
        return outcome
    }

    // MARK: - Live Activity

    func updateLoanActivity(_ loan: Loan, bookTitle: String, ending: Bool, on db: any Database) async throws -> PushOutcome {
        let loanID = try loan.requireID()
        let tokens = try await PushToken.query(on: db)
            .filter(\.$kind == .liveActivity)
            .filter(\.$loanID == loanID)
            .all()
        let now = Int(Date.now.timeIntervalSince1970)
        let daysLeft = max(0, Calendar.current.dateComponents([.day], from: .now, to: loan.dueAt).day ?? 0)
        let notification = APNSLiveActivityNotification(
            expiration: .none,
            priority: ending ? .immediately : .consideringDevicePower,
            appID: topic,
            contentState: LoanActivityState(title: bookTitle, daysLeft: daysLeft),
            event: ending ? .end : .update,
            timestamp: now,
            dismissalDate: ending ? .immediately : .none,
            staleDate: now + 6 * 3600
        )
        var outcome = PushOutcome()
        for token in tokens {
            let apns = await client(for: token.environment)
            await record(&outcome, for: token, on: db) {
                try await apns.sendLiveActivityNotification(notification, deviceToken: token.token)
            }
        }
        return outcome
    }

    /// Starts a Live Activity remotely with a push-to-start token the app registered earlier.
    func startLoanActivity(_ loan: Loan, bookTitle: String, on db: any Database) async throws -> PushOutcome {
        let loanID = try loan.requireID()
        let tokens = try await PushToken.query(on: db)
            .filter(\.$user.$id == loan.$user.id)
            .filter(\.$kind == .pushToStart)
            .all()
        let notification = APNSStartLiveActivityNotification(
            expiration: .none,
            priority: .immediately,
            appID: topic,
            contentState: LoanActivityState(title: bookTitle, daysLeft: 2),
            timestamp: Int(Date.now.timeIntervalSince1970),
            attributes: LoanAttributes(loanID: loanID),
            attributesType: "LoanAttributes",
            alert: .init(title: .raw("Loan due soon"), body: .raw(bookTitle))
        )
        var outcome = PushOutcome()
        for token in tokens {
            let apns = await client(for: token.environment)
            await record(&outcome, for: token, on: db) {
                try await apns.sendStartLiveActivityNotification(notification, pushToStartToken: token.token)
            }
        }
        return outcome
    }

    // MARK: - Outcome handling

    /// Runs one send and turns every result into a counter and, for dead tokens, a delete.
    private func record(
        _ outcome: inout PushOutcome,
        for token: PushToken,
        on db: any Database,
        _ send: () async throws -> APNSResponse
    ) async {
        do {
            let response = try await send()
            outcome.delivered += 1
            logger.debug("accepted by APNs", metadata: ["apns-id": "\(response.apnsID?.uuidString ?? "-")"])
        } catch let error as APNSError {
            switch error.responseStatus {
            case 410:
                // Unregistered or ExpiredToken. Keep the row if the app re-registered after APNs last saw it fail.
                if let failedAt = error.timestamp, token.registeredAt > failedAt {
                    outcome.dropped += 1
                } else {
                    try? await token.delete(on: db)
                    outcome.pruned += 1
                }
            case 429, 500, 503:
                outcome.retryable += 1
                logger.warning("APNs asked us to retry", metadata: ["status": "\(error.responseStatus)", "reason": "\(error.reason?.reason ?? "-")"])
            case 403:
                // A wrong key, team or topic is a deployment bug, not a bad device.
                outcome.dropped += 1
                logger.error("APNs rejected our credentials", metadata: ["reason": "\(error.reason?.reason ?? "-")"])
            default:
                outcome.dropped += 1
                logger.warning("APNs rejected a notification", metadata: ["status": "\(error.responseStatus)", "reason": "\(error.reason?.reason ?? "-")"])
            }
        } catch {
            // Connection reset, TLS or DNS trouble: the notification may not have reached APNs at all.
            outcome.retryable += 1
            logger.warning("could not reach APNs", metadata: ["error": "\(error)"])
        }
    }
}

outcome: PushOutcome(delivered: 1, pruned: 1, retryable: 0, dropped: 0)
outcome (re-registered): PushOutcome(delivered: 1, pruned: 0, retryable: 0, dropped: 1)

@testable import App
import APNS
import APNSCore
import APNSTestServer
import Crypto
import Fluent
import FluentSQLiteDriver
import Foundation
import Testing
import Vapor
import VaporAPNS
import VaporTesting

/// Signs every request in as the given user, standing in for the app's real authenticator.
struct SignedInAs: AsyncMiddleware {
    let user: User
    func respond(to request: Request, chainingTo next: any AsyncResponder) async throws -> Response {
        request.auth.login(user)
        return try await next.respond(to: request)
    }
}

@Suite("APNs push", .serialized)
struct PushTests {
    /// An app wired to a local mock APNs server instead of api.sandbox.push.apple.com.
    private func withApp(_ body: (Application, APNSTestServer, User, Loan) async throws -> Void) async throws {
        let server = APNSTestServer()
        try await server.start()
        let app = try await Application.make(.testing)
        do {
            app.databases.use(.sqlite(.memory), as: .sqlite)
            app.migrations.add(CreateBookshelf())
            app.migrations.add(CreatePushToken())
            try await app.autoMigrate()

            // A throwaway key generated here; real keys come from the environment.
            let config = APNSClientConfiguration(
                authenticationMethod: .jwt(
                    privateKey: P256.Signing.PrivateKey(),
                    keyIdentifier: "ABC123DEFG",
                    teamIdentifier: "DEF123GHIJ"
                ),
                environment: .custom(url: "http://127.0.0.1", port: server.port)
            )
            await app.apns.containers.use(
                config,
                eventLoopGroupProvider: .shared(app.eventLoopGroup),
                responseDecoder: JSONDecoder(),
                requestEncoder: JSONEncoder(),
                as: .development
            )
            app.pushSettings = PushSettings(topic: "com.example.bookshelf")

            let author = Author(name: "Ursula K. Le Guin"); try await author.save(on: app.db)
            let book = Book(title: "The Dispossessed", isbn: "978-0-06-051275-4", publishedYear: 1974,
                            priceCents: 1299, status: .published, authorID: try author.requireID())
            try await book.save(on: app.db)
            let librarian = User(email: "lib@example.com", passwordHash: "x", role: .librarian)
            try await librarian.save(on: app.db)
            let loan = Loan(bookID: try book.requireID(), userID: try librarian.requireID(),
                            loanedAt: .now, dueAt: .now.addingTimeInterval(2 * 86_400))
            try await loan.save(on: app.db)

            try app.grouped(SignedInAs(user: librarian)).register(collection: PushController())
            try await body(app, server, librarian, loan)
        } catch {
            try? await app.asyncShutdown()
            try? await server.shutdown()
            throw error
        }
        try await app.asyncShutdown()
        try await server.shutdown()
    }

    private let goodToken = String(repeating: "ab", count: 32)

    @Test("register is idempotent and rejects non-hex tokens")
    func registration() async throws {
        try await withApp { app, _, _, _ in
            let body = RegisterPushTokenRequest(token: goodToken.uppercased(), environment: .development)
            try await app.test(.PUT, "v1/devices", beforeRequest: { try $0.content.encode(body) }) { res async in
                #expect(res.status == .created)
            }
            try await app.test(.PUT, "v1/devices", beforeRequest: { try $0.content.encode(body) }) { res async in
                #expect(res.status == .ok)
            }
            #expect(try await PushToken.query(on: app.db).count() == 1)
            #expect(try await PushToken.query(on: app.db).first()?.token == goodToken)

            let bad = RegisterPushTokenRequest(token: "not-a-token", environment: .development)
            try await app.test(.PUT, "v1/devices", beforeRequest: { try $0.content.encode(bad) }) { res async in
                #expect(res.status == .unprocessableEntity)
            }
            try await app.test(.DELETE, "v1/devices/\(goodToken)") { res async in
                #expect(res.status == .noContent)
            }
            #expect(try await PushToken.query(on: app.db).count() == 0)
        }
    }

    @Test("the reminder reaches APNs with the right headers and payload")
    func alert() async throws {
        try await withApp { app, server, user, loan in
            try await PushToken(userID: user.requireID(), token: goodToken, environment: .development).save(on: app.db)
            try await app.test(.POST, "v1/loans/\(loan.requireID())/remind") { res async throws in
                #expect(res.status == .ok)
                print("remind ->", res.body.string)
                #expect(try res.content.decode(RemindResponse.self).delivered == 1)
            }
            let sent = try #require(server.getSentNotifications().first)
            print("headers: type=\(sent.pushType ?? "-") topic=\(sent.topic ?? "-") priority=\(sent.priority ?? "-") expiration=\(sent.expiration ?? "-") collapse=\(sent.collapseID ?? "-")")
            print("authorization:", String((sent.authorization ?? "").prefix(10)) + "...")
            print("payload:", String(decoding: sent.payload, as: UTF8.self))
            #expect(sent.pushType == "alert")
            #expect(sent.topic == "com.example.bookshelf")
            #expect(sent.priority == "10")
            #expect(sent.deviceToken == goodToken)
        }
    }

    @Test("a 410 deletes the token unless the app registered after the failure")
    func unregistered() async throws {
        try await withApp { app, server, user, loan in
            let dead = String(repeating: "f", count: 64) // the mock server answers 410 for this token
            try await PushToken(userID: user.requireID(), token: dead, environment: .development,
                                registeredAt: Date(timeIntervalSince1970: 1_000_000_000)).save(on: app.db)
            try await PushToken(userID: user.requireID(), token: goodToken, environment: .development).save(on: app.db)
            let push = try PushService(app: app, logger: app.logger)
            let outcome = try await push.sendLoanDue(loan, bookTitle: "The Dispossessed", on: app.db)
            print("outcome:", outcome)
            #expect(outcome.delivered == 1)
            #expect(outcome.pruned == 1)
            #expect(try await PushToken.query(on: app.db).filter(\.$token == dead).count() == 0)

            // The same 410, but the app re-registered after APNs last saw the failure: keep the row.
            try await PushToken(userID: user.requireID(), token: dead, environment: .development).save(on: app.db)
            let again = try await push.sendLoanDue(loan, bookTitle: "The Dispossessed", on: app.db)
            print("outcome (re-registered):", again)
            #expect(again.pruned == 0)
            #expect(again.dropped == 1)
            _ = server
        }
    }

    // ... the 503, background/Live Activity and key-loading tests follow the same shape.
}

✔ Test "register is idempotent and rejects non-hex tokens" passed
✔ Test "the reminder reaches APNs with the right headers and payload" passed
✔ Test "a 410 deletes the token unless the app registered after the failure" passed
✔ Test "a 503 is counted as retryable and the token is kept" passed
✔ Test "background and Live Activity pushes use their own push types" passed
✔ Test "the signing key loader reads a file and rejects garbage" passed
✔ Test run with 6 tests in 1 suite passed