heygo/tests/private-redirect.test.ts

import { describe, expect, it } from 'vitest';
import worker from '../worker/index';
import { AUTH_SESSION_COOKIE_NAME, hashSessionToken } from '../worker/auth';

type LinkRow = {
  id: string;
  alias: string;
  scope: 'public' | 'private';
  status: 'active' | 'archived' | 'deleted';
  link_type: 'redirect' | 'custom';
  target_url: string | null;
  content_markdown: string | null;
  click_count: number;
  owner_user_id?: string;
};

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

type RunCall = {
  sql: string;
  params: unknown[];
};

class FakeD1Database {
  readonly preparedSql: string[] = [];
  readonly runCalls: RunCall[] = [];

  constructor(
    private readonly links: LinkRow[],
    private readonly sessions: SessionRow[],
  ) {}

  prepare(sql: string): FakeD1PreparedStatement {
    this.preparedSql.push(sql);
    return new FakeD1PreparedStatement(this, sql);
  }

  async batch(statements: FakeD1PreparedStatement[]): Promise<D1Result[]> {
    return Promise.all(statements.map((statement) => statement.run()));
  }

  findSession(hash: string): SessionRow | null {
    return this.sessions.find((candidate) => candidate.session_token_hash === hash) ?? null;
  }

  findPrivateActive(ownerUserId: string, alias: string): Omit<LinkRow, 'scope' | 'status' | 'owner_user_id'> | null {
    const row = this.links.find((candidate) => {
      return (
        candidate.scope === 'private' &&
        candidate.status === 'active' &&
        candidate.owner_user_id === ownerUserId &&
        candidate.alias === alias
      );
    });

    if (!row) {
      return null;
    }

    return {
      id: row.id,
      alias: row.alias,
      link_type: row.link_type,
      target_url: row.target_url,
      content_markdown: row.content_markdown,
      click_count: row.click_count,
    };
  }

  findPrivateActiveById(ownerUserId: string, id: string): Omit<LinkRow, 'scope' | 'status' | 'owner_user_id'> | null {
    const row = this.links.find((candidate) => {
      return (
        candidate.scope === 'private' &&
        candidate.status === 'active' &&
        candidate.owner_user_id === ownerUserId &&
        candidate.id === id
      );
    });

    if (!row) {
      return null;
    }

    return {
      id: row.id,
      alias: row.alias,
      link_type: row.link_type,
      target_url: row.target_url,
      content_markdown: row.content_markdown,
      click_count: row.click_count,
    };
  }

  findPublicActiveById(id: string): Omit<LinkRow, 'scope' | 'status'> | null {
    const row = this.links.find((candidate) => {
      return candidate.scope === 'public' && candidate.status === 'active' && candidate.id === id;
    });

    if (!row) {
      return null;
    }

    return {
      id: row.id,
      alias: row.alias,
      link_type: row.link_type,
      target_url: row.target_url,
      content_markdown: row.content_markdown,
      click_count: row.click_count,
    };
  }
}

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

  constructor(
    private readonly db: FakeD1Database,
    private readonly sql: string,
    params: unknown[] = [],
  ) {
    this.params = params;
  }

  bind(...params: unknown[]): FakeD1PreparedStatement {
    return new FakeD1PreparedStatement(this.db, this.sql, params);
  }

  async first<T>(): Promise<T | null> {
    if (this.sql.includes('session_token_hash')) {
      const hash = String(this.params[0]);
      const row = this.db.findSession(hash);
      if (!row) {
        return null;
      }
      return {
        id: row.id,
        email: row.email,
        name: row.name,
        image_url: row.image_url,
        role: row.role,
        expires_at: row.expires_at,
      } as T;
    }

    if (this.sql.includes("scope='private'")) {
      expect(this.sql).toContain("status='active'");
      expect(this.sql).toContain('owner_user_id=?');
      expect(this.sql).toContain('LIMIT 1');

      if (this.sql.includes('alias=?')) {
        const ownerUserId = String(this.params[0]);
        const alias = String(this.params[1]);
        return this.db.findPrivateActive(ownerUserId, alias) as T;
      }

      expect(this.sql).toContain('id=?');
      const id = String(this.params[0]);
      const ownerUserId = String(this.params[1]);
      return this.db.findPrivateActiveById(ownerUserId, id) as T;
    }

    if (this.sql.includes("scope='public'")) {
      expect(this.sql).toContain("status='active'");
      expect(this.sql).toContain('id=?');
      const id = String(this.params[0]);
      return this.db.findPublicActiveById(id) as T;
    }

    return null;
  }

  async run(): Promise<D1Result> {
    this.db.runCalls.push({ sql: this.sql, params: this.params });
    return { success: true, meta: {} } as D1Result;
  }
}

