1
import { env, exports as workerExports } from "cloudflare:workers";
2
import { afterEach, beforeAll, describe, expect, it, vi } from "vitest";
3
 
4
import {
5
  D1_BOOKMARK_COOKIE_MAX_AGE_SECONDS,
6
  D1_BOOKMARK_COOKIE_HEADER_NAME,
7
  D1_BOOKMARK_HEADER,
8
  parseBookmarkCookie,
9
} from "./util/d1Bookmark";
10
import { ensureD1Migrations } from "./util/d1Setup";
11
import { setupRepoForTests } from "./util/repoSeed";
12
 
13
// The D1 Sessions middleware decision is observed through a spy that
14
// wraps `env.DB.withSession`. The wrapper records the chosen anchor and
15
// optionally substitutes a deterministic outbound bookmark so emit-side
16
// assertions don't depend on whether workerd's local D1 advances a real
17
// bookmark for read-only queries. `D1SessionBookmark` and
18
// `D1SessionConstraint` are global ambient types from
19
// `worker-configuration.d.ts`. `undefined` is included to match the
20
// original signature shape even though the middleware always passes a
21
// concrete value.
22
type D1SessionAnchor = D1SessionBookmark | D1SessionConstraint | undefined;
23
 
24
type InstrumentOptions = {
25
  // When set, the wrapped `getBookmark()` returns this value verbatim
26
  // instead of delegating to the real session. `null` is meaningful (no
27
  // emit). `undefined` (default) delegates to the real session.
28
  fakeOutboundBookmark?: string | null;
29
};
30
 
31
type Instrumentation = {
32
  readonly anchors: ReadonlyArray<D1SessionAnchor>;
33
  readonly sessionPrepareCalls: number;
34
  readonly getBookmarkCalls: number;
35
  setFakeBookmark(bookmark: string | null | undefined): void;
36
  reset(): void;
37
  restore(): void;
38
};
39
 
40
function instrumentSessions(options: InstrumentOptions = {}): Instrumentation {
41
  const anchors: D1SessionAnchor[] = [];
42
  let sessionPrepareCalls = 0;
43
  let getBookmarkCalls = 0;
44
  let fakeBookmark: string | null | undefined = options.fakeOutboundBookmark;
45
  const original = env.DB.withSession.bind(env.DB);
46
  const spy = vi.spyOn(env.DB, "withSession").mockImplementation((anchor) => {
47
    anchors.push(anchor);
48
    const session = original(anchor);
49
    const wrapped: D1DatabaseSession = {
50
      prepare(query) {
51
        sessionPrepareCalls += 1;
52
        return session.prepare(query);
53
      },
54
      batch<T = unknown>(statements: D1PreparedStatement[]) {
55
        return session.batch<T>(statements);
56
      },
57
      getBookmark() {
58
        getBookmarkCalls += 1;
59
        if (fakeBookmark !== undefined) return fakeBookmark;
60
        return session.getBookmark();
61
      },
62
    };
63
    return wrapped;
64
  });
65
  return {
66
    get anchors() {
67
      return anchors;
68
    },
69
    get sessionPrepareCalls() {
70
      return sessionPrepareCalls;
71
    },
72
    get getBookmarkCalls() {
73
      return getBookmarkCalls;
74
    },
75
    setFakeBookmark(bookmark) {
76
      fakeBookmark = bookmark;
77
    },
78
    reset() {
79
      anchors.length = 0;
80
      sessionPrepareCalls = 0;
81
      getBookmarkCalls = 0;
82
    },
83
    restore() {
84
      spy.mockRestore();
85
    },
86
  };
87
}
88
 
89
beforeAll(async () => {
90
  await ensureD1Migrations(env);
91
});
92
 
93
let activeInstrumentation: Instrumentation | null = null;
94
 
95
afterEach(() => {
96
  activeInstrumentation?.restore();
97
  activeInstrumentation = null;
98
});
99
 
100
function instrument(options: InstrumentOptions = {}): Instrumentation {
101
  const handle = instrumentSessions(options);
102
  activeInstrumentation = handle;
103
  return handle;
104
}
105
 
106
const ALWAYS_PRIMARY = "first-primary";
107
const ALWAYS_REPLICA = "first-unconstrained";
108
 
