fix(relay): read control-error in every state and pace repeated rate limits (#27160)

Co-authored-by: m4air <m4air@m4airs-MacBook-Air.local>

src/main/runtime/relay/relay-auth-coordinator-readiness.test.ts

@@ -4,7 +4,7 @@ import { DesktopRelayService } from './desktop-relay-service'
 import { RelayAuthCoordinator } from './relay-auth-coordinator'
 import type { RelayAuthContext } from './relay-auth-identity'
 import { RelayHttpError } from './relay-http-client'
-import { relayUnavailableReasonFor } from './relay-readiness'
+import { relayRetryFloorMs, relayUnavailableReasonFor } from './relay-readiness'
 import { RelaySessionBroker } from './relay-session-broker'
 
 const context: RelayAuthContext = {
@@ -172,6 +172,41 @@ describe('relay readiness', () => {
     }
   )
 
+  it('floors a limit-coded or repeated rate limit, but keeps a lone 4429 fast (#21203)', () => {
+    const limitCoded = new Error('relay_control_error_limit_exceeded')
+    const closed4429 = new Error('relay_control_closed_4429')
+    expect(relayUnavailableReasonFor(limitCoded)).toBe('control_error_limit_exceeded')
+    expect(relayRetryFloorMs(limitCoded, 0)).toBe(5_000)
+    expect(relayRetryFloorMs(closed4429, 0)).toBe(0)
+    expect(relayRetryFloorMs(closed4429, 1)).toBe(5_000)
+    expect(relayRetryFloorMs(closed4429, 2)).toBe(10_000)
+    expect(relayRetryFloorMs(closed4429, 9)).toBe(30_000)
+    expect(relayRetryFloorMs(new Error('relay_control_error_unknown_control_message'), 3)).toBe(0)
+    expect(relayRetryFloorMs(new RelayHttpError('assignment', 503, 7_000), 0)).toBe(7_000)
+  })
+
+  it('paces a coordinator whose control keeps closing 4429', async () => {
+    vi.useFakeTimers()
+    const openBroker = vi.fn(async () => {
+      throw new Error('relay_control_closed_4429')
+    })
+    const coordinator = new RelayAuthCoordinator({
+      readContext: async () => context,
+      openBroker,
+      onStatus: () => {},
+      random: () => 0
+    })
+    coordinator.reconcile()
+    // A lone 4429 retries at once (zero jitter); the repeat then waits 5 s.
+    await vi.advanceTimersByTimeAsync(0)
+    expect(openBroker).toHaveBeenCalledTimes(2)
+    await vi.advanceTimersByTimeAsync(4_999)
+    expect(openBroker).toHaveBeenCalledTimes(2)
+    await vi.advanceTimersByTimeAsync(1)
+    expect(openBroker).toHaveBeenCalledTimes(3)
+    coordinator.stop()
+  })
+
   it('classifies transport failures into actionable reasons', () => {
     const tls = new TypeError('fetch failed', {
       cause: Object.assign(new Error('unable to get local issuer certificate'), {

src/main/runtime/relay/relay-auth-coordinator.ts

@@ -5,9 +5,10 @@ import {
 import { relayStatusCellUrl } from '../../../shared/mobile-relay-status'
 import type { RelayBrokerStatus } from './relay-session-broker'
 import type { RelayAccessTokenRefresh } from './relay-session-broker-contract'
-import { RelayHttpError, shouldRetryRelayConnectionError } from './relay-http-client'
+import { shouldRetryRelayConnectionError } from './relay-http-client'
 import {
   RelayReadinessWaiters,
+  relayRetryFloorMs,
   relayUnavailableReasonFor,
   type RelayUnavailable
 } from './relay-readiness'
@@ -278,7 +279,7 @@ export class RelayAuthCoordinator {
         this.unavailableReason = relayUnavailableReasonFor(error)
         this.publish('offline')
         if (shouldRetryRelayConnectionError(error)) {
-          const retryAfterMs = error instanceof RelayHttpError ? (error.retryAfterMs ?? 0) : 0
+          const retryAfterMs = relayRetryFloorMs(error, this.retry.attempts)
           this.scheduleRetry(epoch, retryIdentityKey, retryAfterMs)
         }
       }

src/main/runtime/relay/relay-control-client.test.ts

@@ -666,6 +666,29 @@ describe('RelayControlClient scripted-socket lifecycle', () => {
     warn.mockRestore()
   })
 
+  it('names a control-error during the proof instead of failing it as a bad credential', async () => {
+    const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
+    const { client, socket } = scriptedControl()
+    const closeCodes: number[] = []
+    socket.script = (message, ws) => {
+      if (message.type === 'host-hello') {
+        // An extra key from a newer relay must not turn the frame into an unknown one.
+        ws.deliver({ type: 'control-error', code: 'limit_exceeded', retryAfterMs: 1_000 })
+        expect(ws.readyState).toBe(1)
+        ws.close(MOBILE_RELAY_CLOSE_CODE.LIMIT_EXCEEDED)
+      }
+    }
+    const close = socket.close.bind(socket)
+    socket.close = (code?: number) => {
+      closeCodes.push(code ?? 1000)
+      close(code)
+    }
+
+    await expect(client.connect()).rejects.toThrow('relay_control_error_limit_exceeded')
+    expect(closeCodes).toEqual([MOBILE_RELAY_CLOSE_CODE.LIMIT_EXCEEDED])
+    warn.mockRestore()
+  })
+
   it('still opens a connection the relay handed over before it asked us to drain', async () => {
     const { client, socket, onConnectionOpen } = scriptedControl()
     await client.connect()

src/main/runtime/relay/relay-control-client.ts

@@ -5,6 +5,7 @@ import { MOBILE_RELAY_CLOSE_CODE } from '../../../shared/mobile-relay-close-code
 import type { RelayHostCloseReason } from '../../../shared/relay-host-close-reason'
 import {
   RelayConnectionOpenMessageSchema,
+  RelayControlErrorMessageSchema,
   RelayDrainMessageSchema,
   RelayHostChallengeMessageSchema,
   RelayHostHelloAckMessageSchema,
@@ -38,6 +39,8 @@ export class RelayControlClient {
   private connectResolve: ((ack: RelayHostHelloAckMessage) => void) | null = null
   private connectReject: ((error: Error) => void) | null = null
   private connectTimer: ReturnType<typeof setTimeout> | null = null
+  // The relay's own reason for the close that usually follows an unanswered control-error.
+  private controlErrorCode: string | null = null
 
   constructor(options: RelayControlClientOptions) {
     this.options = options
@@ -190,6 +193,17 @@ export class RelayControlClient {
       this.failProtocol('invalid control JSON')
       return
     }
+    const controlError = RelayControlErrorMessageSchema.safeParse(message)
+    // Why every state: in proving it used to fail the proof (4401, blaming the credential).
+    if (controlError.success) {
+      if (!this.requests.resolveMessage(message)) {
+        this.controlErrorCode = /^[a-z0-9_]{1,48}$/.test(controlError.data.code)
+          ? controlError.data.code
+          : 'unrecognized'
+        console.warn(`[relay] control-error code=${this.controlErrorCode} state=${this.state}`)
+      }
+      return
+    }
     if (this.state === 'proving') {
       this.handleProofMessage(message)
       return
@@ -293,7 +307,13 @@ export class RelayControlClient {
     this.state = 'closed'
     this.liveness.stop()
     if (wasConnecting) {
-      this.connectReject?.(new Error(`relay_control_closed_${code}`))
+      this.connectReject?.(
+        new Error(
+          this.controlErrorCode
+            ? `relay_control_error_${this.controlErrorCode}`
+            : `relay_control_closed_${code}`
+        )
+      )
       this.clearConnectPromise()
     }
     this.requests.rejectAll(new Error(`relay_control_closed_${code}`))

src/main/runtime/relay/relay-control-protocol.ts

@@ -131,13 +131,13 @@ export const RelayDeviceResumeConfirmedMessageSchema = DeviceResumeConfirmedSche
   type: z.literal('device-resume-confirmed')
 }).strict()
 
-export const RelayControlErrorMessageSchema = z
-  .object({
-    type: z.literal('control-error'),
-    reqId: OpaqueIdSchema.optional(),
-    code: z.string().min(1).max(128)
-  })
-  .strict()
+// Not strict: a relay that adds a field (e.g. a retry hint) must not make every
+// shipped desktop drop its control-error as unrecognized. Unknown keys are stripped.
+export const RelayControlErrorMessageSchema = z.object({
+  type: z.literal('control-error'),
+  reqId: OpaqueIdSchema.optional(),
+  code: z.string().min(1).max(128)
+})
 
 export type RelayPendingConnection = z.infer<typeof PendingConnectionSchema>
 export type RelayHostHelloAckMessage = z.infer<typeof RelayHostHelloAckMessageSchema>

src/main/runtime/relay/relay-readiness.ts

@@ -71,6 +71,26 @@ export function relayUnavailableReasonFor(error: unknown): RelayUnavailableReaso
   return message ? `control_${message[1].slice('control_'.length)}` : 'network'
 }
 
+const RATE_LIMIT_FLOOR_BASE_MS = 5_000
+const RATE_LIMIT_FLOOR_MAX_MS = 30_000
+
+// The earliest a reconnect may retry. A lone 4429 is the relay's "retry this
+// cell" (also sent for a transient relay-side failure) and keeps the fast retry;
+// a limit-coded refusal or a 4429 repeated after a failed retry waits 5 s, then
+// 10 s, ... up to 30 s, so a rate-limited desktop stops keeping itself limited.
+export function relayRetryFloorMs(error: unknown, priorAttempts: number): number {
+  if (error instanceof RelayHttpError) {
+    return error.retryAfterMs ?? 0
+  }
+  const reason = relayUnavailableReasonFor(error)
+  const limitCoded = reason.startsWith('control_error_') && /limit|rate/.test(reason)
+  if (!limitCoded && !(reason === 'rate_limited' && priorAttempts > 0)) {
+    return 0
+  }
+  const doublings = Math.max(0, limitCoded ? priorAttempts : priorAttempts - 1)
+  return Math.min(RATE_LIMIT_FLOOR_MAX_MS, RATE_LIMIT_FLOOR_BASE_MS * 2 ** doublings)
+}
+
 // Bounds a readiness wait by its caller's deadline and lets a fence release it,
 // without cancelling the shared reconnect work other callers still need.
 export class RelayReadinessWaiters {

src/main/runtime/relay/relay-retry-schedule.ts

@@ -15,6 +15,11 @@ export class RelayRetrySchedule {
     return this.timer !== null
   }
 
+  // Consecutive scheduled retries since the last reset (a success or a fresh demand).
+  get attempts(): number {
+    return this.attempt
+  }
+
   get retryAt(): number | null {
     return this.dueAt
   }