fix(ios): bound every enrollment replacement without throttling losses

Save each replacement attempt before generating a key, so a failed
attestation or save waits instead of minting a key on every retry.
Mark a request answered 404 or 409 as rejected, since its replay then
returns 404, and throttle that 404 like a 401. A replay's 404 for an
unmarked request is an ordinary loss and still replaces at once without
counting. A 400 keeps the journal again: it can be a temporary policy
rejection of a request that committed. Match the rejected journal by
challenge so a stored copy is never resumed again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Tom Brow <tomb@block.xyz>

mobile/ios/BuzzPushKit/Sources/BuzzPushKit/BuzzDevPushEnrollmentDriver.swift

@@ -407,12 +407,13 @@ public final class BuzzDevPushEnrollmentDriver {
     try await enroll(deviceToken: deviceToken, relayURL: relayURL, replacing: nil)
   }
 
-  /// `rejected` is a journal the gateway terminally rejected. It stays durable,
-  /// with its rejection count, until a fresh request overwrites it.
+  /// `rejected` names a journal the gateway terminally rejected. It stays
+  /// durable, with its replacement attempt recorded, until a fresh request
+  /// overwrites it.
   private func enroll(
     deviceToken: Data,
     relayURL: URL,
-    replacing rejected: BuzzPushPendingEnrollmentRecord?
+    replacing rejected: Replacement?
   ) async throws -> BuzzPushEndpointGrantRecord {
     precondition(!deviceToken.isEmpty, "The APNs device token must not be empty")
     let relayOrigin = try Self.relayOrigin(relayURL)
@@ -430,7 +431,10 @@ public final class BuzzDevPushEnrollmentDriver {
       gatewayOrigin: gatewayOrigin,
       relayOrigin: relayOrigin.text
     )
-    let pendingEnrollment = journal == rejected ? nil : journal
+    // Match the rejected request by its challenge, so a stored copy that differs
+    // in any other field cannot be resumed again.
+    let pendingEnrollment =
+      rejected != nil && journal?.challengeId == rejected?.journal.challengeId ? nil : journal
     if let pending = pendingEnrollment,
       pending.relayPubkey != relayPubkey || pending.endpointHash != endpointHash
         || pending.expiresAt <= nowSeconds
@@ -439,12 +443,16 @@ public final class BuzzDevPushEnrollmentDriver {
       // replacing it. A token/key change does not make a response-lost request
       // safe to discard, nor does it require revoking sibling installations.
       if pending.expiresAt > nowSeconds {
-        if case .rejected = try await resumeEnrollment(
+        if case .rejected(let replacement) = try await resumeEnrollment(
           pending,
           storedRecords: storedRecords,
           relayMetadataPubkey: nil
         ) {
-          return try await enroll(deviceToken: deviceToken, relayURL: relayURL, replacing: pending)
+          return try await enroll(
+            deviceToken: deviceToken,
+            relayURL: relayURL,
+            replacing: replacement
+          )
         }
       } else {
         try store.removePendingEnrollment(
@@ -559,7 +567,7 @@ public final class BuzzDevPushEnrollmentDriver {
         keyId: attestation.keyId,
         attestation: attestation.attestation,
         preparedAt: nowSeconds,
-        rejectedReplacements: rejected.map { ($0.rejectedReplacements ?? 0) + 1 }
+        rejectedReplacements: rejected?.rejectedReplacements
       )
       // The exact signed request is durable before the first network attempt.
       // Saving it is also what replaces a rejected journal.
@@ -574,8 +582,12 @@ public final class BuzzDevPushEnrollmentDriver {
       ) {
       case .completed(let record):
         return record
-      case .rejected:
-        return try await enroll(deviceToken: deviceToken, relayURL: relayURL, replacing: pending)
+      case .rejected(let replacement):
+        return try await enroll(
+          deviceToken: deviceToken,
+          relayURL: relayURL,
+          replacing: replacement
+        )
       }
     }
     return try await completeEnrollment(
@@ -588,7 +600,14 @@ public final class BuzzDevPushEnrollmentDriver {
   private enum Resumption {
     case completed(BuzzPushEndpointGrantRecord)
     /// The gateway terminally rejected the journal, which may now be replaced.
-    case rejected
+    case rejected(Replacement)
+  }
+
+  private struct Replacement {
+    /// The rejected journal as saved, with the replacement attempt recorded.
+    let journal: BuzzPushPendingEnrollmentRecord
+    /// The rejection count the fresh request carries.
+    let rejectedReplacements: Int
   }
 
   private func resumeEnrollment(
@@ -604,37 +623,46 @@ public final class BuzzDevPushEnrollmentDriver {
           relayMetadataPubkey: relayMetadataPubkey
         )
       )
-    } catch BuzzDevPushEnrollmentError.unexpectedStatus(
-      route: "v1/installations", _, let actual, _
-    ) where pending.gatewayInstallationHandle == nil && replacesRejectedJournal(pending, actual) {
+    } catch let error as BuzzDevPushEnrollmentError {
       // Only a terminal response permits replacing the old journal with a fresh
       // challenge and key. A fresh request's failure propagates instead of
       // recursively retrying forever.
-      return .rejected
+      guard case .unexpectedStatus(route: "v1/installations", _, let actual, _) = error,
+        pending.gatewayInstallationHandle == nil,
+        let replacement = try replacement(of: pending, status: actual)
+      else { throw error }
+      return .rejected(replacement)
     }
   }
 
   /// 404 means the exact replay found no committed installation and its
-  /// challenge is gone. The gateway returns 400 and 401 before looking up a
-  /// committed installation, so the same bytes can never recover one, even
-  /// when an earlier attempt committed and its response was lost. Every
-  /// replacement generates and attests a new App Attest key, so only a first
-  /// 404 replaces at once. Otherwise the journal must be old enough, with the
-  /// wait doubling for each consecutive replacement. That bounds new keys when
-  /// a replacement itself fails, for example with 409 against a live
-  /// installation for the same token, whose next replay then returns 404.
-  private func replacesRejectedJournal(
-    _ pending: BuzzPushPendingEnrollmentRecord,
-    _ status: Int
-  ) -> Bool {
-    guard [400, 401, 404].contains(status) else { return false }
+  /// challenge is gone, which an ordinary lost request also causes, so a
+  /// journal with no recorded rejection is replaced at once. 401 means the
+  /// gateway rejected the attestation before looking up a committed
+  /// installation, so the same bytes can never recover one. Every replacement
+  /// generates and attests a new App Attest key, so a 401, or a 404 for a
+  /// journal with a recorded rejection (such as a 409 against a live
+  /// installation for the same token, whose consumed challenge then replays as
+  /// 404), replaces only once the journal is old enough, with the wait
+  /// doubling for each rejection. Each attempt is saved before any key is
+  /// generated, so a replacement that fails also waits. A 400 can be a
+  /// temporary policy rejection of a request that did commit, so it keeps the
+  /// journal.
+  private func replacement(
+    of pending: BuzzPushPendingEnrollmentRecord,
+    status: Int
+  ) throws -> Replacement? {
     let replacements = pending.rejectedReplacements ?? 0
-    guard status != 404 || replacements > 0, let preparedAt = pending.preparedAt else {
-      return true
-    }
+    let isRejection = status == 401 || (status == 404 && replacements > 0)
+    guard isRejection || status == 404 else { return nil }
     let nowSeconds = Int64(now().timeIntervalSince1970)
-    let wait = Self.rejectedEnrollmentReplacementSeconds << min(replacements, 20)
-    return nowSeconds - preparedAt >= wait
+    if isRejection, let preparedAt = pending.preparedAt {
+      let wait = Self.rejectedEnrollmentReplacementSeconds << min(replacements, 20)
+      guard nowSeconds - preparedAt >= wait else { return nil }
+    }
+    let attempt = pending.with(preparedAt: nowSeconds, rejectedReplacements: replacements + 1)
+    try store.savePendingEnrollment(attempt)
+    return Replacement(journal: attempt, rejectedReplacements: isRejection ? replacements + 1 : 0)
   }
 
   private func completeEnrollment(
@@ -668,12 +696,28 @@ public final class BuzzDevPushEnrollmentDriver {
           route: "pending development enrollment"
         )
       }
-      installation = try await enrollInstallation(
-        challenge: Challenge(id: challengeUUID, value: challengeValue),
-        endpoint: endpoint,
-        expiresAt: pending.expiresAt,
-        attestation: BuzzDevAttestation(keyId: keyId, attestation: attestation)
-      )
+      do {
+        installation = try await enrollInstallation(
+          challenge: Challenge(id: challengeUUID, value: challengeValue),
+          endpoint: endpoint,
+          expiresAt: pending.expiresAt,
+          attestation: BuzzDevAttestation(keyId: keyId, attestation: attestation)
+        )
+      } catch let error as BuzzDevPushEnrollmentError {
+        // A request answered 409 or 404 consumed or outlived its challenge, so
+        // its replay returns 404. Record the rejection so that 404 waits instead
+        // of replacing at once. A replay's own 404 is decided on the journal as
+        // loaded, so this mark does not throttle an ordinary lost request.
+        if case .unexpectedStatus(route: "v1/installations", _, let actual, _) = error,
+          actual == 404 || actual == 409,
+          (pending.rejectedReplacements ?? 0) == 0
+        {
+          try store.savePendingEnrollment(
+            pending.with(preparedAt: pending.preparedAt, rejectedReplacements: 1)
+          )
+        }
+        throw error
+      }
       pending = BuzzPushPendingEnrollmentRecord(
         gatewayOrigin: gatewayOrigin,
         relayOrigin: pending.relayOrigin,
@@ -1115,3 +1159,26 @@ private struct RelayKeys {
   let pushPubkey: String
   let metadataPubkey: String?
 }
+
+extension BuzzPushPendingEnrollmentRecord {
+  /// This journal with only its replacement bookkeeping changed.
+  fileprivate func with(preparedAt: Int64?, rejectedReplacements: Int) -> Self {
+    Self(
+      gatewayOrigin: gatewayOrigin,
+      relayOrigin: relayOrigin,
+      relayPubkey: relayPubkey,
+      endpoint: endpoint,
+      endpointHash: endpointHash,
+      expiresAt: expiresAt,
+      installationId: installationId,
+      gatewayInstallationHandle: gatewayInstallationHandle,
+      challengeId: challengeId,
+      challenge: challenge,
+      keyId: keyId,
+      attestation: attestation,
+      delegationGeneration: delegationGeneration,
+      preparedAt: preparedAt,
+      rejectedReplacements: rejectedReplacements
+    )
+  }
+}

mobile/ios/BuzzPushKit/Sources/BuzzPushKit/BuzzPushPendingEnrollmentRecord.swift

@@ -28,10 +28,13 @@ public struct BuzzPushPendingEnrollmentRecord: Codable, Equatable, Sendable {
   public let attestation: String?
   /// Highest delegation generation reserved before a network request.
   public let delegationGeneration: Int64
-  /// Unix time the journaled installation request was prepared; nil when there
-  /// is none or the journal predates this field.
+  /// Unix time the journaled installation request was prepared, or its latest
+  /// replacement attempt began; nil when there is none or the journal predates
+  /// this field.
   public let preparedAt: Int64?
-  /// Consecutive gateway-rejected journals this one replaced; nil means none.
+  /// Gateway rejections and replacement attempts recorded for this journal and
+  /// the rejected journals it replaced, which lengthen the wait before the
+  /// next replacement; nil or 0 means none.
   public let rejectedReplacements: Int?
 
   /// Creates a gateway-scoped durable record from authenticated enrollment state.

mobile/ios/BuzzPushKit/Tests/BuzzPushKitTests/BuzzDevPushEnrollmentDriverTests.swift

@@ -395,21 +395,26 @@ final class BuzzDevPushEnrollmentDriverTests: XCTestCase {
   func testRejectedEnrollmentReplayIsReplacedOnlyOnTerminalStatus() async throws {
     // A journaled request the gateway rejects must not be replayed for its whole
     // lifetime, nor replaced (with a fresh App Attest key) on every retry; a
-    // transient failure must keep the journal. A fresh 409 models a request that
-    // committed before its response was lost and was rejected once its
-    // attestation certificate expired: the replacement conflicts with the live
-    // installation and its consumed challenge then returns 404.
+    // transient failure, or a 400 that a later policy change could lift, must
+    // keep the journal. A fresh status of 0 drops the replacement request, whose
+    // replay then finds its challenge expired: an ordinary loss, not a
+    // rejection. A fresh 409 conflicts with a live installation for the same
+    // token, for example one whose request committed before its response was
+    // lost and was rejected once its attestation certificate expired; its
+    // consumed challenge then returns 404, as does a fresh 404's replay.
     let interval = BuzzDevPushEnrollmentDriver.rejectedEnrollmentReplacementSeconds
-    let cases: [(replay: Int, fresh: Int?, elapsed: Int64, legacy: Bool)] = [
-      (400, 201, interval, false), (401, 201, interval, false), (404, 201, 0, false),
-      (401, 201, 0, true), (401, nil, interval - 1, false), (401, 401, interval, false),
-      (401, 409, interval, false), (503, nil, interval, false),
+    let cases: [(replay: Int, fresh: Int?, elapsed: Int64, legacy: Bool, keys: Int)] = [
+      (400, nil, interval, false, 1), (401, 201, interval, false, 2),
+      (404, 201, 0, false, 2), (404, 0, 0, false, 3), (404, 404, 0, false, 2),
+      (404, 409, 0, false, 2), (401, 201, 0, true, 2),
+      (401, nil, interval - 1, false, 1), (401, 401, interval, false, 2),
+      (401, 409, interval, false, 2), (503, nil, interval, false, 1),
     ]
-    for (replayStatus, freshStatus, elapsed, legacy) in cases {
-      let fresh = freshStatus.map(String.init) ?? "none"
+    for (replayStatus, freshStatus, elapsed, legacy, keys) in cases {
+      let fresh = freshStatus.map { $0 == 0 ? "dropped" : String($0) } ?? "none"
       let label = "replay \(replayStatus), fresh \(fresh) after \(elapsed)s, legacy \(legacy)"
       // Status for a later exact replay of a fresh request.
-      let freshReplayStatus = freshStatus == 409 ? 404 : freshStatus
+      let freshReplayStatus = [0, 409].contains(freshStatus) ? 404 : freshStatus
       let store = MemoryGrantStore()
       let appAttest = RecordingAppAttest()
       var clock = Self.now
@@ -420,9 +425,7 @@ final class BuzzDevPushEnrollmentDriverTests: XCTestCase {
         lifetimeSeconds: 30 * 86_400
       )
       var challengeCount = 0
-      var failChallenge = false
       var replayRequests = 0
-      var freshRequests = 0
       var sentChallenges = Set<String>()
       URLProtocolStub.handler = { request in
         switch (request.httpMethod, request.url?.absoluteString) {
@@ -431,10 +434,6 @@ final class BuzzDevPushEnrollmentDriverTests: XCTestCase {
             "push": ["keys": [["pubkey": Self.relayPubkey, "current": true]]]
           ])
         case ("POST", "http://push.example/v1/installations/challenges"):
-          if failChallenge {
-            failChallenge = false
-            throw URLError(.notConnectedToInternet)
-          }
           challengeCount += 1
           return Self.response(request, status: 200, json: [
             "challenge_id": [Self.firstChallengeId, Self.secondChallengeId].indices
@@ -451,13 +450,13 @@ final class BuzzDevPushEnrollmentDriverTests: XCTestCase {
             if replayRequests == 1 { throw URLError(.networkConnectionLost) }
             return Self.response(request, status: replayStatus, json: ["error": "rejected"])
           }
-          freshRequests += 1
           let challengeId = try XCTUnwrap(body["challenge_id"] as? String)
           let isReplay = !sentChallenges.insert(challengeId).inserted
           let status = try XCTUnwrap(
             isReplay ? freshReplayStatus : freshStatus,
             "\(label) must keep the journal"
           )
+          if status == 0 { throw URLError(.networkConnectionLost) }
           guard status == 201 else {
             return Self.response(request, status: status, json: ["error": "rejected"])
           }
@@ -504,43 +503,71 @@ final class BuzzDevPushEnrollmentDriverTests: XCTestCase {
             (attempt == 0 ? freshStatus : freshReplayStatus) ?? replayStatus,
             label
           )
-          XCTAssertEqual(
-            store.pending.first?.challengeId,
-            freshStatus == nil ? Self.firstChallengeId : Self.secondChallengeId,
-            label
-          )
+          let challengeId = try XCTUnwrap(store.pending.first?.challengeId, label)
+          if freshStatus == nil {
+            XCTAssertEqual(challengeId, Self.firstChallengeId, label)
+          } else {
+            XCTAssertNotEqual(challengeId, Self.firstChallengeId, label)
+          }
+        } catch let error as URLError where freshStatus == 0 {
+          // Each retry replaces the lost replacement at once and loses it again.
+          XCTAssertEqual(error.code, .networkConnectionLost, label)
+          XCTAssertNotEqual(store.pending.first?.challengeId, Self.firstChallengeId, label)
         }
       }
-      XCTAssertEqual(freshRequests, freshStatus == nil ? 0 : freshStatus == 201 ? 1 : 2, label)
-      XCTAssertEqual(appAttest.preparedAttestations.count, freshStatus == nil ? 1 : 2, label)
+      XCTAssertEqual(appAttest.preparedAttestations.count, keys, label)
+      if freshStatus == 0 {
+        XCTAssertEqual(store.pending.first?.rejectedReplacements ?? 0, 0, label)
+        // An ordinary replacement that fails after generating its key still
+        // records the attempt, so the next retry waits instead of generating
+        // another key, however old the journal was.
+        clock += 4 * interval
+        appAttest.failAttestation = true
+        do {
+          _ = try await driver.enroll(deviceToken: token, relayURL: Self.relayURL)
+          XCTFail("Expected the simulated attestation failure: \(label)")
+        } catch {
+          XCTAssertEqual((error as NSError).domain, NSURLErrorDomain, label)
+        }
+        await Self.assertEnrollmentStatus(driver, token, 404, label)
+        XCTAssertEqual(appAttest.preparedAttestations.count, keys + 1, label)
+      }
 
-      guard let freshStatus, let freshReplayStatus, freshStatus != 201 else { continue }
+      guard let freshStatus, let freshReplayStatus, [401, 404, 409].contains(freshStatus) else {
+        continue
+      }
       // Each consecutive replacement doubles the wait before the next one.
       XCTAssertEqual(store.pending.first?.rejectedReplacements, 1, label)
       clock += interval
       await Self.assertEnrollmentStatus(driver, token, freshReplayStatus, label)
       XCTAssertEqual(appAttest.preparedAttestations.count, 2, label)
 
-      // A replacement that fails before its request is durable keeps the
-      // rejected journal and its count, even across a driver restart.
+      // A replacement that fails after generating its key keeps the rejected
+      // journal with the attempt recorded, so retries wait instead of
+      // generating more keys, even across a driver restart.
       clock += interval
-      failChallenge = true
+      appAttest.failAttestation = true
       do {
         _ = try await driver.enroll(deviceToken: token, relayURL: Self.relayURL)
-        XCTFail("Expected the simulated challenge failure: \(label)")
+        XCTFail("Expected the simulated attestation failure: \(label)")
       } catch {
         XCTAssertEqual((error as NSError).domain, NSURLErrorDomain, label)
       }
-      XCTAssertEqual(store.pending.first?.rejectedReplacements, 1, label)
+      XCTAssertEqual(appAttest.preparedAttestations.count, 3, label)
+      XCTAssertEqual(store.pending.first?.challengeId, Self.secondChallengeId, label)
+      XCTAssertEqual(store.pending.first?.rejectedReplacements, 2, label)
       let restarted = try makeDriver(
         store: store,
         appAttest: appAttest,
         now: { Date(timeIntervalSince1970: TimeInterval(clock)) },
         lifetimeSeconds: 30 * 86_400
       )
-      await Self.assertEnrollmentStatus(restarted, token, freshStatus, label)
+      await Self.assertEnrollmentStatus(restarted, token, freshReplayStatus, label)
       XCTAssertEqual(appAttest.preparedAttestations.count, 3, label)
-      XCTAssertEqual(store.pending.first?.rejectedReplacements, 2, label)
+      clock += 4 * interval
+      await Self.assertEnrollmentStatus(restarted, token, freshStatus, label)
+      XCTAssertEqual(appAttest.preparedAttestations.count, 4, label)
+      XCTAssertEqual(store.pending.first?.rejectedReplacements, 3, label)
     }
   }
 
@@ -1487,6 +1514,8 @@ private final class RecordingAppAttest: BuzzDevAppAttesting {
   var clientData: [Data] = []
   var preparedAttestations: [BuzzDevAttestation] = []
   var assertionKeyIds: [String] = []
+  /// Fails the next attestation after its key was generated.
+  var failAttestation = false
 
   func prepareAttestation() async throws -> BuzzDevAttestation {
     let prepared = BuzzDevAttestation(
@@ -1501,6 +1530,10 @@ private final class RecordingAppAttest: BuzzDevAppAttesting {
     _ prepared: BuzzDevAttestation,
     clientData: Data
   ) async throws -> BuzzDevAttestation {
+    if failAttestation {
+      failAttestation = false
+      throw URLError(.notConnectedToInternet)
+    }
     self.clientData.append(clientData)
     return prepared
   }