heygo/worker/routes/api.oauth.ts

import {
  AUTH_SESSION_COOKIE_NAME,
  getSessionCookieAttributes,
  hashSessionToken,
  isConfiguredAdminEmail,
} from '../auth';
import type { Env } from '../env';

const GOOGLE_AUTH_URL = 'https://accounts.google.com/o/oauth2/v2/auth';
const GOOGLE_TOKEN_URL = 'https://oauth2.googleapis.com/token';
const GOOGLE_USERINFO_URL = 'https://openidconnect.googleapis.com/v1/userinfo';
const GOOGLE_SCOPES = 'openid email profile';
const OAUTH_STATE_COOKIE_NAME = 'heygo_oauth_state';
const SESSION_TTL_DAYS = 30;

const FIND_OAUTH_USER = `SELECT users.id, users.email, users.name, users.image_url, users.role
FROM oauth_accounts JOIN users ON users.id = oauth_accounts.user_id
WHERE oauth_accounts.provider = ? AND oauth_accounts.provider_account_id = ?
LIMIT 1`;
const FIND_USER_BY_EMAIL = `SELECT id, email, name, image_url, role FROM users WHERE lower(email) = lower(?) LIMIT 1`;
const CREATE_USER = `INSERT INTO users (id, email, name, image_url, role) VALUES (?, ?, ?, ?, ?)`;
const UPDATE_USER_PROFILE = `UPDATE users SET name = COALESCE(?, name), image_url = COALESCE(?, image_url), updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') WHERE id = ?`;
const UPDATE_USER_ROLE = `UPDATE users SET role = ?, updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') WHERE id = ?`;
const CREATE_OAUTH_ACCOUNT = `INSERT OR IGNORE INTO oauth_accounts (id, user_id, provider, provider_account_id) VALUES (?, ?, ?, ?)`;
const CREATE_SESSION = `INSERT INTO sessions (id, user_id, session_token_hash, expires_at) VALUES (?, ?, ?, ?)`;

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

type GoogleTokenResponse = {
  access_token?: string;
  token_type?: string;
  expires_in?: number;
  id_token?: string;
  error?: string;
  error_description?: string;
};

type GoogleUserInfo = {
  sub?: string;
  email?: string;
  email_verified?: boolean;
  name?: string;
  picture?: string;
};

export async function handleOAuthApi(request: Request, env: Env): Promise<Response | null> {
  const url = new URL(request.url);

  if (url.pathname === '/api/auth/google') {
    return handleGoogleStart(request, env);
  }

  if (url.pathname === '/api/auth/google/callback') {
    return handleGoogleCallback(request, env);
  }

  return null;
}

async function handleGoogleStart(request: Request, env: Env): Promise<Response> {
  if (request.method !== 'GET') {
    return Response.json({ error: 'Method not allowed' }, { status: 405 });
  }

  const credentialError = requireGoogleCredentials(env);
  if (credentialError) {
    return credentialError;
  }

  const requestUrl = new URL(request.url);
  const state = generateToken();
  const redirectUri = googleRedirectUri(env);
  const loginHint = requestUrl.searchParams.get('login_hint') ?? env.ADMIN_EMAILS?.split(/[\s,;]+/)[0] ?? '';
  const googleUrl = new URL(GOOGLE_AUTH_URL);
  googleUrl.searchParams.set('client_id', env.GOOGLE_CLIENT_ID!);
  googleUrl.searchParams.set('redirect_uri', redirectUri);
  googleUrl.searchParams.set('response_type', 'code');
  googleUrl.searchParams.set('scope', GOOGLE_SCOPES);
  googleUrl.searchParams.set('state', state);
  googleUrl.searchParams.set('include_granted_scopes', 'true');
  if (loginHint) {
    googleUrl.searchParams.set('login_hint', loginHint);
  }

  return new Response(null, {
    status: 302,
    headers: {
      location: googleUrl.toString(),
      'set-cookie': oauthStateCookie(request.url, env, state),
      'cache-control': 'no-store',
    },
  });
}

async function handleGoogleCallback(request: Request, env: Env): Promise<Response> {
  if (request.method !== 'GET') {
    return Response.json({ error: 'Method not allowed' }, { status: 405 });
  }

  const credentialError = requireGoogleCredentials(env);
  if (credentialError) {
    return credentialError;
  }

  const url = new URL(request.url);
  const error = url.searchParams.get('error');
  if (error) {
    return redirectWithAuthError(env, `Google sign-in failed: ${error}`);
  }

  const code = url.searchParams.get('code');
  const state = url.searchParams.get('state');
  const storedState = parseCookieHeader(request.headers.get('cookie'))[OAUTH_STATE_COOKIE_NAME];
  if (!code || !state || !storedState || !timingSafeEqual(state, storedState)) {
    return redirectWithAuthError(env, 'Invalid sign-in state. Please try again.');
  }

  const token = await exchangeGoogleCode(env, code);
  if (!token.access_token) {
    return redirectWithAuthError(env, token.error_description ?? token.error ?? 'Google token exchange failed.');
  }

  const profile = await fetchGoogleUserInfo(token.access_token);
  if (!profile.sub || !profile.email || profile.email_verified !== true) {
    return redirectWithAuthError(env, 'Google account email is not verified.');
  }

  const user = await findOrCreateOAuthUser(env, profile);
  const sessionToken = generateToken();
  const tokenHash = await hashSessionToken(sessionToken);
  const expiresAt = new Date(Date.now() + SESSION_TTL_DAYS * 24 * 60 * 60 * 1000).toISOString();
  await env.DB.prepare(CREATE_SESSION).bind(crypto.randomUUID(), user.id, tokenHash, expiresAt).run();

  const sessionCookie = `${AUTH_SESSION_COOKIE_NAME}=${sessionToken}; ${getSessionCookieAttributes(url, env.COOKIE_DOMAIN).join('; ')}`;
  const clearStateCookie = oauthStateCookie(request.url, env, '', true);
  return new Response(null, {
    status: 302,
    headers: [
      ['location', `${env.APP_BASE_URL}/app/private`],
      ['set-cookie', sessionCookie],
      ['set-cookie', clearStateCookie],
      ['cache-control', 'no-store'],
    ],
  });
}

