Bind Better Auth sessions to devices

cloudflare/migrations/0007_better_auth_session_device_context.sql

@@ -0,0 +1,11 @@
+CREATE TABLE IF NOT EXISTS better_auth_session_device_context (
+  session_id TEXT NOT NULL PRIMARY KEY,
+  user_id TEXT NOT NULL,
+  device_id TEXT NOT NULL,
+  updated_at INTEGER NOT NULL,
+  FOREIGN KEY (session_id) REFERENCES better_auth_session (id) ON DELETE CASCADE,
+  FOREIGN KEY (user_id, device_id) REFERENCES user_devices (user_id, device_id) ON DELETE CASCADE
+);
+
+CREATE INDEX IF NOT EXISTS idx_better_auth_session_device_context_user_device
+  ON better_auth_session_device_context (user_id, device_id);

cloudflare/src/auth.ts

@@ -6,9 +6,15 @@ const BEARER_TOKEN_PATTERN = /^[A-Za-z0-9._~+/=-]{32,4096}$/;
 const SUBJECT_ID_PATTERN = /^[a-zA-Z0-9._:-]{3,128}$/;
 const DEVICE_ID_PATTERN = /^[a-zA-Z0-9._:-]{3,128}$/;
 const BETTER_AUTH_SESSION_QUERY = `
-  SELECT id, userId, expiresAt
+  SELECT
-  FROM better_auth_session
+    session.id,
-  WHERE token = ?
+    session.userId,
+    session.expiresAt,
+    device_context.device_id AS deviceId
+  FROM better_auth_session AS session
+  LEFT JOIN better_auth_session_device_context AS device_context
+    ON device_context.session_id = session.id
+  WHERE session.token = ?
 `;
 
 export interface AuthContext {
@@ -23,6 +29,7 @@ interface BetterAuthSessionRow extends Record<string, unknown> {
   id: unknown;
   userId: unknown;
   expiresAt: unknown;
+  deviceId?: unknown;
 }
 
 export type AuthErrorCode =
@@ -115,6 +122,10 @@ async function readBetterAuthSessionContext(
     tokenHash,
     expiresAt: timestampField(row, "expiresAt"),
   };
+  const deviceId = optionalStringField(row, "deviceId");
+  if (deviceId !== undefined) {
+    session.deviceId = deviceIdValue(deviceId);
+  }
   if (Date.parse(session.expiresAt) <= now.getTime()) {
     throw new AuthError("session_expired");
   }

cloudflare/src/devices.ts

@@ -166,6 +166,24 @@ const DEVICE_REVOKE_QUERY = `
     revoked_at = COALESCE(revoked_at, ?)
   WHERE user_id = ? AND device_id = ? AND revoked_at IS NULL
 `;
+const SESSION_DEVICE_CONTEXT_UPSERT_QUERY = `
+  INSERT INTO better_auth_session_device_context (
+    session_id,
+    user_id,
+    device_id,
+    updated_at
+  )
+  SELECT ?, ?, ?, ?
+  WHERE EXISTS (
+    SELECT 1
+    FROM better_auth_session
+    WHERE id = ? AND userId = ?
+  )
+  ON CONFLICT(session_id) DO UPDATE SET
+    user_id = excluded.user_id,
+    device_id = excluded.device_id,
+    updated_at = excluded.updated_at
+`;
 
 export async function deviceListDocument(
   env: Env,
@@ -208,6 +226,7 @@ export async function registerDeviceDocument(
   if (row === null) {
     throw new DevicePersistenceError("device_registration_missing");
   }
+  await bindSessionDeviceContext(env, context, registration.deviceId, nowSeconds);
 
   return {
     version: 1,
@@ -216,6 +235,17 @@ export async function registerDeviceDocument(
   };
 }
 
+async function bindSessionDeviceContext(
+  env: Env,
+  context: AuthContext,
+  deviceId: string,
+  nowSeconds: number,
+): Promise<void> {
+  await env.ELY_DB.prepare(SESSION_DEVICE_CONTEXT_UPSERT_QUERY)
+    .bind(context.sessionId, context.userId, deviceId, nowSeconds, context.sessionId, context.userId)
+    .run();
+}
+
 export async function approveDeviceDocument(
   request: Request,
   env: Env,

cloudflare/tests/api_controls.test.ts

@@ -213,7 +213,12 @@ describe("api controls", () => {
         rateLimitKeys,
         d1: testD1Database({
           firstRows: [
-            { id: "session-01", userId: "user-01", expiresAt: "2099-01-01T00:00:00.000Z" },
+            {
+              id: "session-01",
+              userId: "user-01",
+              expiresAt: "2099-01-01T00:00:00.000Z",
+              deviceId: "device-01",
+            },
           ],
           queries: d1Queries,
           binds: d1Binds,
@@ -223,12 +228,23 @@ describe("api controls", () => {
       ["GET"],
       (context) =>
         Promise.resolve(
-          jsonResponse({ user_id: context.userId, session_id: context.sessionId }, 200),
+          jsonResponse(
+            {
+              user_id: context.userId,
+              session_id: context.sessionId,
+              device_id: context.deviceId,
+            },
+            200,
+          ),
         ),
     );
 
     assert.equal(response.status, 200);
-    assert.deepEqual(await response.json(), { user_id: "user-01", session_id: "session-01" });
+    assert.deepEqual(await response.json(), {
+      user_id: "user-01",
+      session_id: "session-01",
+      device_id: "device-01",
+    });
     assert.deepEqual(rateLimitKeys, [`local:devices.list:bearer:${tokenHash}`]);
     assert.deepEqual(kvReads, [authSessionCacheKvKey("local", tokenHash)]);
     assert.equal(d1Queries.length, 1);
@@ -243,6 +259,7 @@ describe("api controls", () => {
       "",
       "user-01",
     ]);
+    assert.equal(auditEvents[0]?.blobs?.[7], "device-01");
   });
 
   it("rejects expired authenticated sessions", async () => {

cloudflare/tests/devices_routes.test.ts

@@ -207,6 +207,7 @@ describe("device routes", () => {
     assert.ok(d1.queries[0]?.includes("INSERT INTO user_devices"));
     assert.ok(d1.queries[0]?.includes("ON CONFLICT(user_id, idempotency_key) DO NOTHING"));
     assert.ok(d1.queries[1]?.includes("WHERE user_id = ? AND idempotency_key = ?"));
+    assert.ok(d1.queries[2]?.includes("better_auth_session_device_context"));
     assert.deepEqual(d1.binds[0]?.slice(0, 5), [
       "user-01",
       "device-01",
@@ -218,6 +219,7 @@ describe("device routes", () => {
     assert.equal(typeof d1.binds[0]?.[6], "number");
     assert.equal(d1.binds[0]?.[7], IDEMPOTENCY_KEY);
     assert.deepEqual(d1.binds[1], ["user-01", IDEMPOTENCY_KEY]);
+    assert.deepEqual(d1.binds[2]?.slice(0, 3), ["session-01", "user-01", "device-01"]);
   });
 
   it("rejects invalid device registration payloads before D1 writes", async () => {

cloudflare/tests/migrations.test.ts

@@ -13,6 +13,7 @@ const EXPECTED_MIGRATIONS = [
   "0004_releases.sql",
   "0005_audit.sql",
   "0006_better_auth.sql",
+  "0007_better_auth_session_device_context.sql",
 ];
 const USER_SCOPED_TABLES = [
   "user_devices",
@@ -39,6 +40,7 @@ describe("D1 migrations", () => {
         "audit_events",
         "better_auth_account",
         "better_auth_session",
+        "better_auth_session_device_context",
         "better_auth_user",
         "better_auth_verification",
         "device_approvals",
@@ -104,6 +106,15 @@ describe("D1 migrations", () => {
         ]),
         [],
       );
+      assert.deepEqual(
+        requiredColumns(databasePath, "better_auth_session_device_context", [
+          "session_id",
+          "user_id",
+          "device_id",
+          "updated_at",
+        ]),
+        [],
+      );
     });
   });