From 036c317fd6d610621ef9d04a0a279d7935efd583 Mon Sep 17 00:00:00 2001
From: Lewis <lu5a@proton.me>
Date: Sun, 24 May 2026 12:47:33 +0300
Subject: [PATCH] fix(migrate): oauth refreshing, embedded db first invite code

Lewis: May this revision serve well! <lu5a@proton.me>
---
 crates/tranquil-pds/src/state.rs              |  38 +++---
 .../components/dashboard/AdminContent.svelte  |   5 +-
 frontend/src/lib/migration/atproto-client.ts  |  90 ++++++++++++-
 frontend/src/lib/migration/flow.svelte.ts     | 125 ++++++++++--------
 frontend/src/lib/migration/storage.ts         |   1 +
 frontend/src/lib/migration/types.ts           |   2 +
 frontend/src/tests/migration/storage.test.ts  |   1 +
 7 files changed, 183 insertions(+), 79 deletions(-)

diff --git a/crates/tranquil-pds/src/state.rs b/crates/tranquil-pds/src/state.rs
index 186ffec..61a9641 100644
--- a/crates/tranquil-pds/src/state.rs
+++ b/crates/tranquil-pds/src/state.rs
@@ -214,10 +214,10 @@ impl AppState {
     pub async fn new(shutdown: CancellationToken) -> Result<Self, Box<dyn Error>> {
         let cfg = tranquil_config::get();
 
-        match cfg.storage.repo_backend() {
+        let mut state = match cfg.storage.repo_backend() {
             tranquil_config::RepoBackend::TranquilStore => {
                 tracing::info!("tranquil-store repo backend active. EXPERIMENTAL!");
-                Ok(Self::from_store(shutdown).await)
+                Self::from_store(shutdown).await
             }
             tranquil_config::RepoBackend::Postgres => {
                 let database_url = &cfg.database.url;
@@ -247,28 +247,22 @@ impl AppState {
                     .await
                     .map_err(|e| format!("Failed to run migrations: {}", e))?;
 
-                let bootstrap_invite_code = match (
-                    cfg.server.invite_code_required,
-                    sqlx::query_scalar!("SELECT COUNT(*) FROM users")
-                        .fetch_one(&db)
-                        .await,
-                ) {
-                    (true, Ok(Some(0))) => {
-                        let code = crate::util::gen_invite_code();
-                        tracing::info!(
-                            "No users exist and invite codes are required. Bootstrap invite code: {}",
-                            code
-                        );
-                        Some(code)
-                    }
-                    _ => None,
-                };
-
-                let mut state = Self::from_db(db, shutdown).await;
-                state.bootstrap_invite_code = bootstrap_invite_code;
-                Ok(state)
+                Self::from_db(db, shutdown).await
             }
+        };
+
+        if cfg.server.invite_code_required
+            && state.repos.user.count_users().await.unwrap_or(1) == 0
+        {
+            let code = crate::util::gen_invite_code();
+            tracing::info!(
+                "No users exist and invite codes are required. Bootstrap invite code: {}",
+                code
+            );
+            state.bootstrap_invite_code = Some(code);
         }
+
+        Ok(state)
     }
 
     pub async fn from_db(db: PgPool, shutdown: CancellationToken) -> Self {
diff --git a/frontend/src/components/dashboard/AdminContent.svelte b/frontend/src/components/dashboard/AdminContent.svelte
index eb6f139..a964b9c 100644
--- a/frontend/src/components/dashboard/AdminContent.svelte
+++ b/frontend/src/components/dashboard/AdminContent.svelte
@@ -1,5 +1,5 @@
 <script lang="ts">
-  import { onMount } from 'svelte'
+  import { onMount, onDestroy } from 'svelte'
   import { _ } from '../../lib/i18n'
   import { api, ApiError } from '../../lib/api'
   import { toast } from '../../lib/toast.svelte'
@@ -85,9 +85,10 @@
 
   onMount(async () => {
     await Promise.all([loadStats(), loadServerConfig(), loadUsers(true), loadSignalStatus()])
-    return () => stopSignalPolling()
   })
 
+  onDestroy(() => stopSignalPolling())
+
   async function loadStats() {
     loading = true
     try {
diff --git a/frontend/src/lib/migration/atproto-client.ts b/frontend/src/lib/migration/atproto-client.ts
index 9cbb8ce..fd6e9a0 100644
--- a/frontend/src/lib/migration/atproto-client.ts
+++ b/frontend/src/lib/migration/atproto-client.ts
@@ -37,6 +37,8 @@ export class AtprotoClient {
   private dpopKeyPair: DPoPKeyPair | null = null;
   private dpopNonce: string | null = null;
   private isRefreshing = false;
+  private oauthTokenEndpoint: string | null = null;
+  private oauthClientId: string | null = null;
 
   constructor(pdsUrl: string) {
     this.baseUrl = pdsUrl.replace(/\/$/, "");
@@ -66,10 +68,26 @@ export class AtprotoClient {
     this.dpopKeyPair = keyPair;
   }
 
+  setOAuthRefreshContext(tokenEndpoint: string, clientId: string) {
+    this.oauthTokenEndpoint = tokenEndpoint;
+    this.oauthClientId = clientId;
+  }
+
   private async tryRefreshToken(): Promise<boolean> {
     if (!this.refreshToken || this.isRefreshing) return false;
     this.isRefreshing = true;
     try {
+      if (this.dpopKeyPair && this.oauthTokenEndpoint && this.oauthClientId) {
+        const tokens = await refreshSourceOAuthToken(this.oauthTokenEndpoint, {
+          refreshToken: this.refreshToken,
+          clientId: this.oauthClientId,
+          dpopKeyPair: this.dpopKeyPair,
+          nonce: this.dpopNonce ?? undefined,
+        });
+        this.accessToken = tokens.access_token;
+        this.refreshToken = tokens.refresh_token ?? this.refreshToken;
+        return true;
+      }
       const session = await this.refreshSessionInternal(this.refreshToken);
       this.accessToken = session.accessJwt;
       this.refreshToken = session.refreshJwt;
@@ -209,7 +227,7 @@ export class AtprotoClient {
         message: res.statusText,
       }));
 
-      const isTokenExpired = res.status === 401 &&
+      const isTokenExpired = (res.status === 401 || res.status === 400) &&
         (err.error === "ExpiredToken" || err.error === "invalid_token" ||
           (err.message && err.message.includes("expired")));
 
@@ -292,6 +310,7 @@ export class AtprotoClient {
     );
 
     this.accessToken = session.accessJwt;
+    this.refreshToken = session.refreshJwt;
     return session;
   }
 
@@ -448,6 +467,7 @@ export class AtprotoClient {
 
     const session = (await res.json()) as Session;
     this.accessToken = session.accessJwt;
+    this.refreshToken = session.refreshJwt;
     return session;
   }
 
@@ -589,6 +609,7 @@ export class AtprotoClient {
       },
     );
     this.accessToken = session.accessJwt;
+    this.refreshToken = session.refreshJwt;
     return session;
   }
 
@@ -925,6 +946,73 @@ export async function exchangeOAuthCode(
   return res.json();
 }
 
+export async function refreshSourceOAuthToken(
+  tokenEndpoint: string,
+  params: {
+    refreshToken: string;
+    clientId: string;
+    dpopKeyPair: DPoPKeyPair;
+    nonce?: string;
+  },
+): Promise<OAuthTokenResponse> {
+  const body = new URLSearchParams({
+    grant_type: "refresh_token",
+    refresh_token: params.refreshToken,
+    client_id: params.clientId,
+  });
+
+  const makeRequest = async (nonce?: string): Promise<Response> => {
+    const dpopProof = await createDPoPProof(
+      params.dpopKeyPair,
+      "POST",
+      tokenEndpoint,
+      nonce,
+    );
+
+    return fetch(tokenEndpoint, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/x-www-form-urlencoded",
+        "DPoP": dpopProof,
+      },
+      body: body.toString(),
+    });
+  };
+
+  let res = await makeRequest(params.nonce);
+
+  if (!res.ok) {
+    const err = await res.json().catch(() => ({
+      error: "token_error",
+      error_description: res.statusText,
+    }));
+
+    if (err.error === "use_dpop_nonce") {
+      const dpopNonce = res.headers.get("DPoP-Nonce");
+      if (dpopNonce) {
+        res = await makeRequest(dpopNonce);
+        if (!res.ok) {
+          const retryErr = await res.json().catch(() => ({
+            error: "token_error",
+            error_description: res.statusText,
+          }));
+          throw new Error(
+            retryErr.error_description || retryErr.error ||
+              "Token refresh failed",
+          );
+        }
+        return res.json();
+      }
+    }
+
+    throw new Error(
+      err.error_description || err.error || "Token refresh failed",
+    );
+  }
+
+  return res.json();
+}
+
 export async function resolveDidDocument(did: string): Promise<DidDocument> {
   if (did.startsWith("did:plc:")) {
     const res = await fetch(`https://plc.directory/${did}`);
diff --git a/frontend/src/lib/migration/flow.svelte.ts b/frontend/src/lib/migration/flow.svelte.ts
index 46ac95d..f1357d3 100644
--- a/frontend/src/lib/migration/flow.svelte.ts
+++ b/frontend/src/lib/migration/flow.svelte.ts
@@ -74,6 +74,7 @@ export function createInboundMigrationFlow() {
     passkeySetupToken: null,
     oauthCodeVerifier: null,
     localAccessToken: null,
+    localRefreshToken: null,
     generatedAppPassword: null,
     generatedAppPasswordName: null,
     handlePreservation: "new",
@@ -169,7 +170,8 @@ export function createInboundMigrationFlow() {
       redirectUri: getMigrationOAuthRedirectUri(),
       codeChallenge,
       state: oauthState,
-      scope: "atproto identity:* rpc:com.atproto.server.createAccount?aud=*",
+      scope:
+        "atproto identity:* account:repo?action=manage rpc:com.atproto.server.createAccount?aud=*",
       dpopJkt: dpopKeyPair.thumbprint,
       loginHint: state.sourceHandle,
     });
@@ -268,6 +270,10 @@ export function createInboundMigrationFlow() {
     sourceClient.setAccessToken(tokenResponse.access_token);
     sourceClient.setRefreshToken(tokenResponse.refresh_token ?? null);
     sourceClient.setDPoPKeyPair(dpopKeyPair);
+    sourceClient.setOAuthRefreshContext(
+      metadata.token_endpoint,
+      getMigrationOAuthClientId(),
+    );
 
     cleanupOAuthSessionData();
 
@@ -288,22 +294,38 @@ export function createInboundMigrationFlow() {
         if (state.localAccessToken) {
           localClient.setAccessToken(state.localAccessToken);
         }
+        if (state.localRefreshToken) {
+          localClient.setRefreshToken(state.localRefreshToken);
+        }
         if (state.authMethod === "passkey" && state.passkeySetupToken) {
           setStep("passkey-setup");
           migrationLog(
             "handleOAuthCallback: Resuming passkey flow at passkey-setup",
           );
         } else {
-          setStep("email-verify");
-          migrationLog(
-            "handleOAuthCallback: Resuming at email-verify for re-auth",
-          );
+          const alreadyVerified = await localClient
+            .checkChannelVerified(state.sourceDid, state.verificationChannel)
+            .catch(() => false);
+          if (alreadyVerified) {
+            migrationLog(
+              "handleOAuthCallback: Already verified, skipping email-verify",
+            );
+            await proceedAfterVerification();
+          } else {
+            setStep("email-verify");
+            migrationLog(
+              "handleOAuthCallback: Resuming at email-verify for re-auth",
+            );
+          }
         }
       } else if (targetStep === "email-verify") {
         localClient = createLocalClient();
         if (state.localAccessToken) {
           localClient.setAccessToken(state.localAccessToken);
         }
+        if (state.localRefreshToken) {
+          localClient.setRefreshToken(state.localRefreshToken);
+        }
         setStep("email-verify");
         migrationLog("handleOAuthCallback: Resuming at email-verify");
       } else {
@@ -459,6 +481,7 @@ export function createInboundMigrationFlow() {
         });
         localClient.setAccessToken(session.accessJwt);
         state.localAccessToken = session.accessJwt;
+        state.localRefreshToken = session.refreshJwt;
       }
 
       setProgress({ currentOperation: "Exporting repository..." });
@@ -646,50 +669,48 @@ export function createInboundMigrationFlow() {
     );
   }
 
+  async function proceedAfterVerification(): Promise<void> {
+    if (state.authMethod === "passkey") {
+      setStep("passkey-setup");
+      return;
+    }
+
+    if (!localClient!.getAccessToken()) {
+      await localClient!.loginDeactivated(
+        state.targetEmail,
+        state.targetPassword,
+      );
+    }
+
+    if (!sourceClient) {
+      setStep("source-handle");
+      setError(
+        "Email verified! Please log in to your old account again to complete the migration.",
+      );
+      return;
+    }
+
+    if (state.sourceDid.startsWith("did:web:")) {
+      const credentials = await localClient!.getRecommendedDidCredentials();
+      state.targetVerificationMethod =
+        credentials.verificationMethods?.atproto || null;
+      setStep("did-web-update");
+    } else {
+      await sourceClient.requestPlcOperationSignature();
+      setStep("plc-token");
+    }
+  }
+
   const verificationPoller = createEmailVerificationPoller({
     async checkVerified() {
       if (!localClient) return false;
-      if (state.verificationChannel === "email") {
-        return localClient.checkEmailVerified(state.targetEmail);
-      }
       return localClient.checkChannelVerified(
         state.sourceDid,
         state.verificationChannel,
       );
     },
     async onVerified() {
-      if (state.authMethod === "passkey") {
-        migrationLog(
-          "checkEmailVerifiedAndProceed: Email verified, proceeding to passkey setup",
-        );
-        setStep("passkey-setup");
-        return;
-      }
-
-      if (!localClient!.getAccessToken()) {
-        await localClient!.loginDeactivated(
-          state.targetEmail,
-          state.targetPassword,
-        );
-      }
-
-      if (!sourceClient) {
-        setStep("source-handle");
-        setError(
-          "Email verified! Please log in to your old account again to complete the migration.",
-        );
-        return;
-      }
-
-      if (state.sourceDid.startsWith("did:web:")) {
-        const credentials = await localClient!.getRecommendedDidCredentials();
-        state.targetVerificationMethod =
-          credentials.verificationMethods?.atproto || null;
-        setStep("did-web-update");
-      } else {
-        await sourceClient.requestPlcOperationSignature();
-        setStep("plc-token");
-      }
+      await proceedAfterVerification();
     },
   });
 
@@ -776,16 +797,10 @@ export function createInboundMigrationFlow() {
         });
         setProgress({ deactivated: true });
       } catch (deactivateErr) {
-        const err = deactivateErr as Error & {
-          error?: string;
-          status?: number;
-        };
-        migrationLog("Step 5 FAILED: Could not deactivate on source PDS", {
-          durationMs: Date.now() - deactivateStart,
-          error: err.message,
-          errorCode: err.error,
-          status: err.status,
-        });
+        console.error(
+          "[MIGRATION] Failed to deactivate old account on source PDS",
+          deactivateErr,
+        );
       }
 
       migrationLog("submitPlcToken SUCCESS: Migration complete", {
@@ -859,10 +874,10 @@ export function createInboundMigrationFlow() {
         });
         setProgress({ deactivated: true });
       } catch (deactivateErr) {
-        const err = deactivateErr as Error & { error?: string };
-        migrationLog("Could not deactivate on source PDS", {
-          error: err.message,
-        });
+        console.error(
+          "[MIGRATION] Failed to deactivate old account on source PDS",
+          deactivateErr,
+        );
       }
 
       migrationLog("completeDidWebMigration SUCCESS");
@@ -963,6 +978,7 @@ export function createInboundMigrationFlow() {
       passkeySetupToken: null,
       oauthCodeVerifier: null,
       localAccessToken: null,
+      localRefreshToken: null,
       generatedAppPassword: null,
       generatedAppPasswordName: null,
       handlePreservation: "new",
@@ -988,6 +1004,7 @@ export function createInboundMigrationFlow() {
     state.targetEmail = stored.targetEmail;
     state.authMethod = stored.authMethod ?? "password";
     state.localAccessToken = stored.localAccessToken ?? null;
+    state.localRefreshToken = stored.localRefreshToken ?? null;
     state.progress = {
       ...createInitialProgress(),
       ...stored.progress,
diff --git a/frontend/src/lib/migration/storage.ts b/frontend/src/lib/migration/storage.ts
index 21a2f30..7374160 100644
--- a/frontend/src/lib/migration/storage.ts
+++ b/frontend/src/lib/migration/storage.ts
@@ -23,6 +23,7 @@ export function saveMigrationState(state: MigrationState): void {
     authMethod: state.authMethod,
     passkeySetupToken: state.passkeySetupToken ?? undefined,
     localAccessToken: state.localAccessToken ?? undefined,
+    localRefreshToken: state.localRefreshToken ?? undefined,
     progress: {
       repoExported: state.progress.repoExported,
       repoImported: state.progress.repoImported,
diff --git a/frontend/src/lib/migration/types.ts b/frontend/src/lib/migration/types.ts
index e2c0028..40ce930 100644
--- a/frontend/src/lib/migration/types.ts
+++ b/frontend/src/lib/migration/types.ts
@@ -74,6 +74,7 @@ export interface InboundMigrationState {
   passkeySetupToken: string | null;
   oauthCodeVerifier: string | null;
   localAccessToken: string | null;
+  localRefreshToken: string | null;
   generatedAppPassword: string | null;
   generatedAppPasswordName: string | null;
   needsReauth?: boolean;
@@ -135,6 +136,7 @@ export interface StoredMigrationState {
   authMethod?: AuthMethod;
   passkeySetupToken?: string;
   localAccessToken?: string;
+  localRefreshToken?: string;
   progress: {
     repoExported: boolean;
     repoImported: boolean;
diff --git a/frontend/src/tests/migration/storage.test.ts b/frontend/src/tests/migration/storage.test.ts
index ba9d0fd..5f9a3b1 100644
--- a/frontend/src/tests/migration/storage.test.ts
+++ b/frontend/src/tests/migration/storage.test.ts
@@ -84,6 +84,7 @@ function createInboundState(
     passkeySetupToken: null,
     oauthCodeVerifier: null,
     localAccessToken: null,
+    localRefreshToken: null,
     generatedAppPassword: null,
     generatedAppPasswordName: null,
     handlePreservation: "new",