archify-vscode-ext/src/sync/auth.ts

import * as crypto from 'node:crypto';
import * as http from 'node:http';
import * as vscode from 'vscode';
import { SyncApi, SyncApiError } from './api';
import { Account, TokenResponse } from './protocol';

export const AUTH_PROVIDER_ID = 'archify-sync';
const SECRET_KEY = 'archifySync.sessions';
const SIGN_IN_TIMEOUT_MS = 5 * 60_000;

interface StoredSession {
  id: string;
  server: string;
  accessToken: string;
  /** Epoch ms. */
  expiresAt: number;
  refreshToken: string;
  account: Account;
}

interface PendingSignIn {
  resolve: (code: string) => void;
  reject: (error: Error) => void;
}

/**
 * "Archify Sync" in the Accounts menu: one session per server. Sign-in runs
 * in the browser and comes back through the extension's URI handler with a
 * one-time code, exchanged with a PKCE verifier. Tokens live in SecretStorage.
 * Refresh tokens are single use, so refreshes are serialised per server and
 * read the latest stored session first (another window may have refreshed).
 */
export class SyncAuth implements vscode.AuthenticationProvider, vscode.Disposable {
  private readonly changed = new vscode.EventEmitter<vscode.AuthenticationProviderAuthenticationSessionsChangeEvent>();
  readonly onDidChangeSessions = this.changed.event;
  private readonly pending = new Map<string, PendingSignIn>();
  private readonly refreshing = new Map<string, Promise<string | undefined>>();
  private readonly disposables: vscode.Disposable[];

  constructor(private readonly context: vscode.ExtensionContext, private readonly log: vscode.LogOutputChannel) {
    this.disposables = [
      vscode.authentication.registerAuthenticationProvider(AUTH_PROVIDER_ID, 'Archify Sync', this, { supportsMultipleAccounts: true }),
      // Another window signed in or out.
      context.secrets.onDidChange((e) => {
        if (e.key === SECRET_KEY) this.changed.fire({ added: [], removed: [], changed: [] });
      }),
    ];
  }

  private async load(): Promise<StoredSession[]> {
    try {
      return JSON.parse((await this.context.secrets.get(SECRET_KEY)) ?? '[]') as StoredSession[];
    } catch {
      return [];
    }
  }

  private async save(sessions: StoredSession[]): Promise<void> {
    await this.context.secrets.store(SECRET_KEY, JSON.stringify(sessions));
  }

  private toSession(s: StoredSession): vscode.AuthenticationSession {
    return {
      id: s.id,
      accessToken: s.accessToken,
      account: { id: s.account.id, label: `${s.account.display_name} (${new URL(s.server).host})` },
      scopes: [s.server],
    };
  }

  async getSessions(scopes?: readonly string[]): Promise<vscode.AuthenticationSession[]> {
    const server = scopes?.[0];
    return (await this.load()).filter((s) => !server || s.server === server).map((s) => this.toSession(s));
  }

  async createSession(scopes: readonly string[]): Promise<vscode.AuthenticationSession> {
    const server = scopes[0];
    if (!server) throw new Error('Archify Sync: say which server to sign in to.');
    return this.toSession(await this.signIn(server));
  }

  async removeSession(id: string): Promise<void> {
    const sessions = await this.load();
    const gone = sessions.find((s) => s.id === id);
    if (!gone) return;
    await this.save(sessions.filter((s) => s.id !== id));
    new SyncApi(gone.server).logout(gone.refreshToken).catch(() => undefined);
    this.changed.fire({ added: [], removed: [this.toSession(gone)], changed: [] });
  }

  async session(server: string): Promise<StoredSession | undefined> {
    return (await this.load()).find((s) => s.server === server);
  }

  async servers(): Promise<string[]> {
    return [...new Set((await this.load()).map((s) => s.server))];
  }

  /** Signs in through the browser; with an invite, the page also offers to make an account. */
  async signIn(server: string, invite?: string): Promise<StoredSession> {
    const verifier = crypto.randomBytes(32).toString('base64url');
    const challenge = crypto.createHash('sha256').update(verifier).digest('base64url');
    const state = crypto.randomBytes(16).toString('base64url');
    // On the desktop the browser comes back to a port on this machine: no
    // click needed, and it reaches this window. Remote and web editors use
    // the extension's URI handler (vscode://…/auth) instead.
    const loopback = vscode.env.remoteName === undefined && vscode.env.uiKind === vscode.UIKind.Desktop ? await this.listen().catch(() => undefined) : undefined;
    const redirectUri = loopback
      ? `http://127.0.0.1:${loopback.port}/auth`
      : (await vscode.env.asExternalUri(vscode.Uri.parse(`${vscode.env.uriScheme}://${this.context.extension.id}/auth`))).toString(true);
    const query = new URLSearchParams({
      response_type: 'code',
      redirect_uri: redirectUri,
      state,
      code_challenge: challenge,
      code_challenge_method: 'S256',
    });
    if (invite) query.set('invite', invite);

    const code = await vscode.window.withProgress(
      { location: vscode.ProgressLocation.Notification, title: `Archify: signing in to ${new URL(server).host} in the browser…`, cancellable: true },
      (_progress, cancel) =>
        new Promise<string>((resolve, reject) => {
          const timer = setTimeout(() => reject(new Error('The sign-in took too long. Try again.')), SIGN_IN_TIMEOUT_MS);
          const done = () => {
            clearTimeout(timer);
            this.pending.delete(state);
            loopback?.close();
          };
          this.pending.set(state, {
            resolve: (c) => (done(), resolve(c)),
            reject: (e) => (done(), reject(e)),
          });
          cancel.onCancellationRequested(() => this.pending.get(state)?.reject(new Error('Sign-in cancelled.')));
          void vscode.env.openExternal(vscode.Uri.parse(`${server}/authorize?${query}`, true)).then((opened) => {
            if (!opened) this.pending.get(state)?.reject(new Error('The browser did not open.'));
          });
        }),
    );
    const tokens = await new SyncApi(server).exchangeCode(code, verifier, redirectUri);
    const session = await this.store(server, tokens);
    await this.context.globalState.update('archifySync.used', true);
    this.log.info(`sync: signed in to ${server} as ${tokens.account.name}`);
    return session;
  }

