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 { 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 | 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 | 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 | 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(): Promise { 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 { this.db.runCalls.push({ sql: this.sql, params: this.params }); return { success: true, meta: {} } as D1Result; } } class FakeExecutionContext { readonly promises: Promise[] = []; waitUntil(promise: Promise): 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 { 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[0], env as unknown as Parameters[1], ctx as unknown as Parameters[2], ); return { response, db, ctx }; } async function userSession(token: string, userId: string): Promise { 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', 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('