heygo/worker/routes/private-redirect.ts

import type { Env } from '../env';
import { getCurrentUser } from '../auth';
import { validateAlias } from '../lib/aliases';
import { renderCustomLinkHtml } from '../lib/custom-link';
import {
  privateAliasNotFoundResponse,
  privateAppUrl,
  privateHtmlResponse,
  privateLoginRequiredResponse,
  privateRedirectResponse,
  publicBadRequestResponse,
  withPrivateNoStoreHeaders,
} from '../lib/responses';
import { resolveTemplateUrl, TemplateResolutionError } from '../lib/templates';
import { parseAliasPath } from './redirect';

type PrivateLinkRow = {
  id: string;
  alias: string;
  link_type: 'redirect' | 'custom';
  target_url: string | null;
  content_markdown: string | null;
  click_count: number;
};

const PRIVATE_LINK_QUERY = `SELECT id, alias, link_type, target_url, content_markdown, click_count
FROM links
WHERE scope='private' AND status='active' AND owner_user_id=? AND alias=?
LIMIT 1`;

const CLICK_COUNT_UPDATE = `UPDATE links SET click_count = click_count + 1, updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') WHERE id=?`;

export function isHeygoPrivateHost(hostname: string, privateHost: string): boolean {
  return hostname === privateHost;
}

export async function handlePrivateShortlink(
  request: Request,
  env: Env,
  ctx: Pick<ExecutionContext, 'waitUntil'>,
): Promise<Response> {
  const url = new URL(request.url);

  // Root: authenticated users go to the private links app; everyone else gets a login prompt.
  if (url.pathname === '/') {
    const user = await getCurrentUser(request, env);
    if (!user) {
      return privateLoginRequiredResponse(env.APP_BASE_URL);
    }
    return privateRedirectResponse(privateAppUrl(env.APP_BASE_URL), 302);
  }

  const pathParts = parseAliasPath(url.pathname);
  if (!pathParts) {
    return privateNotFoundOrLogin(request, env);
  }

  const aliasValidation = validateAlias(pathParts.alias);
  if (!aliasValidation.ok || aliasValidation.value !== pathParts.alias.toLowerCase()) {
    return privateNotFoundOrLogin(request, env);
  }

  const user = await getCurrentUser(request, env);
  if (!user) {
    return privateLoginRequiredResponse(env.APP_BASE_URL);
  }

  const link = await env.DB.prepare(PRIVATE_LINK_QUERY)
    .bind(user.id, aliasValidation.value)
    .first<PrivateLinkRow>();
  if (!link) {
    return privateAliasNotFoundResponse(env.APP_BASE_URL);
  }

  ctx.waitUntil(recordClick(env.DB, link.id));

  if (link.link_type === 'redirect') {
    if (!link.target_url) {
      return privateAliasNotFoundResponse(env.APP_BASE_URL);
    }

    try {
      const location = resolveTemplateUrl({
        targetUrl: link.target_url,
        pathParam: pathParts.pathParam,
        query: url.searchParams,
      });

      return privateRedirectResponse(location, 302);
    } catch (error) {
      if (error instanceof TemplateResolutionError) {
        return withPrivateNoStoreHeaders(publicBadRequestResponse(error.message));
      }

      throw error;
    }
  }

  if (link.link_type === 'custom') {
    return privateHtmlResponse(renderCustomLinkHtml(link), { status: 200 });
  }

  return privateAliasNotFoundResponse(env.APP_BASE_URL);
}

async function privateNotFoundOrLogin(request: Request, env: Env): Promise<Response> {
  const user = await getCurrentUser(request, env);
  if (!user) {
    return privateLoginRequiredResponse(env.APP_BASE_URL);
  }
  return privateAliasNotFoundResponse(env.APP_BASE_URL);
}

async function recordClick(db: D1Database, linkId: string): Promise<void> {
  try {
    await db.prepare(CLICK_COUNT_UPDATE).bind(linkId).run();
  } catch {
    // Analytics must never block or break a private redirect/render.
  }
}