1
import { createExecutionContext, env } from "cloudflare:test";
2
import { beforeEach, describe, expect, it } from "vitest";
3
import worker from "@/worker/index";
4
import { createWebSocketTicket } from "@/worker/services/inbox";
5
import { resetWorkerState, seedDomain, seedInbox } from "./api/helpers";
6
 
7
function websocketRequest(address: string, ticket?: string, origin = "https://flamemail.devbin.tools") {
8
  const url = new URL("https://flamemail.devbin.tools/ws");
9
  url.searchParams.set("address", address);
10
  if (ticket) {
11
    url.searchParams.set("ticket", ticket);
12
  }
13
 
14
  return new Request(url, {
15
    headers: {
16
      origin,
17
      upgrade: "websocket",
18
    },
19
  }) as Parameters<typeof worker.fetch>[0];
20
}
21
 
22
describe("worker websocket admission", () => {
23
  beforeEach(async () => {
24
    await resetWorkerState();
25
  });
26
 
27
  it("rejects websocket upgrades with an invalid origin", async () => {
28
    const response = await worker.fetch(
29
      websocketRequest("reader@mail.test", "wst_invalid", "https://evil.example"),
30
      env,
31
      createExecutionContext(),
32
    );
33
 
34
    expect(response.status).toBe(403);
35
    await expect(response.text()).resolves.toBe("Forbidden");
36
  });
37
 
38
  it("rejects websocket upgrades with an invalid or missing ticket", async () => {
39
    const response = await worker.fetch(websocketRequest("reader@mail.test"), env, createExecutionContext());
40
 
41
    expect(response.status).toBe(401);
42
    await expect(response.text()).resolves.toBe("Unauthorized");
43
  });
44
 
45
  it("rejects websocket upgrades when the inbox does not exist", async () => {
46
    const ticket = await createWebSocketTicket(env, "reader@mail.test", {
47
      type: "admin",
48
      sub: "00000000-0000-4000-8000-000000000001",
49
    });
50
 
51
    const response = await worker.fetch(websocketRequest("reader@mail.test", ticket), env, createExecutionContext());
52
 
53
    expect(response.status).toBe(404);
54
    await expect(response.text()).resolves.toBe("Inbox not found");
55
  });
56
 
57
  it("rejects websocket upgrades when a user session ticket targets another inbox", async () => {
58
    await seedDomain("mail.test", true);
59
    await seedInbox({
60
      address: "reader@mail.test",
61
    });
62
    const ticket = await createWebSocketTicket(env, "reader@mail.test", {
63
      type: "user",
64
      address: "other@mail.test",
65
    });
66
 
67
    const response = await worker.fetch(websocketRequest("reader@mail.test", ticket), env, createExecutionContext());
68
 
69
    expect(response.status).toBe(403);
70
    await expect(response.text()).resolves.toBe("Forbidden");
71
  });
72
 
73
  it("accepts websocket upgrades for a valid one-time ticket", async () => {
74
    await seedDomain("mail.test", true);
75
    await seedInbox({
76
      address: "reader@mail.test",
77
    });
78
    const ticket = await createWebSocketTicket(env, "reader@mail.test", {
79
      type: "user",
80
      address: "reader@mail.test",
81
    });
82
 
83
    const response = await worker.fetch(websocketRequest("reader@mail.test", ticket), env, createExecutionContext());
84
 
85
    expect(response.status).toBe(101);
86
    expect(await env.SESSIONS.get(`ws-ticket:${ticket}`)).toBeNull();
87
  });
88
 
89
  it("rejects reusing a consumed websocket ticket", async () => {
90
    await seedDomain("mail.test", true);
91
    await seedInbox({
92
      address: "reader@mail.test",
93
    });
94
    const ticket = await createWebSocketTicket(env, "reader@mail.test", {
95
      type: "admin",
96
      sub: "00000000-0000-4000-8000-000000000001",
97
    });
98
 
99
    const firstResponse = await worker.fetch(
100
      websocketRequest("reader@mail.test", ticket),
101
      env,
102
      createExecutionContext(),
103
    );
104
    const secondResponse = await worker.fetch(
105
      websocketRequest("reader@mail.test", ticket),
106
      env,
107
      createExecutionContext(),
108
    );
109
 
110
    expect(firstResponse.status).toBe(101);
111
    expect(secondResponse.status).toBe(401);
112
    await expect(secondResponse.text()).resolves.toBe("Unauthorized");
113
  });
114
});