class FakeExecutionContext {
  readonly promises: Promise<unknown>[] = [];

  waitUntil(promise: Promise<unknown>): void {
    this.promises.push(promise);
  }

  passThroughOnException(): void {}
}

function futureIso(): string {
  return new Date(Date.now() + 60 * 60 * 1000).toISOString();
}

type EnvBundle = {
  env: {
    DB: D1Database;
    PUBLIC_HOST: string;
    PRIVATE_HOST: string;
    APP_BASE_URL: string;
    COOKIE_DOMAIN: string;
  };
  db: FakeD1Database;
  ctx: FakeExecutionContext;
};

function makeEnv(links: LinkRow[] = [], sessions: SessionRow[] = []): EnvBundle {
  const db = new FakeD1Database(links, sessions);
  return {
    env: {
      DB: db as unknown as D1Database,
      PUBLIC_HOST: 'heygo.cc',
      PRIVATE_HOST: 'my.heygo.cc',
      APP_BASE_URL: 'https://heygo.cc',
      COOKIE_DOMAIN: '.heygo.cc',
    },
    db,
    ctx: new FakeExecutionContext(),
  };
}

async function sessionCookie(token: string): Promise<string> {
  return `${AUTH_SESSION_COOKIE_NAME}=${token}`;
}

async function fetchWorker(
  url: string,
  opts: { links?: LinkRow[]; sessions?: SessionRow[]; cookie?: string } = {},
) {
  const { env, db, ctx } = makeEnv(opts.links ?? [], opts.sessions ?? []);
  const headers = new Headers();
  if (opts.cookie) {
    headers.set('cookie', opts.cookie);
  }
  const response = await worker.fetch(
    new Request(url, { headers }) as unknown as Parameters<typeof worker.fetch>[0],
    env as unknown as Parameters<typeof worker.fetch>[1],
    ctx as unknown as Parameters<typeof worker.fetch>[2],
  );
  return { response, db, ctx };
}

async function userSession(token: string, userId: string): Promise<SessionRow> {
  const hash = await hashSessionToken(token);
  return {
    id: userId,
    email: `${userId}@heygo.cc`,
    name: userId,
    image_url: null,
    role: 'user',
    expires_at: futureIso(),
    session_token_hash: hash,
  };
}

function expectPrivateNoStoreHeaders(response: Response): void {
  expect(response.headers.get('cache-control')).toBe('no-store');
  expect(
    response.headers
      .get('vary')
      ?.split(',')
      .map((value) => value.trim().toLowerCase()),
  ).toContain('cookie');
}

function expectNoPrivateNoStoreHeaders(response: Response): void {
  expect(response.headers.get('cache-control')).not.toBe('no-store');
  expect(
    response.headers
      .get('vary')
      ?.split(',')
      .map((value) => value.trim().toLowerCase()) ?? [],
  ).not.toContain('cookie');
}