function requireGoogleCredentials(env: Env): Response | null {
  if (!env.GOOGLE_CLIENT_ID || !env.GOOGLE_CLIENT_SECRET) {
    return Response.json(
      { error: 'Google OAuth is not configured' },
      { status: 503, headers: { 'cache-control': 'no-store' } },
    );
  }
  return null;
}

async function exchangeGoogleCode(env: Env, code: string): Promise<GoogleTokenResponse> {
  const body = new URLSearchParams({
    code,
    client_id: env.GOOGLE_CLIENT_ID!,
    client_secret: env.GOOGLE_CLIENT_SECRET!,
    redirect_uri: googleRedirectUri(env),
    grant_type: 'authorization_code',
  });

  const response = await fetch(GOOGLE_TOKEN_URL, {
    method: 'POST',
    headers: { 'content-type': 'application/x-www-form-urlencoded' },
    body,
  });
  return (await response.json()) as GoogleTokenResponse;
}

async function fetchGoogleUserInfo(accessToken: string): Promise<GoogleUserInfo> {
  const response = await fetch(GOOGLE_USERINFO_URL, {
    headers: { authorization: `Bearer ${accessToken}` },
  });
  if (!response.ok) {
    return {};
  }
  return (await response.json()) as GoogleUserInfo;
}

async function findOrCreateOAuthUser(env: Env, profile: GoogleUserInfo): Promise<OAuthUserRow> {
  const provider = 'google';
  const providerAccountId = profile.sub!;
  const email = profile.email!;
  const configuredAdmin = isConfiguredAdminEmail(email, env);
  const role = configuredAdmin ? 'admin' : 'user';

  let user = await env.DB.prepare(FIND_OAUTH_USER).bind(provider, providerAccountId).first<OAuthUserRow>();
  if (!user) {
    user = await env.DB.prepare(FIND_USER_BY_EMAIL).bind(email).first<OAuthUserRow>();
  }

  if (!user) {
    user = {
      id: crypto.randomUUID(),
      email,
      name: profile.name ?? email.split('@')[0],
      image_url: profile.picture ?? null,
      role,
    };
    await env.DB.prepare(CREATE_USER).bind(user.id, user.email, user.name, user.image_url, user.role).run();
  } else {
    await env.DB.prepare(UPDATE_USER_PROFILE).bind(profile.name ?? null, profile.picture ?? null, user.id).run();
    if (configuredAdmin && user.role !== 'admin') {
      await env.DB.prepare(UPDATE_USER_ROLE).bind('admin', user.id).run();
      user = { ...user, role: 'admin' };
    }
  }

  await env.DB.prepare(CREATE_OAUTH_ACCOUNT)
    .bind(crypto.randomUUID(), user.id, provider, providerAccountId)
    .run();

  return user;
}

function googleRedirectUri(env: Env): string {
  return `${env.APP_BASE_URL}/api/auth/google/callback`;
}

function oauthStateCookie(requestUrl: string | URL, env: Env, value: string, expired = false): string {
  const attrs = getSessionCookieAttributes(requestUrl, env.COOKIE_DOMAIN);
  if (expired) {
    attrs.push('Expires=Thu, 01 Jan 1970 00:00:00 GMT');
  } else {
    attrs.push('Max-Age=600');
  }
  return `${OAUTH_STATE_COOKIE_NAME}=${value}; ${attrs.join('; ')}`;
}

function parseCookieHeader(cookieHeader: string | null): Record<string, string> {
  const cookies: Record<string, string> = {};
  if (!cookieHeader) {
    return cookies;
  }
  for (const part of cookieHeader.split(';')) {
    const index = part.indexOf('=');
    if (index === -1) {
      continue;
    }
    cookies[part.slice(0, index).trim()] = part.slice(index + 1).trim();
  }
  return cookies;
}

function generateToken(): string {
  const bytes = new Uint8Array(32);
  crypto.getRandomValues(bytes);
  return Array.from(bytes, (byte) => byte.toString(16).padStart(2, '0')).join('');
}

function timingSafeEqual(a: string, b: string): boolean {
  if (a.length !== b.length) {
    return false;
  }
  let diff = 0;
  for (let i = 0; i < a.length; i += 1) {
    diff |= a.charCodeAt(i) ^ b.charCodeAt(i);
  }
  return diff === 0;
}

function redirectWithAuthError(env: Env, message: string): Response {
  const location = new URL(`${env.APP_BASE_URL}/app/login`);
  location.searchParams.set('error', message);
  return Response.redirect(location.toString(), 302);
}