heygo/tests/api.oauth.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import worker from '../worker/index';
import { AUTH_SESSION_COOKIE_NAME } from '../worker/auth';
import type { Env } from '../worker/env';

type UserRow = {
  id: string;
  email: string;
  name: string | null;
  image_url: string | null;
  role: 'user' | 'admin';
};

type OAuthAccountRow = {
  id: string;
  user_id: string;
  provider: string;
  provider_account_id: string;
};

type SessionRow = {
  id: string;
  user_id: string;
  session_token_hash: string;
  expires_at: string;
};

class FakeOAuthD1 {
  readonly users: UserRow[] = [];
  readonly oauthAccounts: OAuthAccountRow[] = [];
  readonly sessions: SessionRow[] = [];

  prepare(sql: string): FakeOAuthStatement {
    return new FakeOAuthStatement(this, sql);
  }
}

class FakeOAuthStatement {
  private params: unknown[] = [];

  constructor(
    private readonly db: FakeOAuthD1,
    private readonly sql: string,
  ) {}

  bind(...params: unknown[]): this {
    this.params = params;
    return this;
  }

  async first<T>(): Promise<T | null> {
    if (this.sql.includes('FROM oauth_accounts JOIN users')) {
      const [provider, providerAccountId] = this.params.map(String);
      const account = this.db.oauthAccounts.find(
        (row) => row.provider === provider && row.provider_account_id === providerAccountId,
      );
      const user = account ? this.db.users.find((row) => row.id === account.user_id) : undefined;
      return (user ?? null) as T | null;
    }

    if (this.sql.includes('FROM users WHERE lower(email)')) {
      const email = String(this.params[0]).toLowerCase();
      const user = this.db.users.find((row) => row.email.toLowerCase() === email);
      return (user ?? null) as T | null;
    }

    return null;
  }

  async run(): Promise<D1Result> {
    if (this.sql.startsWith('INSERT INTO users')) {
      const [id, email, name, imageUrl, role] = this.params;
      this.db.users.push({
        id: String(id),
        email: String(email),
        name: name == null ? null : String(name),
        image_url: imageUrl == null ? null : String(imageUrl),
        role: role as 'user' | 'admin',
      });
    }

    if (this.sql.startsWith('UPDATE users SET name')) {
      const [name, imageUrl, id] = this.params;
      const user = this.db.users.find((row) => row.id === id);
      if (user) {
        if (name != null) user.name = String(name);
        if (imageUrl != null) user.image_url = String(imageUrl);
      }
    }

    if (this.sql.startsWith('UPDATE users SET role')) {
      const [role, id] = this.params;
      const user = this.db.users.find((row) => row.id === id);
      if (user) {
        user.role = role as 'user' | 'admin';
      }
    }

    if (this.sql.startsWith('INSERT OR IGNORE INTO oauth_accounts')) {
      const [id, userId, provider, providerAccountId] = this.params.map(String);
      const exists = this.db.oauthAccounts.some(
        (row) => row.provider === provider && row.provider_account_id === providerAccountId,
      );
      if (!exists) {
        this.db.oauthAccounts.push({ id, user_id: userId, provider, provider_account_id: providerAccountId });
      }
    }

    if (this.sql.startsWith('INSERT INTO sessions')) {
      const [id, userId, sessionTokenHash, expiresAt] = this.params.map(String);
      this.db.sessions.push({
        id,
        user_id: userId,
        session_token_hash: sessionTokenHash,
        expires_at: expiresAt,
      });
    }

    return { success: true, meta: {} } as D1Result;
  }
}

class FakeExecutionContext {
  waitUntil(): void {}
  passThroughOnException(): void {}
}

function makeEnv(db: FakeOAuthD1): Env {
  return {
    DB: db as unknown as D1Database,
    PUBLIC_HOST: 'dev.heygo.cc',
    PRIVATE_HOST: 'my.dev.heygo.cc',
    APP_BASE_URL: 'https://dev.heygo.cc',
    COOKIE_DOMAIN: '.heygo.cc',
    ADMIN_EMAILS: 'wahyd4@gmail.com',
    GOOGLE_CLIENT_ID: 'google-client-id',
    GOOGLE_CLIENT_SECRET: 'google-client-secret',
  };
}