  /** A one-off HTTP listener on 127.0.0.1 for the browser's way back; closed when the sign-in ends. */
  private listen(): Promise<{ port: number; close: () => void }> {
    return new Promise((resolve, reject) => {
      const server = http.createServer((req, res) => {
        const url = new URL(req.url ?? '/', 'http://127.0.0.1');
        const known = url.pathname === '/auth' && this.pending.has(url.searchParams.get('state') ?? '');
        res.writeHead(known ? 200 : 404, { 'content-type': 'text/html; charset=utf-8', 'cache-control': 'no-store' });
        res.end(
          known
            ? '<!doctype html><meta charset="utf-8"><title>Archify Sync</title><body style="font:15px system-ui;margin:48px auto;max-width:420px"><h2>Signed in</h2><p>You can close this tab and go back to the editor.</p></body>'
            : 'Not found',
        );
        if (known) this.completeSignIn(url.searchParams);
      });
      server.once('error', reject);
      server.listen(0, '127.0.0.1', () => {
        const address = server.address();
        if (!address || typeof address === 'string') return reject(new Error('no port'));
        resolve({ port: address.port, close: () => server.close() });
      });
    });
  }

  /** The browser came back: `…/auth?code=&state=`. */
  handleAuthUri(uri: vscode.Uri): void {
    this.completeSignIn(new URLSearchParams(uri.query));
  }

  private completeSignIn(query: URLSearchParams): void {
    const state = query.get('state') ?? '';
    const code = query.get('code');
    const pending = this.pending.get(state);
    if (!pending) {
      this.log.warn('sync: a sign-in came back that this window did not start');
      return;
    }
    if (code) pending.resolve(code);
    else pending.reject(new Error(query.get('error_description') ?? 'The sign-in was not completed.'));
  }

  private async store(server: string, tokens: TokenResponse, id?: string): Promise<StoredSession> {
    const session: StoredSession = {
      id: id ?? crypto.randomUUID(),
      server,
      accessToken: tokens.access_token,
      expiresAt: Date.now() + tokens.expires_in * 1000,
      refreshToken: tokens.refresh_token,
      account: tokens.account,
    };
    const others = (await this.load()).filter((s) => s.server !== server);
    await this.save([...others, session]);
    const event = id ? { added: [], removed: [], changed: [this.toSession(session)] } : { added: [this.toSession(session)], removed: [], changed: [] };
    this.changed.fire(event);
    return session;
  }

  /** A valid access token for the server, refreshed when close to expiry; undefined without a session. */
  accessToken(server: string, force = false): Promise<string | undefined> {
    const running = this.refreshing.get(server);
    if (running) return running;
    const work = this.freshToken(server, force).finally(() => this.refreshing.delete(server));
    this.refreshing.set(server, work);
    return work;
  }

  private async freshToken(server: string, force: boolean): Promise<string | undefined> {
    const session = await this.session(server);
    if (!session) return undefined;
    if (!force && session.expiresAt - 60_000 > Date.now()) return session.accessToken;
    try {
      const tokens = await new SyncApi(server).refresh(session.refreshToken);
      return (await this.store(server, tokens, session.id)).accessToken;
    } catch (error) {
      if (error instanceof SyncApiError && error.status === 401) {
        // Expired, or signed out elsewhere: the person signs in again.
        const latest = await this.session(server);
        if (latest && latest.refreshToken !== session.refreshToken) return latest.accessToken;
        await this.removeSession(session.id);
        this.log.warn(`sync: the session for ${server} ended; sign in again`);
        return undefined;
      }
      // Offline: keep the old token; the connection waits and tries again.
      this.log.info(`sync: could not refresh the token for ${server} (${String(error)})`);
      return session.accessToken;
    }
  }

  dispose(): void {
    for (const pending of this.pending.values()) pending.reject(new Error('Closed.'));
    for (const d of this.disposables) d.dispose();
    this.changed.dispose();
  }
}