Fix concurrent APNs JWT generation and preserve error reasons

hub/src/push-ios/apnsClient.test.ts

@@ -1,4 +1,4 @@
-import { afterAll, beforeAll, describe, expect, it } from 'bun:test'
+import { afterAll, beforeAll, describe, expect, it, spyOn } from 'bun:test'
 import http2 from 'node:http2'
 import { generateKeyPairSync } from 'node:crypto'
 import * as jose from 'jose'
@@ -50,6 +50,33 @@ describe('ApnsJwtProvider', () => {
         expect(jose.decodeJwt(second).iat).toBe(Math.floor(later / 1000))
         expect(APNS_JWT_MAX_AGE_MS).toBe(45 * 60 * 1000)
     })
+
+    it('shares one JWT across concurrent callers on startup and refresh', async () => {
+        const provider = new ApnsJwtProvider(TEST_KEY_P8, 'KEYID12345', 'TEAMID9999')
+        const t0 = 1_700_000_000_000
+        const first = await Promise.all(Array.from({ length: 8 }, () => provider.getToken(t0)))
+        expect(new Set(first).size).toBe(1)
+
+        const later = t0 + APNS_JWT_MAX_AGE_MS + 1
+        const refreshed = await Promise.all(Array.from({ length: 8 }, () => provider.getToken(later)))
+        expect(new Set(refreshed).size).toBe(1)
+        expect(refreshed[0]).not.toBe(first[0])
+        expect(await provider.getToken(later)).toBe(refreshed[0]!)
+    })
+
+    it('can retry signing after a failed token generation', async () => {
+        const provider = new ApnsJwtProvider(TEST_KEY_P8, 'KEYID12345', 'TEAMID9999')
+        const sign = spyOn(jose.SignJWT.prototype, 'sign')
+            .mockRejectedValueOnce(new Error('test signing failure'))
+        try {
+            await expect(provider.getToken()).rejects.toThrow('test signing failure')
+            const token = await provider.getToken()
+            expect(jose.decodeProtectedHeader(token).kid).toBe('KEYID12345')
+            expect(sign).toHaveBeenCalledTimes(2)
+        } finally {
+            sign.mockRestore()
+        }
+    })
 })
 
 describe('buildApnsRequestBody', () => {

hub/src/push-ios/apnsClient.ts

@@ -20,7 +20,7 @@ export const APNS_JWT_MAX_AGE_MS = 45 * 60 * 1000
  * Header: `{ alg: "ES256", kid: <keyId> }`; claims: `{ iss: <teamId>, iat }`.
  */
 export class ApnsJwtProvider {
-    private cached: { token: string; issuedAtMs: number } | null = null
+    private cached: { token: Promise<string>; issuedAtMs: number } | null = null
 
     constructor(
         /** PKCS8 PEM contents of the .p8 APNs auth key. */
@@ -33,14 +33,27 @@ export class ApnsJwtProvider {
         if (this.cached && nowMs - this.cached.issuedAtMs < APNS_JWT_MAX_AGE_MS) {
             return this.cached.token
         }
+        // Share the signing work across concurrent device fan-out, including
+        // on startup and when the previous JWT needs refreshing.
+        const token = this.signToken(nowMs)
+        this.cached = { token, issuedAtMs: nowMs }
+        try {
+            return await token
+        } catch (error) {
+            if (this.cached?.token === token) {
+                this.cached = null
+            }
+            throw error
+        }
+    }
+
+    private async signToken(nowMs: number): Promise<string> {
         const key = await jose.importPKCS8(this.keyP8, 'ES256')
-        const token = await new jose.SignJWT({})
+        return new jose.SignJWT({})
             .setProtectedHeader({ alg: 'ES256', kid: this.keyId })
             .setIssuer(this.teamId)
             .setIssuedAt(Math.floor(nowMs / 1000))
             .sign(key)
-        this.cached = { token, issuedAtMs: nowMs }
-        return token
     }
 }
 

relay/src/apns.test.ts

@@ -5,9 +5,9 @@
  * if these tests ever start failing on a Bun upgrade, swap the transport
  * behind the ApnsClient interface (see relay/src/apns.ts header comment).
  */
-import { afterEach, beforeEach, describe, expect, test } from 'bun:test'
+import { afterEach, beforeEach, describe, expect, spyOn, test } from 'bun:test'
 import { createServer } from 'node:http2'
-import { decodeProtectedHeader, jwtVerify } from 'jose'
+import { decodeProtectedHeader, jwtVerify, SignJWT } from 'jose'
 import {
     APNS_JWT_MAX_AGE_MS,
     ApnsJwtProvider,
@@ -111,6 +111,33 @@ describe('ApnsJwtProvider', () => {
         // with identical claims.
         expect(second).not.toBe(first)
     })
+
+    test('shares one JWT across concurrent callers on startup and refresh', async () => {
+        let nowMs = 1_700_000_000_000
+        const provider = makeProvider(() => nowMs)
+        const first = await Promise.all(Array.from({ length: 8 }, () => provider.getToken()))
+        expect(new Set(first).size).toBe(1)
+
+        nowMs += APNS_JWT_MAX_AGE_MS + 1
+        const refreshed = await Promise.all(Array.from({ length: 8 }, () => provider.getToken()))
+        expect(new Set(refreshed).size).toBe(1)
+        expect(refreshed[0]).not.toBe(first[0])
+        expect(await provider.getToken()).toBe(refreshed[0]!)
+    })
+
+    test('can retry signing after a failed token generation', async () => {
+        const provider = makeProvider(() => 1_700_000_000_000)
+        const sign = spyOn(SignJWT.prototype, 'sign')
+            .mockRejectedValueOnce(new Error('test signing failure'))
+        try {
+            await expect(provider.getToken()).rejects.toThrow('test signing failure')
+            const token = await provider.getToken()
+            expect(decodeProtectedHeader(token).kid).toBe(KEY_ID)
+            expect(sign).toHaveBeenCalledTimes(2)
+        } finally {
+            sign.mockRestore()
+        }
+    })
 })
 
 describe('Http2ApnsClient against a node:http2 mock', () => {
@@ -204,6 +231,17 @@ describe('Http2ApnsClient against a node:http2 mock', () => {
             .toBe(mock.requests[1]!.headers.authorization)
     })
 
+    test('shares one authorization token when first pushes arrive concurrently', async () => {
+        const c = makeClient()
+        const outcomes = await Promise.all(Array.from({ length: 8 }, (_, i) => c.push(
+            basePush({ deviceToken: i.toString(16).padStart(64, '0') })
+        )))
+        expect(outcomes.every((result) => result.kind === 'delivered')).toBe(true)
+        expect(mock.requests.length).toBe(8)
+        expect(mock.sessionCount).toBe(1)
+        expect(new Set(mock.requests.map((request) => request.headers.authorization)).size).toBe(1)
+    })
+
     test('maps APNs 410 Unregistered', async () => {
         mock.queueStatus(410, { reason: 'Unregistered', timestamp: 1_660_000_000_000 })
         const result = await makeClient().push(basePush())

relay/src/apns.ts

@@ -2,11 +2,11 @@
  * APNs client for the HAPI push relay: ES256 provider-token auth (jose) plus
  * a persistent HTTP/2 connection (node:http2).
  *
- * NOTE - deliberate duplication: hub/src/apns/ (work package P1) carries the
+ * NOTE - deliberate duplication: hub/src/push-ios/apnsClient.ts carries the
  * hub's own APNs JWT + HTTP/2 client for hubs that self-configure an APNs
  * key. The two copies are kept separate on purpose so the relay stays a
  * standalone, dependency-light deployable that never imports hub code. If
- * you fix a protocol bug here, check the twin in hub/src/apns/.
+ * you fix a protocol bug here, check the twin in hub/src/push-ios/.
  *
  * Bun compatibility: verified on Bun 1.3.14 - Bun's node:http2 client talks
  * to a real node:http2 server (the tests in apns.test.ts exercise exactly
@@ -44,7 +44,7 @@ export type ApnsJwtProviderOptions = {
 }
 
 export class ApnsJwtProvider {
-    private cached: { token: string; issuedAtMs: number } | undefined
+    private cached: { token: Promise<string>; issuedAtMs: number } | undefined
     private keyPromise: Promise<ApnsSigningKey> | undefined
 
     constructor(private readonly options: ApnsJwtProviderOptions) {}
@@ -54,14 +54,27 @@ export class ApnsJwtProvider {
         if (this.cached !== undefined && nowMs - this.cached.issuedAtMs <= APNS_JWT_MAX_AGE_MS) {
             return this.cached.token
         }
+        // Cache the signing work before awaiting it so concurrent pushes use
+        // the same JWT, including on startup and when the old token expires.
+        const token = this.signToken(nowMs)
+        this.cached = { token, issuedAtMs: nowMs }
+        try {
+            return await token
+        } catch (error) {
+            if (this.cached?.token === token) {
+                this.cached = undefined
+            }
+            throw error
+        }
+    }
+
+    private async signToken(nowMs: number): Promise<string> {
         const key = await this.getKey()
-        const token = await new SignJWT({})
+        return new SignJWT({})
             .setProtectedHeader({ alg: 'ES256', kid: this.options.keyId })
             .setIssuer(this.options.teamId)
             .setIssuedAt(Math.floor(nowMs / 1000))
             .sign(key)
-        this.cached = { token, issuedAtMs: nowMs }
-        return token
     }
 
     /** Drop the cached token so the next getToken() mints a fresh one. */

relay/src/index.ts

@@ -172,12 +172,16 @@ function mapApnsResult(result: ApnsPushResult): { response: Response; outcome: s
     if (unregistered) {
         return {
             response: pushError(410, 'unregistered'),
-            outcome: `unregistered (apns ${result.status})`
+            outcome: `unregistered (apns ${result.status} ${result.reason})`
         }
     }
     if (result.status === 429) {
-        // APNs itself throttled this device token; surface as retryable.
-        return { response: pushError(429, 'rate_limited'), outcome: 'apns-throttled' }
+        // APNs can throttle delivery or provider-token updates. Keep the
+        // reason in the operator log so the two causes can be distinguished.
+        return {
+            response: pushError(429, 'rate_limited'),
+            outcome: `apns-throttled (apns ${result.status} ${result.reason})`
+        }
     }
     // Everything else (APNs 5xx, and 4xx caused by relay config such as
     // BadTopic / auth problems) is an upstream failure from the hub's view.

relay/src/push.test.ts

@@ -174,6 +174,7 @@ describe('POST /v1/push forwarding', () => {
         { name: 'APNs 400 BadTopic (relay misconfig)', result: { kind: 'rejected', status: 400, reason: 'BadTopic' }, status: 502, code: 'upstream' },
         { name: 'APNs 403 InvalidProviderToken', result: { kind: 'rejected', status: 403, reason: 'InvalidProviderToken' }, status: 502, code: 'upstream' },
         { name: 'APNs 429 TooManyRequests', result: { kind: 'rejected', status: 429, reason: 'TooManyRequests' }, status: 429, code: 'rate_limited' },
+        { name: 'APNs 429 TooManyProviderTokenUpdates', result: { kind: 'rejected', status: 429, reason: 'TooManyProviderTokenUpdates' }, status: 429, code: 'rate_limited' },
         { name: 'APNs 500', result: { kind: 'rejected', status: 500, reason: 'InternalServerError' }, status: 502, code: 'upstream' },
         { name: 'APNs 503', result: { kind: 'rejected', status: 503, reason: 'ServiceUnavailable' }, status: 502, code: 'upstream' },
         { name: 'network failure', result: { kind: 'transport-error', message: 'connect ECONNREFUSED' }, status: 502, code: 'upstream' }
@@ -195,6 +196,23 @@ describe('POST /v1/push forwarding', () => {
         }
     })
 
+    test.each([
+        { status: 429, reason: 'TooManyRequests', outcome: 'apns-throttled' },
+        { status: 429, reason: 'TooManyProviderTokenUpdates', outcome: 'apns-throttled' },
+        { status: 400, reason: 'BadDeviceToken', outcome: 'unregistered' },
+        { status: 410, reason: 'Unregistered', outcome: 'unregistered' }
+    ])('logs the APNs reason for $status $reason', async ({ status, reason, outcome }) => {
+        const { app, apns, logs } = makeHarness()
+        apns.nextResults.push({ kind: 'rejected', status, reason })
+        await app.handle(
+            pushRequest({ platform: 'ios', token: TOKEN, envelope: ENVELOPE }),
+            '1.2.3.4'
+        )
+        expect(logs).toEqual([
+            `[relay] push token=${hashedTokenPrefix(TOKEN)} outcome=${outcome} (apns ${status} ${reason})`
+        ])
+    })
+
     test('never logs the envelope or the raw device token', async () => {
         const { app, logs } = makeHarness()
         await app.handle(