heygo/tests/redirect.test.ts

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

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;
};

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

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

  constructor(private readonly rows: LinkRow[]) {}

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

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

    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 {
  constructor(
    private readonly db: FakeD1Database,
    private readonly sql: string,
    private readonly params: unknown[] = [],
  ) {}

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

  async first<T>(): Promise<T | null> {
    expect(this.sql).toContain("scope='public'");
    expect(this.sql).toContain("status='active'");
    expect(this.sql).toContain('alias=?');

    const alias = String(this.params[0]);
    return this.db.findPublicActive(alias) as T | 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 makeEnv(rows: LinkRow[] = []) {
  const db = new FakeD1Database(rows);
  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,
  };
}

async function fetchWorker(url: string, rows: LinkRow[] = []) {
  const { env, db } = makeEnv(rows);
  const ctx = new FakeExecutionContext();
  const response = await worker.fetch(
    new Request(url) 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 };
}

describe('public heygo.cc shortlink redirects', () => {
  it('redirects a public redirect link and records analytics asynchronously', async () => {
    const { response, ctx, db } = await fetchWorker('https://heygo.cc/docs', [
      {
        id: 'link_1',
        alias: 'docs',
        scope: 'public',
        status: 'active',
        link_type: 'redirect',
        target_url: 'https://example.com/docs',
        content_markdown: null,
        click_count: 0,
      },
    ]);

    expect(response.status).toBe(302);
    expect(response.headers.get('location')).toBe('https://example.com/docs');
    expect(ctx.promises).toHaveLength(1);

    await Promise.all(ctx.promises);
    expect(db.runCalls).toHaveLength(1);
    expect(db.runCalls[0].sql).toContain('UPDATE links SET click_count = click_count + 1');
    expect(db.runCalls[0].params).toEqual(['link_1']);
  });

  it('renders a public custom link as escaped HTML', async () => {
    const { response } = await fetchWorker('https://heygo.cc/about', [
      {
        id: 'link_2',
        alias: 'about',
        scope: 'public',
        status: 'active',
        link_type: 'custom',
        target_url: null,
        content_markdown: '# About Heygo\n<script>alert("x")</script>',
        click_count: 5,
      },
    ]);

    expect(response.status).toBe(200);
    expect(response.headers.get('content-type')).toContain('text/html');

    const html = await response.text();
    expect(html).toContain('About Heygo');
    expect(html).not.toContain('<script>');
    expect(html).toContain('&lt;script&gt;');
  });

  it('returns a public 404 HTML response for an unknown alias on heygo.cc', async () => {
    const { response } = await fetchWorker('https://heygo.cc/missing');

    expect(response.status).toBe(404);
    expect(response.headers.get('content-type')).toContain('text/html');
    await expect(response.text()).resolves.toContain('Not found');
  });

  it('resolves template parameters from query, pathParam, and defaults', async () => {
    const { response } = await fetchWorker('https://heygo.cc/search/path%20value?lang=fr', [
      {
        id: 'link_3',
        alias: 'search',
        scope: 'public',
        status: 'active',
        link_type: 'redirect',
        target_url: 'https://search.example/{query,default=hello}?lang={lang,default=en}&page={page,default=1}',
        content_markdown: null,
        click_count: 0,
      },
    ]);

    expect(response.status).toBe(302);
    expect(response.headers.get('location')).toBe('https://search.example/path%20value?lang=fr&page=1');
  });

  it('does not return private links through the heygo.cc public route', async () => {
    const { response } = await fetchWorker('https://heygo.cc/secret', [
      {
        id: 'link_4',
        alias: 'secret',
        scope: 'private',
        status: 'active',
        link_type: 'redirect',
        target_url: 'https://example.com/secret',
        content_markdown: null,
        click_count: 0,
      },
    ]);

    expect(response.status).toBe(404);
  });

  it('keeps /api/health working and does not treat it as an alias', async () => {
    const { response, db } = await fetchWorker('https://heygo.cc/api/health', [
      {
        id: 'link_5',
        alias: 'api',
        scope: 'public',
        status: 'active',
        link_type: 'redirect',
        target_url: 'https://example.com/api',
        content_markdown: null,
        click_count: 0,
      },
    ]);

    expect(response.status).toBe(200);
    await expect(response.json()).resolves.toEqual({ ok: true, service: 'heygo-worker' });
    expect(db.preparedSql).toHaveLength(0);
  });

  it('reserves /app/* and /admin/* without D1 alias lookup', async () => {
    for (const path of ['/app/foo', '/admin/foo']) {
      const { response, db } = await fetchWorker(`https://heygo.cc${path}`, [
        {
          id: `link_${path}`,
          alias: path.split('/')[1],
          scope: 'public',
          status: 'active',
          link_type: 'redirect',
          target_url: 'https://example.com/reserved',
          content_markdown: null,
          click_count: 0,
        },
      ]);

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