async function fetchWorker(path: string, env: Env, init: RequestInit = {}): Promise<Response> {
  const request = new Request(`https://dev.heygo.cc${path}`, init);
  return worker.fetch(
    request as unknown as Parameters<typeof worker.fetch>[0],
    env as unknown as Parameters<typeof worker.fetch>[1],
    new FakeExecutionContext() as unknown as Parameters<typeof worker.fetch>[2],
  );
}

describe('Google OAuth API', () => {
  beforeEach(() => {
    vi.stubGlobal(
      'fetch',
      vi.fn(async (input: RequestInfo | URL) => {
        const url = String(input);
        if (url === 'https://oauth2.googleapis.com/token') {
          return Response.json({ access_token: 'google-access-token', token_type: 'Bearer', expires_in: 3600 });
        }
        if (url === 'https://openidconnect.googleapis.com/v1/userinfo') {
          return Response.json({
            sub: 'google-user-1',
            email: 'wahyd4@gmail.com',
            email_verified: true,
            name: 'Wahyd',
            picture: 'https://example.com/avatar.png',
          });
        }
        return new Response('not found', { status: 404 });
      }),
    );
  });

  afterEach(() => {
    vi.unstubAllGlobals();
  });

  it('redirects to Google with a state cookie', async () => {
    const db = new FakeOAuthD1();
    const env = makeEnv(db);

    const response = await fetchWorker('/api/auth/google', env);

    expect(response.status).toBe(302);
    const location = new URL(response.headers.get('location')!);
    expect(location.origin + location.pathname).toBe('https://accounts.google.com/o/oauth2/v2/auth');
    expect(location.searchParams.get('client_id')).toBe('google-client-id');
    expect(location.searchParams.get('redirect_uri')).toBe('https://dev.heygo.cc/api/auth/google/callback');
    expect(location.searchParams.get('scope')).toBe('openid email profile');
    expect(location.searchParams.get('state')).toMatch(/^[0-9a-f]{64}$/);
    expect(response.headers.get('set-cookie')).toContain('heygo_oauth_state=');
    expect(response.headers.get('set-cookie')).toContain('HttpOnly');
  });

  it('creates an admin user and session on callback', async () => {
    const db = new FakeOAuthD1();
    const env = makeEnv(db);
    const state = 'a'.repeat(64);

    const response = await fetchWorker(
      `/api/auth/google/callback?code=oauth-code&state=${state}`,
      env,
      { headers: { cookie: `heygo_oauth_state=${state}` } },
    );

    expect(response.status).toBe(302);
    expect(response.headers.get('location')).toBe('https://dev.heygo.cc/app/private');
    expect(response.headers.get('set-cookie')).toContain(`${AUTH_SESSION_COOKIE_NAME}=`);
    expect(db.users).toHaveLength(1);
    expect(db.users[0]).toMatchObject({
      email: 'wahyd4@gmail.com',
      name: 'Wahyd',
      image_url: 'https://example.com/avatar.png',
      role: 'admin',
    });
    expect(db.oauthAccounts).toHaveLength(1);
    expect(db.oauthAccounts[0]).toMatchObject({
      user_id: db.users[0].id,
      provider: 'google',
      provider_account_id: 'google-user-1',
    });
    expect(db.sessions).toHaveLength(1);
  });

  it('rejects a callback with an invalid state', async () => {
    const db = new FakeOAuthD1();
    const env = makeEnv(db);

    const response = await fetchWorker(
      '/api/auth/google/callback?code=oauth-code&state=wrong',
      env,
      { headers: { cookie: `heygo_oauth_state=${'a'.repeat(64)}` } },
    );

    expect(response.status).toBe(302);
    expect(response.headers.get('location')).toContain('/app/login');
    expect(db.users).toHaveLength(0);
    expect(db.sessions).toHaveLength(0);
  });
});