fix(auth): prevent device registration replay

cloudflare/src/device_schema.ts

@@ -99,6 +99,13 @@ export class DevicePermissionError extends Error {
   }
 }
 
+export class DeviceConflictError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "DeviceConflictError";
+  }
+}
+
 export class DevicePersistenceError extends Error {
   constructor(message: string) {
     super(message);

cloudflare/src/devices.ts

@@ -6,10 +6,12 @@ import {
   type DeviceApprovalRow,
   type DeviceListDocument,
   type DeviceRegistrationDocument,
+  type DeviceRegistrationRequest,
   type DeviceRevocationDocument,
   type DeviceRevocationRequest,
   type DeviceRevocationRow,
   type DeviceRow,
+  DeviceConflictError,
   DevicePermissionError,
   DevicePersistenceError,
   approvedDeviceDocument,
@@ -24,6 +26,7 @@ import {
 } from "./device_schema.js";
 
 export {
+  DeviceConflictError,
   DevicePermissionError,
   DevicePersistenceError,
   DeviceSchemaError,
@@ -61,7 +64,7 @@ const DEVICE_REGISTER_QUERY = `
     revoked_at,
     idempotency_key
   ) VALUES (?, ?, ?, ?, ?, 'pending', ?, NULL, ?, NULL, ?)
-  ON CONFLICT(user_id, idempotency_key) DO NOTHING
+  ON CONFLICT DO NOTHING
 `;
 const DEVICE_BY_IDEMPOTENCY_KEY_QUERY = `
   SELECT
@@ -208,7 +211,7 @@ export async function registerDeviceDocument(
     throw new DevicePermissionError("device_context_mismatch");
   }
 
-  await env.ELY_DB.prepare(DEVICE_REGISTER_QUERY)
+  const writeResult = await env.ELY_DB.prepare(DEVICE_REGISTER_QUERY)
     .bind(
       context.userId,
       registration.deviceId,
@@ -220,21 +223,76 @@ export async function registerDeviceDocument(
       registration.idempotencyKey,
     )
     .run();
+  const insertedRows = changedRowCount(writeResult);
+  if (insertedRows > 1) {
+    throw new DevicePersistenceError("device_registration_write_count_invalid");
+  }
   const row = await env.ELY_DB.prepare(DEVICE_BY_IDEMPOTENCY_KEY_QUERY)
     .bind(context.userId, registration.idempotencyKey)
     .first<DeviceRow>();
   if (row === null) {
+    if (insertedRows === 0) {
+      throw new DeviceConflictError("device_registration_conflict");
+    }
     throw new DevicePersistenceError("device_registration_missing");
   }
-  await bindSessionDeviceContext(env, context, registration.deviceId, nowSeconds);
+  const device = deviceDocument(row, registration.deviceId);
+  if (!registrationMatches(device, registration)) {
+    throw new DeviceConflictError("device_registration_conflict");
+  }
+  if (insertedRows === 0) {
+    if (
+      context.deviceId === undefined ||
+      context.deviceId !== device.device_id ||
+      device.revoked_at !== null
+    ) {
+      throw new DeviceConflictError("device_registration_conflict");
+    }
+    return {
+      version: 1,
+      user_id: context.userId,
+      device,
+    };
+  }
+  if (device.approval_status !== "pending" || device.revoked_at !== null) {
+    throw new DevicePersistenceError("device_registration_state_invalid");
+  }
+  await bindSessionDeviceContext(env, context, device.device_id, nowSeconds);
 
   return {
     version: 1,
     user_id: context.userId,
-    device: deviceDocument(row, registration.deviceId),
+    device,
   };
 }
 
+function registrationMatches(
+  device: DeviceRegistrationDocument["device"],
+  registration: DeviceRegistrationRequest,
+): boolean {
+  return (
+    device.device_id === registration.deviceId &&
+    device.public_key === registration.publicKey &&
+    device.device_name === registration.deviceName &&
+    device.platform === registration.platform
+  );
+}
+
+function changedRowCount(result: unknown): number {
+  if (typeof result !== "object" || result === null || !("meta" in result)) {
+    throw new DevicePersistenceError("device_registration_write_result_invalid");
+  }
+  const meta = result.meta;
+  if (typeof meta !== "object" || meta === null || !("changes" in meta)) {
+    throw new DevicePersistenceError("device_registration_write_result_invalid");
+  }
+  const changes = meta.changes;
+  if (typeof changes !== "number" || !Number.isSafeInteger(changes) || changes < 0) {
+    throw new DevicePersistenceError("device_registration_write_result_invalid");
+  }
+  return changes;
+}
+
 async function bindSessionDeviceContext(
   env: Env,
   context: AuthContext,

cloudflare/src/index.ts

@@ -11,6 +11,7 @@ import {
 } from "./account_deletion.js";
 import { handleBetterAuthRoute } from "./better_auth.js";
 import {
+  DeviceConflictError,
   DevicePermissionError,
   DevicePersistenceError,
   DeviceSchemaError,
@@ -85,6 +86,13 @@ export async function handleRequest(request: Request, env: Env): Promise<Respons
             "Cache-Control": "no-store",
           });
         } catch (error) {
+          if (error instanceof DeviceConflictError) {
+            return jsonResponse(
+              { error: "device_registration_conflict" },
+              409,
+              { "Cache-Control": "no-store" },
+            );
+          }
           if (error instanceof DevicePermissionError) {
             return jsonResponse(
               { error: "device_context_mismatch" },

cloudflare/tests/devices_routes.test.ts

@@ -208,7 +208,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[0]?.includes("ON CONFLICT 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), [
@@ -273,6 +273,113 @@ describe("device routes", () => {
     assert.deepEqual(kvPuts, []);
   });
 
+  it("rejects an unbound session replaying any existing device registration", async () => {
+    for (const status of ["pending", "approved", "revoked"] as const) {
+      const existingDevice = {
+        device_id: "device-01",
+        public_key: PUBLIC_KEY,
+        device_name: "MacBook Pro",
+        platform: "macOS",
+        approval_status: status,
+        created_at: 1_780_000_000,
+        approved_at: status === "pending" ? null : 1_780_000_010,
+        last_active_at: 1_780_000_020,
+        revoked_at: status === "revoked" ? 1_780_000_030 : null,
+      };
+      const d1 = testD1Database({
+        firstRows: [existingDevice],
+        runChanges: [0],
+        sessionRow: {
+          id: "session-02",
+          userId: "user-01",
+          expiresAt: "2099-01-01T00:00:00.000Z",
+          deviceId: null,
+        },
+      });
+
+      const response = await handleRequest(
+        new Request("https://elydora.test/api/devices/register", {
+          method: "POST",
+          headers: {
+            authorization: `Bearer ${ACCESS_TOKEN}`,
+            "content-type": "application/json",
+          },
+          body: JSON.stringify(deviceRegistrationBody()),
+        }),
+        testEnv({ d1 }),
+      );
+
+      assert.equal(response.status, 409, status);
+      assert.deepEqual(await response.json(), { error: "device_registration_conflict" });
+      assert.equal(d1.queries.some((query) => query.includes("session_device_context")), false);
+    }
+  });
+
+  it("allows an exact idempotent retry from the already bound session", async () => {
+    const pendingDevice = {
+      device_id: "device-01",
+      public_key: PUBLIC_KEY,
+      device_name: "MacBook Pro",
+      platform: "macOS",
+      approval_status: "pending",
+      created_at: 1_780_000_000,
+      approved_at: null,
+      last_active_at: 1_780_000_020,
+      revoked_at: null,
+    };
+    const d1 = testD1Database({ firstRows: [pendingDevice], runChanges: [0] });
+
+    const response = await handleRequest(
+      new Request("https://elydora.test/api/devices/register", {
+        method: "POST",
+        headers: {
+          authorization: `Bearer ${ACCESS_TOKEN}`,
+          "content-type": "application/json",
+        },
+        body: JSON.stringify(deviceRegistrationBody()),
+      }),
+      testEnv({ d1 }),
+    );
+
+    assert.equal(response.status, 201);
+    assert.equal(
+      ((await response.json()) as { device: { device_id: string } }).device.device_id,
+      "device-01",
+    );
+    assert.equal(d1.queries.some((query) => query.includes("session_device_context")), false);
+  });
+
+  it("rejects a device id collision with a different idempotency key", async () => {
+    const d1 = testD1Database({
+      firstRows: [],
+      runChanges: [0],
+      sessionRow: {
+        id: "session-02",
+        userId: "user-01",
+        expiresAt: "2099-01-01T00:00:00.000Z",
+        deviceId: null,
+      },
+    });
+    const response = await handleRequest(
+      new Request("https://elydora.test/api/devices/register", {
+        method: "POST",
+        headers: {
+          authorization: `Bearer ${ACCESS_TOKEN}`,
+          "content-type": "application/json",
+        },
+        body: JSON.stringify({
+          ...deviceRegistrationBody(),
+          idempotency_key: "device-register-0002",
+        }),
+      }),
+      testEnv({ d1 }),
+    );
+
+    assert.equal(response.status, 409);
+    assert.deepEqual(await response.json(), { error: "device_registration_conflict" });
+    assert.equal(d1.queries.some((query) => query.includes("session_device_context")), false);
+  });
+
   it("rejects invalid device registration payloads before D1 writes", async () => {
     const tokenHash = await authTokenHash(ACCESS_TOKEN);
     const d1 = testD1Database([]);

cloudflare/tests/devices_test_support.ts

@@ -40,6 +40,7 @@ export interface RecordedR2Put {
 interface TestD1DatabaseOptions {
   allRows?: unknown[];
   firstRows?: unknown[];
+  runChanges?: number[];
   sessionRow?: unknown | null;
 }
 
@@ -114,6 +115,7 @@ export function testD1Database(rows: unknown[] | TestD1DatabaseOptions): Recorde
       ? rows.sessionRow ?? null
       : DEFAULT_AUTH_SESSION_ROW;
   let firstIndex = 0;
+  let runIndex = 0;
   return {
     authBinds,
     authQueries,
@@ -130,6 +132,7 @@ export function testD1Database(rows: unknown[] | TestD1DatabaseOptions): Recorde
         isAuthSessionQuery ? authBinds : binds,
         isAuthSessionQuery,
         sessionRow,
+        () => (!Array.isArray(rows) ? rows.runChanges?.[runIndex++] : undefined) ?? 1,
       );
     },
     batch(statements: ElyD1PreparedStatement[]) {
@@ -159,6 +162,7 @@ function testD1PreparedStatement(
   binds: unknown[][],
   isAuthSessionQuery: boolean,
   sessionRow: unknown | null,
+  nextRunChanges: () => number,
 ): ElyD1PreparedStatement {
   return {
     bind(...values: unknown[]) {
@@ -175,7 +179,7 @@ function testD1PreparedStatement(
       return Promise.resolve({ results: allRows as T[] });
     },
     run() {
-      return Promise.resolve({});
+      return Promise.resolve({ results: [], meta: { changes: nextRunChanges() } });
     },
   };
 }