describe('my.heygo.cc private shortlinks', () => {
  it('returns 404 with a login link when unauthenticated', async () => {
    const { response, db } = await fetchWorker('https://my.heygo.cc/foo', {
      links: [
        {
          id: 'link_1',
          alias: 'foo',
          scope: 'private',
          status: 'active',
          link_type: 'redirect',
          target_url: 'https://example.com/foo',
          content_markdown: null,
          click_count: 0,
          owner_user_id: 'userA',
        },
      ],
    });

    expect(response.status).toBe(404);
    expect(response.headers.get('content-type')).toContain('text/html');
    expectPrivateNoStoreHeaders(response);
    const body = await response.text();
    expect(body).toContain('Sign in to access');
    expect(body).toContain('https://heygo.cc/app/login');
    // No private alias lookup should happen without a session.
    expect(db.preparedSql.filter((sql) => sql.includes("scope='private'"))).toHaveLength(0);
  });

  it('resolves the authenticated user private redirect', async () => {
    const cookie = await sessionCookie('token-a');
    const session = await userSession('token-a', 'userA');
    const { response, ctx, db } = await fetchWorker('https://my.heygo.cc/foo', {
      cookie,
      sessions: [session],
      links: [
        {
          id: 'link_a',
          alias: 'foo',
          scope: 'private',
          status: 'active',
          link_type: 'redirect',
          target_url: 'https://example.com/foo-a',
          content_markdown: null,
          click_count: 0,
          owner_user_id: 'userA',
        },
      ],
    });

    expect(response.status).toBe(302);
    expect(response.headers.get('location')).toBe('https://example.com/foo-a');
    expectPrivateNoStoreHeaders(response);
    expect(ctx.promises).toHaveLength(1);
    await Promise.all(ctx.promises);
    expect(db.runCalls).toHaveLength(2);
    expect(db.runCalls[0].sql).toContain('click_count = click_count + 1');
    expect(db.runCalls[0].params).toEqual(['link_a']);
    expect(db.runCalls[1].sql).toContain('INSERT INTO click_daily');
    expect(db.runCalls[1].params).toEqual(['link_a']);
  });

  it('resolves /links/:id/go for the authenticated owner from the app host', async () => {
    const cookie = await sessionCookie('token-a');
    const session = await userSession('token-a', 'userA');
    const { response, ctx, db } = await fetchWorker('https://heygo.cc/links/link_a/go', {
      cookie,
      sessions: [session],
      links: [
        {
          id: 'link_a',
          alias: 'foo',
          scope: 'private',
          status: 'active',
          link_type: 'redirect',
          target_url: 'https://example.com/foo-a',
          content_markdown: null,
          click_count: 0,
          owner_user_id: 'userA',
        },
      ],
    });

    expect(response.status).toBe(302);
    expect(response.headers.get('location')).toBe('https://example.com/foo-a');
    expectPrivateNoStoreHeaders(response);
    expect(ctx.promises).toHaveLength(1);
    await Promise.all(ctx.promises);
    expect(db.runCalls[0].params).toEqual(['link_a']);
    expect(db.runCalls[1].params).toEqual(['link_a']);
  });

  it('isolates private aliases per user (user B resolves own link, not user A)', async () => {
    const cookieA = await sessionCookie('token-a');
    const cookieB = await sessionCookie('token-b');
    const sessionA = await userSession('token-a', 'userA');
    const sessionB = await userSession('token-b', 'userB');

    const links: LinkRow[] = [
      {
        id: 'link_a',
        alias: 'shared',
        scope: 'private',
        status: 'active',
        link_type: 'redirect',
        target_url: 'https://example.com/a-shared',
        content_markdown: null,
        click_count: 0,
        owner_user_id: 'userA',
      },
      {
        id: 'link_b',
        alias: 'shared',
        scope: 'private',
        status: 'active',
        link_type: 'redirect',
        target_url: 'https://example.com/b-shared',
        content_markdown: null,
        click_count: 0,
        owner_user_id: 'userB',
      },
    ];

    const resA = await fetchWorker('https://my.heygo.cc/shared', {
      cookie: cookieA,
      sessions: [sessionA, sessionB],
      links,
    });
    expect(resA.response.status).toBe(302);
    expect(resA.response.headers.get('location')).toBe('https://example.com/a-shared');

    const resB = await fetchWorker('https://my.heygo.cc/shared', {
      cookie: cookieB,
      sessions: [sessionA, sessionB],
      links,
    });
    expect(resB.response.status).toBe(302);
    expect(resB.response.headers.get('location')).toBe('https://example.com/b-shared');
  });

  it('returns 404 for an authenticated user who does not own the alias (no public fallback)', async () => {
    const cookie = await sessionCookie('token-b');
    const sessionB = await userSession('token-b', 'userB');
    const { response, db } = await fetchWorker('https://my.heygo.cc/only-public', {
      cookie,
      sessions: [sessionB],
      links: [
        {
          id: 'public_link',
          alias: 'only-public',
          scope: 'public',
          status: 'active',
          link_type: 'redirect',
          target_url: 'https://example.com/public',
          content_markdown: null,
          click_count: 0,
        },
        {
          id: 'link_a',
          alias: 'only-public',
          scope: 'private',
          status: 'active',
          link_type: 'redirect',
          target_url: 'https://example.com/a-only',
          content_markdown: null,
          click_count: 0,
          owner_user_id: 'userA',
        },
      ],
    });

    expect(response.status).toBe(404);
    expectPrivateNoStoreHeaders(response);
    const body = await response.text();
    expect(body).toContain('Create this link');
    // Must never query the public scope on my.heygo.cc.
    expect(db.preparedSql.filter((sql) => sql.includes("scope='public'"))).toHaveLength(0);
  });

  it('renders an authenticated private custom link as escaped HTML', async () => {
    const cookie = await sessionCookie('token-a');
    const session = await userSession('token-a', 'userA');
    const { response } = await fetchWorker('https://my.heygo.cc/note', {
      cookie,
      sessions: [session],
      links: [
        {
          id: 'link_custom',
          alias: 'note',
          scope: 'private',
          status: 'active',
          link_type: 'custom',
          target_url: null,
          content_markdown: '# Private Note\n<script>alert("x")</script>',
          click_count: 0,
          owner_user_id: 'userA',
        },
      ],
    });

    expect(response.status).toBe(200);
    expect(response.headers.get('content-type')).toContain('text/html');
    expectPrivateNoStoreHeaders(response);
    const html = await response.text();
    expect(html).toContain('Private Note');
    expect(html).not.toContain('<script>');
    expect(html).toContain('&lt;script&gt;');
  });

  it('returns 404 with Create this private link when the alias is missing for the user', async () => {
    const cookie = await sessionCookie('token-a');
    const session = await userSession('token-a', 'userA');
    const { response } = await fetchWorker('https://my.heygo.cc/does-not-exist', {
      cookie,
      sessions: [session],
    });

    expect(response.status).toBe(404);
    expect(response.headers.get('content-type')).toContain('text/html');
    expectPrivateNoStoreHeaders(response);
    const body = await response.text();
    expect(body).toContain('Create this link');
  });

  it('handles the root path: unauthenticated 404 + login, authenticated 302 to app/private', async () => {
    const unauth = await fetchWorker('https://my.heygo.cc/');
    expect(unauth.response.status).toBe(404);
    expectPrivateNoStoreHeaders(unauth.response);
    expect(await unauth.response.text()).toContain('Sign in to access');

    const cookie = await sessionCookie('token-a');
    const session = await userSession('token-a', 'userA');
    const auth = await fetchWorker('https://my.heygo.cc/', {
      cookie,
      sessions: [session],
    });
    expect(auth.response.status).toBe(302);
    expect(auth.response.headers.get('location')).toBe('https://heygo.cc/app/private');
    expectPrivateNoStoreHeaders(auth.response);
  });

  it('does not trigger D1 private alias lookup for reserved paths', async () => {
    const cookie = await sessionCookie('token-a');
    const session = await userSession('token-a', 'userA');
    const links: LinkRow[] = [
      {
        id: 'link_api',
        alias: 'api',
        scope: 'private',
        status: 'active',
        link_type: 'redirect',
        target_url: 'https://example.com/api',
        content_markdown: null,
        click_count: 0,
        owner_user_id: 'userA',
      },
      {
        id: 'link_app',
        alias: 'app',
        scope: 'private',
        status: 'active',
        link_type: 'redirect',
        target_url: 'https://example.com/app',
        content_markdown: null,
        click_count: 0,
        owner_user_id: 'userA',
      },
      {
        id: 'link_admin',
        alias: 'admin',
        scope: 'private',
        status: 'active',
        link_type: 'redirect',
        target_url: 'https://example.com/admin',
        content_markdown: null,
        click_count: 0,
        owner_user_id: 'userA',
      },
    ];

    for (const path of ['/api/foo', '/app/foo', '/admin/foo']) {
      const { db } = await fetchWorker(`https://my.heygo.cc${path}`, {
        cookie,
        sessions: [session],
        links,
      });
      expect(db.preparedSql.filter((sql) => sql.includes("scope='private'"))).toHaveLength(0);
    }
  });

  it('marks private-host API and reserved index-level responses no-store', async () => {
    const cases = [
      { path: '/api/health', status: 200 },
      { path: '/api/foo', status: 404 },
      { path: '/app/foo', status: 404 },
      { path: '/admin/foo', status: 404 },
    ];

    for (const { path, status } of cases) {
      const { response, db } = await fetchWorker(`https://my.heygo.cc${path}`);

      expect(response.status).toBe(status);
      expectPrivateNoStoreHeaders(response);
      expect(db.preparedSql.filter((sql) => sql.includes("scope='private'"))).toHaveLength(0);
    }
  });

  it('keeps public-host API and reserved index-level response cache headers unchanged', async () => {
    const cases = [
      { path: '/api/health', status: 200 },
      { path: '/api/foo', status: 404 },
      { path: '/app/foo', status: 404 },
      { path: '/admin/foo', status: 404 },
    ];

    for (const { path, status } of cases) {
      const { response, db } = await fetchWorker(`https://heygo.cc${path}`);

      expect(response.status).toBe(status);
      expectNoPrivateNoStoreHeaders(response);
      expect(db.preparedSql).toHaveLength(0);
    }
  });
});