Bind Better Auth sessions to devices

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
-  FROM better_auth_session
-  WHERE token = ?
+  SELECT
+    session.id,
+    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,