109
describe("D1 Sessions middleware", () => {
110
  it("opens exactly one session per request and routes c.var.db through it", async () => {
111
    const owner = `sess-owner-${Math.random().toString(36).slice(2, 8)}`;
112
    await setupRepoForTests(env, owner, "repo");
113
    const probe = instrument();
114
 
115
    const res = await workerExports.default.fetch(`https://example.com/${owner}`);
116
    expect(res.status).toBe(200);
117
    expect(probe.anchors).toHaveLength(1);
118
    // The owner-overview handler runs `findNamespaceBySlug` and at least
119
    // one repository query through `c.var.db`, so the wrapped session
120
    // sees prepares. If the middleware had used `env.DB` directly, the
121
    // counter would still be zero.
122
    expect(probe.sessionPrepareCalls).toBeGreaterThan(0);
123
    expect(probe.getBookmarkCalls).toBe(1);
124
  });
125
 
126
  it("anchors anonymous GET on first-unconstrained when no bookmark is present", async () => {
127
    const owner = `sess-anon-${Math.random().toString(36).slice(2, 8)}`;
128
    await setupRepoForTests(env, owner, "repo");
129
    const probe = instrument();
130
 
131
    await workerExports.default.fetch(`https://example.com/${owner}`);
132
    expect(probe.anchors[0]).toBe(ALWAYS_REPLICA);
133
  });
134
 
135
  it("anchors generic POST on first-primary", async () => {
136
    const probe = instrument();
137
    // No session cookie; the route returns 401 inside the handler. The
138
    // middleware decision is captured before the handler runs.
139
    const res = await workerExports.default.fetch("https://example.com/auth/api/repositories", {
140
      method: "POST",
141
      headers: { "Content-Type": "application/json", Origin: "https://example.com" },
142
      body: "{}",
143
    });
144
    expect(res.status).toBe(401);
145
    expect(probe.anchors[0]).toBe(ALWAYS_PRIMARY);
146
  });
147
 
148
  it("anchors POST /:owner/:repo/git-upload-pack on first-unconstrained (read-shaped POST)", async () => {
149
    const owner = `sess-up-${Math.random().toString(36).slice(2, 8)}`;
150
    const repo = "repo";
151
    await setupRepoForTests(env, owner, repo);
152
    const probe = instrument();
153
 
154
    const res = await workerExports.default.fetch(
155
      `https://example.com/${owner}/${repo}/git-upload-pack`,
156
      {
157
        method: "POST",
158
        headers: {
159
          "Content-Type": "application/x-git-upload-pack-request",
160
          "Git-Protocol": "version=2",
161
        },
162
        body: new Uint8Array([]),
163
      }
164
    );
165
    // 400 because the body has no command; we only care that the
166
    // middleware classified this route correctly.
167
    expect(res.status).toBe(400);
168
    expect(probe.anchors[0]).toBe(ALWAYS_REPLICA);
169
  });
170
 
171
  it("anchors POST /:owner/:repo/git-receive-pack on first-primary (write-shaped POST)", async () => {
172
    const owner = `sess-rcv-${Math.random().toString(36).slice(2, 8)}`;
173
    const repo = "repo";
174
    await setupRepoForTests(env, owner, repo);
175
    const probe = instrument();
176
 
177
    const res = await workerExports.default.fetch(
178
      `https://example.com/${owner}/${repo}/git-receive-pack`,
179
      {
180
        method: "POST",
181
        headers: { "Content-Type": "application/x-git-receive-pack-request" },
182
        body: new Uint8Array([]),
183
      }
184
    );
185
    // Without push credentials the route 401-challenges; that still
186
    // means the middleware classified the path first.
187
    expect(res.status).toBe(401);
188
    expect(probe.anchors[0]).toBe(ALWAYS_PRIMARY);
189
  });
190
 
191
  it("anchors GET /auth/callback on first-primary (write-shaped GET)", async () => {
192
    const probe = instrument();
193
    const res = await workerExports.default.fetch(
194
      "https://example.com/auth/callback?code=&state=",
195
      { redirect: "manual" }
196
    );
197
    // Missing transaction cookie -> redirect to /auth with error. The
198
    // middleware decision still fired.
199
    expect(res.status).toBe(302);
200
    expect(probe.anchors[0]).toBe(ALWAYS_PRIMARY);
201
  });
202
 
203
  it("anchors GET /auth/callback/ (trailing slash) on first-primary", async () => {
204
    // Hono's `strict: false` normalizes the trailing slash before our
205
    // middleware sees `c.req.path`. The defensive normalize in the
206
    // selector backs that up so the classification still holds if the
207
    // option is ever flipped.
208
    const probe = instrument();
209
    const res = await workerExports.default.fetch(
210
      "https://example.com/auth/callback/?code=&state=",
211
      { redirect: "manual" }
212
    );
213
    expect(res.status).toBe(302);
214
    expect(probe.anchors[0]).toBe(ALWAYS_PRIMARY);
215
  });
216
 
217
  it("honors inbound header bookmark on a read-like request", async () => {
218
    const owner = `sess-hdr-${Math.random().toString(36).slice(2, 8)}`;
219
    await setupRepoForTests(env, owner, "repo");
220
    const probe = instrument();
221
 
222
    await workerExports.default.fetch(`https://example.com/${owner}`, {
223
      headers: { [D1_BOOKMARK_HEADER]: "probe-bookmark-aaa" },
224
    });
225
    expect(probe.anchors[0]).toBe("probe-bookmark-aaa");
226
  });
227
 
228
  it("honors inbound cookie bookmark on a read-like request", async () => {
229
    const owner = `sess-ckie-${Math.random().toString(36).slice(2, 8)}`;
230
    await setupRepoForTests(env, owner, "repo");
231
    const probe = instrument();
232
 
233
    await workerExports.default.fetch(`https://example.com/${owner}`, {
234
      headers: {
235
        Cookie: `${D1_BOOKMARK_COOKIE_HEADER_NAME}=${encodeURIComponent("probe-bookmark-bbb")}`,
236
      },
237
    });
238
    expect(probe.anchors[0]).toBe("probe-bookmark-bbb");
239
  });
240
 
241
  it("prefers the inbound header over the inbound cookie", async () => {
242
    const owner = `sess-pref-${Math.random().toString(36).slice(2, 8)}`;
243
    await setupRepoForTests(env, owner, "repo");
244
    const probe = instrument();
245
 
246
    await workerExports.default.fetch(`https://example.com/${owner}`, {
247
      headers: {
248
        [D1_BOOKMARK_HEADER]: "header-wins",
249
        Cookie: `${D1_BOOKMARK_COOKIE_HEADER_NAME}=${encodeURIComponent("cookie-loses")}`,
250
      },
251
    });
252
    expect(probe.anchors[0]).toBe("header-wins");
253
  });
254
 
255
  it("ignores inbound bookmark on a mutating request", async () => {
256
    const probe = instrument();
257
    const res = await workerExports.default.fetch("https://example.com/auth/api/repositories", {
258
      method: "POST",
259
      headers: {
260
        "Content-Type": "application/json",
261
        Origin: "https://example.com",
262
        [D1_BOOKMARK_HEADER]: "should-be-ignored",
263
        Cookie: `${D1_BOOKMARK_COOKIE_HEADER_NAME}=${encodeURIComponent("also-ignored")}`,
264
      },
265
      body: "{}",
266
    });
267
    expect(res.status).toBe(401);
268
    expect(probe.anchors[0]).toBe(ALWAYS_PRIMARY);
269
  });
270
 
271
  it("emits header + cookie when the bookmark advances past the inbound value", async () => {
272
    const owner = `sess-emit-${Math.random().toString(36).slice(2, 8)}`;
273
    await setupRepoForTests(env, owner, "repo");
274
    const fakeBookmark = "advanced-bookmark-001";
275
    const probe = instrument({ fakeOutboundBookmark: fakeBookmark });
276
 
277
    const res = await workerExports.default.fetch(`https://example.com/${owner}`);
278
    expect(probe.anchors[0]).toBe(ALWAYS_REPLICA);
279
    expect(res.headers.get(D1_BOOKMARK_HEADER)).toBe(fakeBookmark);
280
    const cookie = parseBookmarkCookie(res.headers.get("set-cookie"));
281
    expect(cookie?.value).toBe(fakeBookmark);
282
    expect(cookie?.attributes.Path).toBe("/");
283
    expect(cookie?.attributes.HttpOnly).toBe(true);
284
    expect(cookie?.attributes.Secure).toBe(true);
285
    expect(cookie?.attributes.SameSite).toBe("Lax");
286
    expect(cookie?.attributes["Max-Age"]).toBe(String(D1_BOOKMARK_COOKIE_MAX_AGE_SECONDS));
287
  });
288
 
289
  it("does not emit when the route never touches D1", async () => {
290
    // Anonymous `GET /` calls `loadViewer` which short-circuits before
291
    // hitting D1 because there is no session cookie. The session opens
292
    // but no queries run and `getBookmark()` returns null.
293
    const probe = instrument();
294
    const res = await workerExports.default.fetch("https://example.com/");
295
    expect(res.status).toBe(200);
296
    expect(probe.anchors).toHaveLength(1);
297
    expect(probe.sessionPrepareCalls).toBe(0);
298
    expect(res.headers.get(D1_BOOKMARK_HEADER)).toBeNull();
299
    expect(parseBookmarkCookie(res.headers.get("set-cookie"))).toBeNull();
300
  });
301
 
302
  it("does not emit when the outbound bookmark matches the inbound bookmark", async () => {
303
    const matchedBookmark = "no-advance-bookmark";
304
    const probe = instrument({ fakeOutboundBookmark: matchedBookmark });
305
    const res = await workerExports.default.fetch("https://example.com/", {
306
      headers: { [D1_BOOKMARK_HEADER]: matchedBookmark },
307
    });
308
    expect(res.status).toBe(200);
309
    expect(probe.anchors[0]).toBe(matchedBookmark);
310
    expect(res.headers.get(D1_BOOKMARK_HEADER)).toBeNull();
311
    expect(parseBookmarkCookie(res.headers.get("set-cookie"))).toBeNull();
312
  });
313
 
314
  it("rejects malformed inbound bookmark values and falls back to the default anchor", async () => {
315
    const probe = instrument();
316
    // Control character in the header -> rejected; cookie is honored
317
    // because it survives the sanitizer.
318
    await workerExports.default.fetch("https://example.com/", {
319
      headers: {
320
        [D1_BOOKMARK_HEADER]: "bad\x01value",
321
        Cookie: `${D1_BOOKMARK_COOKIE_HEADER_NAME}=${encodeURIComponent("fallback-cookie")}`,
322
      },
323
    });
324
    expect(probe.anchors[0]).toBe("fallback-cookie");
325
  });
326
 
327
  it("rejects oversize inbound bookmark and falls back to first-unconstrained", async () => {
328
    const probe = instrument();
329
    // D1 bookmarks are 59-char Lamport tokens (`8-8-8-32` hex). The cap is
330
    // 256 chars to leave headroom for format evolution; anything past
331
    // that is rejected as transport-malformed.
332
    const oversized = "a".repeat(257);
333
    await workerExports.default.fetch("https://example.com/", {
334
      headers: { [D1_BOOKMARK_HEADER]: oversized },
335
    });
336
    expect(probe.anchors[0]).toBe(ALWAYS_REPLICA);
337
  });
338
 
339
  it("round-trips a bookmark from a write to a follow-up read", async () => {
340
    // The write half: a mutating request returns a bookmark cookie.
341
    // We force a known outbound bookmark so the assertion is decoupled
342
    // from whether the local D1 actually advanced state for the 401
343
    // response that gets returned.
344
    const issuedBookmark = "round-trip-bookmark-xyz";
345
    const writeProbe = instrument({ fakeOutboundBookmark: issuedBookmark });
346
    const writeRes = await workerExports.default.fetch(
347
      "https://example.com/auth/api/repositories",
348
      {
349
        method: "POST",
350
        headers: { "Content-Type": "application/json", Origin: "https://example.com" },
351
        body: "{}",
352
      }
353
    );
354
    expect(writeProbe.anchors[0]).toBe(ALWAYS_PRIMARY);
355
    const cookie = parseBookmarkCookie(writeRes.headers.get("set-cookie"));
356
    expect(cookie?.value).toBe(issuedBookmark);
357
    writeProbe.restore();
358
    activeInstrumentation = null;
359
 
360
    // The read half: the browser would echo the cookie on the next
361
    // navigation. A read-like request consumes it as the session anchor.
362
    const owner = `sess-rt-${Math.random().toString(36).slice(2, 8)}`;
363
    await setupRepoForTests(env, owner, "repo");
364
    const readProbe = instrument();
365
    await workerExports.default.fetch(`https://example.com/${owner}`, {
366
      headers: {
367
        Cookie: `${D1_BOOKMARK_COOKIE_HEADER_NAME}=${encodeURIComponent(issuedBookmark)}`,
368
      },
369
    });
370
    expect(readProbe.anchors[0]).toBe(issuedBookmark);
371
  });
372
});