convex-js PR 32: fix auth races for react native (#35392)

Most of this PR is from #29 from @erquhart

There's one more commit I added to (1) use `authUpdateAttempted` on `AuthError` instead of needing to string match now that the backend feature is in -- we should double check that we're seeing this field come through on the WS with any repros we have and (2) add even more comments to the tests.

(pushed up here because GH permissions are hard)

Co-authored-by: Sarah Shader <sarahshader@gmail.com>
Co-authored-by: Shawn Erquhart <shawn@erquh.art>
GitOrigin-RevId: 446f5ade8dee52d8b8c045e80105c26200458a8c

Author: 3 people

src/browser/sync/authentication_manager.ts

@@ -1,12 +1,14 @@
 import { Logger } from "../logging.js";
 import { LocalSyncState } from "./local_state.js";
-import { AuthError, Transition } from "./protocol.js";
+import { AuthError, IdentityVersion, Transition } from "./protocol.js";
 import jwtDecode from "jwt-decode";
 
 // setTimout uses 32 bit integer, so it can only
 // schedule about 24 days in the future.
 const MAXIMUM_REFRESH_DELAY = 20 * 24 * 60 * 60 * 1000; // 20 days
 
+const MAX_TOKEN_CONFIRMATION_ATTEMPTS = 2;
+
 /**
  * An async function returning the JWT-encoded OpenID Connect Identity Token
  * if available.
@@ -83,21 +85,24 @@ export class AuthenticationManager {
   // Shared by the BaseClient so that the auth manager can easily inspect it
   private readonly syncState: LocalSyncState;
   // Passed down by BaseClient, sends a message to the server
-  private readonly authenticate: (token: string) => void;
+  private readonly authenticate: (token: string) => IdentityVersion;
   private readonly stopSocket: () => Promise<void>;
-  private readonly restartSocket: () => void;
+  private readonly tryRestartSocket: () => void;
   private readonly pauseSocket: () => void;
   private readonly resumeSocket: () => void;
   // Passed down by BaseClient, sends a message to the server
   private readonly clearAuth: () => void;
   private readonly logger: Logger;
   private readonly refreshTokenLeewaySeconds: number;
+  // Number of times we have attempted to confirm the latest token. We retry up
+  // to `MAX_TOKEN_CONFIRMATION_ATTEMPTS` times.
+  private tokenConfirmationAttempts = 0;
   constructor(
     syncState: LocalSyncState,
     callbacks: {
-      authenticate: (token: string) => void;
+      authenticate: (token: string) => IdentityVersion;
       stopSocket: () => Promise<void>;
-      restartSocket: () => void;
+      tryRestartSocket: () => void;
       pauseSocket: () => void;
       resumeSocket: () => void;
       clearAuth: () => void;
@@ -110,7 +115,7 @@ export class AuthenticationManager {
     this.syncState = syncState;
     this.authenticate = callbacks.authenticate;
     this.stopSocket = callbacks.stopSocket;
-    this.restartSocket = callbacks.restartSocket;
+    this.tryRestartSocket = callbacks.tryRestartSocket;
     this.pauseSocket = callbacks.pauseSocket;
     this.resumeSocket = callbacks.resumeSocket;
     this.clearAuth = callbacks.clearAuth;
@@ -138,8 +143,6 @@ export class AuthenticationManager {
         hasRetried: false,
       });
       this.authenticate(token.value);
-      this._logVerbose("resuming WS after auth token fetch");
-      this.resumeSocket();
     } else {
       this.setAuthState({
         state: "initialRefetch",
@@ -148,6 +151,8 @@ export class AuthenticationManager {
       // Try again with `forceRefreshToken: true`
       await this.refetchToken();
     }
+    this._logVerbose("resuming WS after auth token fetch");
+    this.resumeSocket();
   }
 
   onTransition(serverMessage: Transition) {
@@ -176,13 +181,24 @@ export class AuthenticationManager {
     if (this.authState.state === "waitingForServerConfirmationOfFreshToken") {
       this._logVerbose("server confirmed new auth token is valid");
       this.scheduleTokenRefetch(this.authState.token);
+      this.tokenConfirmationAttempts = 0;
       if (!this.authState.hadAuth) {
         this.authState.config.onAuthChange(true);
       }
     }
   }
 
   onAuthError(serverMessage: AuthError) {
+    // If the AuthError is not due to updating the token, and we're currently
+    // waiting on the result of a token update, ignore.
+    if (
+      serverMessage.authUpdateAttempted === false &&
+      (this.authState.state === "waitingForServerConfirmationOfFreshToken" ||
+        this.authState.state === "waitingForServerConfirmationOfCachedToken")
+    ) {
+      this._logVerbose("ignoring non-auth token expired error");
+      return;
+    }
     const { baseVersion } = serverMessage;
     // Versioned AuthErrors are ignored if the client advanced to
     // a newer auth identity
@@ -201,12 +217,14 @@ export class AuthenticationManager {
   // in that we pause the WebSocket so that mutations
   // don't retry with bad auth.
   private async tryToReauthenticate(serverMessage: AuthError) {
-    // We only retry once, to avoid infinite retries
+    this._logVerbose(`attempting to reauthenticate: ${serverMessage.error}`);
     if (
       // No way to fetch another token, kaboom
       this.authState.state === "noAuth" ||
-      // We failed on a fresh token, trying another one won't help
-      this.authState.state === "waitingForServerConfirmationOfFreshToken"
+      // We failed on a fresh token. After a small number of retries, we give up
+      // and clear the auth state to avoid infinite retries.
+      (this.authState.state === "waitingForServerConfirmationOfFreshToken" &&
+        this.tokenConfirmationAttempts >= MAX_TOKEN_CONFIRMATION_ATTEMPTS)
     ) {
       this.logger.error(
         `Failed to authenticate: "${serverMessage.error}", check your server auth config`,
@@ -219,7 +237,13 @@ export class AuthenticationManager {
       }
       return;
     }
-    this._logVerbose("attempting to reauthenticate");
+    if (this.authState.state === "waitingForServerConfirmationOfFreshToken") {
+      this.tokenConfirmationAttempts++;
+      this._logVerbose(
+        `retrying reauthentication, ${MAX_TOKEN_CONFIRMATION_ATTEMPTS - this.tokenConfirmationAttempts} attempts remaining`,
+      );
+    }
+
     await this.stopSocket();
     const token = await this.fetchTokenAndGuardAgainstRace(
       this.authState.config.fetchToken,
@@ -248,7 +272,7 @@ export class AuthenticationManager {
       }
       this.setAndReportAuthFailed(this.authState.config.onAuthChange);
     }
-    this.restartSocket();
+    this.tryRestartSocket();
   }
 
   // Force refetch the token and schedule another refetch
@@ -291,12 +315,12 @@ export class AuthenticationManager {
       }
       this.setAndReportAuthFailed(this.authState.config.onAuthChange);
     }
-    // Resuming in case this refetch was triggered
-    // by an invalid cached token.
+    // Restart in case this refetch was triggered via schedule during
+    // a reauthentication attempt.
     this._logVerbose(
-      "resuming WS after auth token fetch (if currently paused)",
+      "restarting WS after auth token fetch (if currently stopped)",
     );
-    this.resumeSocket();
+    this.tryRestartSocket();
   }
 
   private scheduleTokenRefetch(token: string) {
@@ -348,6 +372,7 @@ export class AuthenticationManager {
       delay = 0;
     }
     const refetchTokenTimeoutId = setTimeout(() => {
+      this._logVerbose("running scheduled token refetch");
       void this.refetchToken();
     }, delay);
     this.setAuthState({
@@ -369,9 +394,15 @@ export class AuthenticationManager {
     },
   ) {
     const originalConfigVersion = ++this.configVersion;
+    this._logVerbose(
+      `fetching token with config version ${originalConfigVersion}`,
+    );
     const token = await fetchToken(fetchArgs);
     if (this.configVersion !== originalConfigVersion) {
       // This is a stale config
+      this._logVerbose(
+        `stale config version, expected ${originalConfigVersion}, got ${this.configVersion}`,
+      );
       return { isFromOutdatedConfig: true };
     }
     return { isFromOutdatedConfig: false, value: token };
@@ -381,6 +412,7 @@ export class AuthenticationManager {
     this.resetAuthState();
     // Bump this in case we are mid-token-fetch when we get stopped
     this.configVersion++;
+    this._logVerbose(`config version bumped to ${this.configVersion}`);
   }
 
   private setAndReportAuthFailed(
@@ -395,6 +427,31 @@ export class AuthenticationManager {
   }
 
   private setAuthState(newAuth: AuthState) {
+    const authStateForLog =
+      newAuth.state === "waitingForServerConfirmationOfFreshToken"
+        ? {
+            hadAuth: newAuth.hadAuth,
+            state: newAuth.state,
+            token: `...${newAuth.token.slice(-7)}`,
+          }
+        : { state: newAuth.state };
+    this._logVerbose(
+      `setting auth state to ${JSON.stringify(authStateForLog)}`,
+    );
+    switch (newAuth.state) {
+      case "waitingForScheduledRefetch":
+      case "notRefetching":
+      case "noAuth":
+        this.tokenConfirmationAttempts = 0;
+        break;
+      case "waitingForServerConfirmationOfFreshToken":
+      case "waitingForServerConfirmationOfCachedToken":
+      case "initialRefetch":
+        break;
+      default: {
+        const _typeCheck: never = newAuth;
+      }
+    }
     if (this.authState.state === "waitingForScheduledRefetch") {
       clearTimeout(this.authState.refetchTokenTimeoutId);
 

src/browser/sync/client.ts

@@ -278,9 +278,10 @@ export class BaseConvexClient {
         authenticate: (token) => {
           const message = this.state.setAuth(token);
           this.webSocketManager.sendMessage(message);
+          return message.baseVersion;
         },
         stopSocket: () => this.webSocketManager.stop(),
-        restartSocket: () => this.webSocketManager.restart(),
+        tryRestartSocket: () => this.webSocketManager.tryRestart(),
         pauseSocket: () => {
           this.webSocketManager.pause();
           this.state.pause();

src/browser/sync/protocol.ts

@@ -293,6 +293,10 @@ export type AuthError = {
   type: "AuthError";
   error: string;
   baseVersion: IdentityVersion;
+  // True if this error is in response to processing a new `Authenticate` message.
+  // Other AuthErrors may occur due to executing a function with expired auth and
+  // should be handled differently.
+  authUpdateAttempted: boolean;
 };
 type FatalError = {
   type: "FatalError";

src/browser/sync/web_socket_manager.ts

@@ -159,6 +159,15 @@ export class WebSocketManager {
     this.connect();
   }
 
+  private setSocketState(state: Socket) {
+    this.socket = state;
+    this._logVerbose(
+      `socket state changed: ${this.socket.state}, paused: ${
+        "paused" in this.socket ? this.socket.paused : undefined
+      }`,
+    );
+  }
+
   private connect() {
     if (this.socket.state === "terminated") {
       return;
@@ -174,11 +183,11 @@ export class WebSocketManager {
 
     const ws = new this.webSocketConstructor(this.uri);
     this._logVerbose("constructed WebSocket");
-    this.socket = {
+    this.setSocketState({
       state: "connecting",
       ws,
       paused: "no",
-    };
+    });
 
     // Kick off server inactivity timer before WebSocket connection is established
     // so we can detect cases where handshake fails.
@@ -191,11 +200,11 @@ export class WebSocketManager {
       if (this.socket.state !== "connecting") {
         throw new Error("onopen called with socket not in connecting state");
       }
-      this.socket = {
+      this.setSocketState({
         state: "ready",
         ws,
         paused: this.socket.paused === "yes" ? "uninitialized" : "no",
-      };
+      });
       this.resetServerInactivityTimeout();
       if (this.socket.paused === "no") {
         this._hasEverConnected = true;
@@ -261,8 +270,14 @@ export class WebSocketManager {
    * @returns Whether the message (might have been) sent.
    */
   sendMessage(message: ClientMessage) {
-    this._logVerbose(`sending message with type ${message.type}`);
-
+    const messageForLog = {
+      type: message.type,
+      ...(message.type === "Authenticate" && message.tokenType === "User"
+        ? {
+            value: `...${message.value.slice(-7)}`,
+          }
+        : {}),
+    };
     if (this.socket.state === "ready" && this.socket.paused === "no") {
       const encodedMessage = encodeClientMessage(message);
       const request = JSON.stringify(encodedMessage);
@@ -275,8 +290,19 @@ export class WebSocketManager {
         this.closeAndReconnect("FailedToSendMessage");
       }
       // We are not sure if this was sent or not.
+      this._logVerbose(
+        `sent message with type ${message.type}: ${JSON.stringify(
+          messageForLog,
+        )}`,
+      );
       return true;
     }
+    this._logVerbose(
+      `message not sent (socket state: ${this.socket.state}, paused: ${"paused" in this.socket ? this.socket.paused : undefined}): ${JSON.stringify(
+        messageForLog,
+      )}`,
+    );
+
     return false;
   }
 
@@ -390,7 +416,7 @@ export class WebSocketManager {
       case "connecting":
       case "ready": {
         const result = this.close();
-        this.socket = { state: "terminated" };
+        this.setSocketState({ state: "terminated" });
         return result;
       }
       default: {
@@ -428,17 +454,16 @@ export class WebSocketManager {
    * Create a new WebSocket after a previous `stop()`, unless `terminate()` was
    * called before.
    */
-  restart(): void {
+  tryRestart(): void {
     switch (this.socket.state) {
       case "stopped":
         break;
       case "terminated":
-        // If we're terminating we ignore restart
-        return;
       case "connecting":
       case "ready":
       case "disconnected":
-        throw new Error("`restart()` is only valid after `stop()`");
+        this.logger.logVerbose("Restart called without stopping first");
+        return;
       default: {
         // Enforce that the switch-case is exhaustive.
         const _: never = this.socket;