diff --git a/.env.example b/.env.example index 527e0b2..a879c11 100644 --- a/.env.example +++ b/.env.example @@ -34,6 +34,9 @@ DATABASE_URL= # Redis REDIS_URL= +# Warn a device to replenish when its unconsumed one-time prekeys drop below +# this count (emits `prekeys_low` over the socket). Defaults to 20. +PREKEY_LOW_THRESHOLD= # Rate limits and quotas (see docs/security/rate-limits.md). # Any bucket in apps/backend/src/config/rateLimits.ts can be overridden with # RATE_LIMIT_=[/]. Defaults apply when unset. diff --git a/apps/backend/docs/e2ee-onboarding.md b/apps/backend/docs/e2ee-onboarding.md index 22f3afe..c13ad29 100644 --- a/apps/backend/docs/e2ee-onboarding.md +++ b/apps/backend/docs/e2ee-onboarding.md @@ -535,6 +535,50 @@ Required guarantees for this path: - signed prekey must still be present if the device remains reachable for session bootstrap - client must treat this as lower-entropy/fallback first-contact establishment and should trigger recipient prekey replenishment UX when possible +### C) Low-prekey warning before exhaustion + +Waiting for exhaustion means every sender in the meantime is downgraded to +3-DH, so the backend warns the owning device *before* it runs dry. + +Two surfaces expose this: + +1. `GET /devices` returns `oneTimePreKeysRemaining` per device — a count of + unconsumed one-time prekeys, `0` when the device has none. Poll-free clients + can read this at startup to decide whether to top up. +2. A `prekeys_low` Socket.IO event is emitted after a bundle fetch drops a + device below the threshold (default `20`, overridable with the + `PREKEY_LOW_THRESHOLD` env var). + +Event payload: + +```json +{ + "deviceId": "uuid", + "oneTimePreKeysRemaining": 19, + "threshold": 20 +} +``` + +Delivery and debounce semantics: + +- emitted only to the `device:{deviceId}` room — the owning device, on whichever + gateway holds its socket. No other device on the account sees it, since only + the owner can generate replacement prekeys. +- fired **at most once per threshold crossing**. A device that keeps serving + bundles while below the threshold is told once, not once per fetch. +- the signal re-arms when `POST /devices/:id/prekeys` brings the device back to + or above the threshold, so a later crossing warns again. Revoking a device + also re-arms it (its prekeys are deleted). +- a device that is offline when the threshold is crossed misses the event; it + should read `oneTimePreKeysRemaining` from `GET /devices` on reconnect. + +Client behavior: + +- on `prekeys_low`, generate and upload a fresh batch via + `POST /devices/:id/prekeys`, respecting the `200` cap +- the upload response echoes `oneTimePreKeysRemaining` so the client can confirm + it is back above the threshold without a follow-up `GET /devices` + ## End-to-end ordering contract For compatibility with the current implementation, clients should rely on this ordering: @@ -564,8 +608,10 @@ For compatibility with the current implementation, clients should rely on this o - device registration/listing/revocation/prekey upload: `apps/backend/src/routes/devices.ts` - recipient key-bundle fetch: `apps/backend/src/routes/users.ts` (`GET /users/:userId/devices/:deviceId/key-bundle`) - E2EE-related schema: `apps/backend/src/db/schema.ts` (`devices`, `devicePrekeys`) +- low-prekey signal + debounce latch: `apps/backend/src/services/prekeyLowSignal.ts` - prekey route tests: `apps/backend/src/__tests__/devices.prekeys.test.ts` - key-bundle route tests: `apps/backend/src/__tests__/users.bundle.test.ts` +- low-prekey signal tests: `apps/backend/src/__tests__/prekeysLow.test.ts` ## Gaps to close for full first-DM support diff --git a/apps/backend/src/__tests__/devices.prekeys.test.ts b/apps/backend/src/__tests__/devices.prekeys.test.ts index 2dce5e6..9d61df0 100644 --- a/apps/backend/src/__tests__/devices.prekeys.test.ts +++ b/apps/backend/src/__tests__/devices.prekeys.test.ts @@ -35,6 +35,14 @@ vi.mock('../db/schema.js', () => ({ vi.mock('../lib/redis.js', () => ({ redis: null })); +// Real threshold + count logic, but the latch release is spied so we can assert +// the upload route re-arms the signal only once the device is back in range. +const mockReleaseLatch = vi.fn().mockResolvedValue(undefined); +vi.mock('../services/prekeyLowSignal.js', async (importOriginal) => ({ + ...(await importOriginal()), + releasePrekeysLowLatch: mockReleaseLatch, +})); + vi.mock('drizzle-orm', () => ({ eq: vi.fn((col: unknown, val: unknown) => ({ col, val })), and: vi.fn((...args: unknown[]) => args), @@ -190,6 +198,30 @@ describe('POST /devices/:id/prekeys', () => { expect(mockInsert).toHaveBeenCalledTimes(2); // signed + OTP }); + it('re-arms the low-prekey signal once replenished to the threshold', async () => { + mockDeviceFindFirst.mockResolvedValue(ACTIVE_DEVICE); + setupOtpCount(50); // post-upload recount lands at or above the threshold + setupInsertChain(); + + const res = await request(makeApp()).post('/devices/device-1/prekeys').send(VALID_BODY); + + expect(res.status).toBe(200); + expect(res.body.oneTimePreKeysRemaining).toBe(50); + expect(mockReleaseLatch).toHaveBeenCalledWith('device-1'); + }); + + it('leaves the signal armed when still below the threshold after upload', async () => { + mockDeviceFindFirst.mockResolvedValue(ACTIVE_DEVICE); + setupOtpCount(3); + setupInsertChain(); + + const res = await request(makeApp()).post('/devices/device-1/prekeys').send(VALID_BODY); + + expect(res.status).toBe(200); + expect(res.body.oneTimePreKeysRemaining).toBe(3); + expect(mockReleaseLatch).not.toHaveBeenCalled(); + }); + it('trims the OTP batch to the remaining cap space', async () => { mockDeviceFindFirst.mockResolvedValue(ACTIVE_DEVICE); setupOtpCount(199); // 1 slot left diff --git a/apps/backend/src/__tests__/prekeysLow.test.ts b/apps/backend/src/__tests__/prekeysLow.test.ts new file mode 100644 index 0000000..c0446b7 --- /dev/null +++ b/apps/backend/src/__tests__/prekeysLow.test.ts @@ -0,0 +1,136 @@ +/** + * Tests for the low one-time-prekey signal. + * + * Covers the debounce contract: `prekeys_low` fires at most once per threshold + * crossing, and replenishing back to the threshold re-arms it. + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +// ── Mocks ───────────────────────────────────────────────────────────────────── + +const mockSelect = vi.fn(); +const mockEmit = vi.fn(); +const mockTo = vi.fn(() => ({ emit: mockEmit })); + +vi.mock('../db/index.js', () => ({ db: { select: mockSelect } })); + +vi.mock('../db/schema.js', () => ({ + devicePrekeys: { deviceId: 'deviceId', keyType: 'keyType', consumed: 'consumed' }, +})); + +vi.mock('drizzle-orm', () => ({ + eq: vi.fn((col: unknown, val: unknown) => ({ col, val })), + and: vi.fn((...args: unknown[]) => args), + count: vi.fn(() => 'count(*)'), +})); + +vi.mock('../lib/redis.js', () => ({ redis: null })); +vi.mock('../lib/socket.js', () => ({ getSocketServer: vi.fn(() => ({ to: mockTo })) })); + +const { + PREKEY_LOW_THRESHOLD, + signalPrekeysLowIfNeeded, + releasePrekeysLowLatch, + countAvailableOneTimePreKeys, + __resetPrekeyLowLatches, +} = await import('../services/prekeyLowSignal.js'); + +const DEVICE = 'device-1'; + +beforeEach(() => { + vi.clearAllMocks(); + __resetPrekeyLowLatches(); +}); + +describe('signalPrekeysLowIfNeeded', () => { + it('defaults the threshold to 20', () => { + expect(PREKEY_LOW_THRESHOLD).toBe(20); + }); + + it('emits prekeys_low to the device room when below the threshold', async () => { + await signalPrekeysLowIfNeeded(DEVICE, 19); + + expect(mockTo).toHaveBeenCalledWith(`device:${DEVICE}`); + expect(mockEmit).toHaveBeenCalledTimes(1); + expect(mockEmit).toHaveBeenCalledWith('prekeys_low', { + deviceId: DEVICE, + oneTimePreKeysRemaining: 19, + threshold: 20, + }); + }); + + it('does not emit while at or above the threshold', async () => { + await signalPrekeysLowIfNeeded(DEVICE, 20); + await signalPrekeysLowIfNeeded(DEVICE, 50); + + expect(mockEmit).not.toHaveBeenCalled(); + }); + + it('emits at most once per threshold crossing', async () => { + for (let remaining = 19; remaining >= 0; remaining--) { + await signalPrekeysLowIfNeeded(DEVICE, remaining); + } + + expect(mockEmit).toHaveBeenCalledTimes(1); + expect(mockEmit).toHaveBeenCalledWith( + 'prekeys_low', + expect.objectContaining({ oneTimePreKeysRemaining: 19 }), + ); + }); + + it('latches per device — one device going low does not silence another', async () => { + await signalPrekeysLowIfNeeded('device-a', 5); + await signalPrekeysLowIfNeeded('device-b', 5); + await signalPrekeysLowIfNeeded('device-a', 4); + + expect(mockEmit).toHaveBeenCalledTimes(2); + }); + + it('stops firing after replenishment, then fires again on the next crossing', async () => { + await signalPrekeysLowIfNeeded(DEVICE, 3); + expect(mockEmit).toHaveBeenCalledTimes(1); + + // Replenished above the threshold — this re-arms the latch and is itself + // silent. + await signalPrekeysLowIfNeeded(DEVICE, 100); + expect(mockEmit).toHaveBeenCalledTimes(1); + + // Draining back down below the threshold signals once more. + await signalPrekeysLowIfNeeded(DEVICE, 19); + await signalPrekeysLowIfNeeded(DEVICE, 18); + expect(mockEmit).toHaveBeenCalledTimes(2); + }); + + it('re-arms after an explicit latch release (e.g. prekey upload)', async () => { + await signalPrekeysLowIfNeeded(DEVICE, 2); + await releasePrekeysLowLatch(DEVICE); + await signalPrekeysLowIfNeeded(DEVICE, 2); + + expect(mockEmit).toHaveBeenCalledTimes(2); + }); + + it('swallows socket-layer failures', async () => { + mockTo.mockImplementationOnce(() => { + throw new Error('socket server exploded'); + }); + + await expect(signalPrekeysLowIfNeeded(DEVICE, 1)).resolves.toBeUndefined(); + }); +}); + +describe('countAvailableOneTimePreKeys', () => { + it('returns the unconsumed one-time prekey count', async () => { + const where = vi.fn().mockResolvedValue([{ total: 7 }]); + mockSelect.mockReturnValue({ from: vi.fn().mockReturnValue({ where }) }); + + expect(await countAvailableOneTimePreKeys(DEVICE)).toBe(7); + }); + + it('returns 0 when the device has none', async () => { + const where = vi.fn().mockResolvedValue([]); + mockSelect.mockReturnValue({ from: vi.fn().mockReturnValue({ where }) }); + + expect(await countAvailableOneTimePreKeys(DEVICE)).toBe(0); + }); +}); diff --git a/apps/backend/src/__tests__/users.bundle.test.ts b/apps/backend/src/__tests__/users.bundle.test.ts index 5f5e0a7..0f37680 100644 --- a/apps/backend/src/__tests__/users.bundle.test.ts +++ b/apps/backend/src/__tests__/users.bundle.test.ts @@ -54,9 +54,15 @@ vi.mock('drizzle-orm', () => ({ ilike: vi.fn(), exists: vi.fn(), isNull: vi.fn((col: unknown) => ({ op: 'isNull', col })), + count: vi.fn(() => 'count(*)'), sql: vi.fn(), })); +const mockSignalPrekeysLow = vi.fn().mockResolvedValue(undefined); +vi.mock('../services/prekeyLowSignal.js', () => ({ + signalPrekeysLowIfNeeded: mockSignalPrekeysLow, +})); + vi.mock('../lib/redis.js', () => ({ redis: null })); vi.mock('../services/presence.js', () => ({ isOnline: vi.fn().mockResolvedValue(false) })); vi.mock('../lib/socket.js', () => ({ getSocketServer: vi.fn(() => null) })); @@ -135,17 +141,26 @@ describe('GET /users/:userId/devices/:deviceId/key-bundle', () => { const updateSet = vi.fn().mockReturnValue({ where: updateWhere }); const lockFor = vi.fn().mockResolvedValue([claimed]); const tx = { - select: vi.fn().mockReturnValue({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - orderBy: vi.fn().mockReturnValue({ - limit: vi.fn().mockReturnValue({ - for: lockFor, + // Two selects run inside the claiming transaction: the row-locked claim, + // then the post-consumption remaining count that drives `prekeys_low`. + select: vi + .fn() + .mockReturnValueOnce({ + from: vi.fn().mockReturnValue({ + where: vi.fn().mockReturnValue({ + orderBy: vi.fn().mockReturnValue({ + limit: vi.fn().mockReturnValue({ + for: lockFor, + }), }), }), }), + }) + .mockReturnValueOnce({ + from: vi.fn().mockReturnValue({ + where: vi.fn().mockResolvedValue([{ total: 4 }]), + }), }), - }), update: vi.fn().mockReturnValue({ set: updateSet }), }; mockTransaction.mockImplementation(async (cb: (txArg: typeof tx) => unknown) => cb(tx)); @@ -167,6 +182,9 @@ describe('GET /users/:userId/devices/:deviceId/key-bundle', () => { // history stays auditable. expect(updateSet).toHaveBeenCalledWith({ consumed: true }); expect(updateWhere).toHaveBeenCalled(); + // The post-consumption count feeds the low-prekey signal, which decides on + // its own whether this crossed the threshold. + expect(mockSignalPrekeysLow).toHaveBeenCalledWith('device-2', 4); }); it('uses skipLocked so concurrent bundle reads only consume one OTP', async () => { @@ -234,5 +252,8 @@ describe('GET /users/:userId/devices/:deviceId/key-bundle', () => { expect(res.status).toBe(200); expect(res.body.oneTimePreKey).toBeNull(); expect(tx.update).not.toHaveBeenCalled(); + // Nothing was consumed, so the count did not move — the crossing already + // signalled on the fetch that drained the last key. + expect(mockSignalPrekeysLow).not.toHaveBeenCalled(); }); }); diff --git a/apps/backend/src/routes/devices.ts b/apps/backend/src/routes/devices.ts index 1bf6278..e0e10a3 100644 --- a/apps/backend/src/routes/devices.ts +++ b/apps/backend/src/routes/devices.ts @@ -25,6 +25,11 @@ import { SignedPreKeyEntrySchema, PreKeyEntrySchema, verifyEd25519Signature } fr import { conversationRoom } from '../services/roomManager.js'; import { redis } from '../lib/redis.js'; import { markDeviceRevoked } from '../services/deviceRevocation.js'; +import { + PREKEY_LOW_THRESHOLD, + countAvailableOneTimePreKeys, + releasePrekeysLowLatch, +} from '../services/prekeyLowSignal.js'; import { actorFromRequest, recordAuditEvent } from '../services/auditLog.js'; export const devicesRouter: RouterType = Router(); @@ -128,6 +133,10 @@ async function revokeDeviceRow(deviceId: string): Promise { await db.update(devices).set({ revokedAt, updatedAt: revokedAt }).where(eq(devices.id, deviceId)); await db.delete(devicePrekeys).where(eq(devicePrekeys.deviceId, deviceId)); + // The device's prekeys are gone; leaving a latch behind would suppress the + // first low signal if this identity key is later re-registered. + await releasePrekeysLowLatch(deviceId); + // Force-disconnect any live socket for this device, on this node or any // other (cross-node via Redis pub/sub — see services/deviceRevocation.ts). markDeviceRevoked(deviceId); @@ -347,10 +356,20 @@ devicesRouter.post('/:id/prekeys', validate(UploadPreKeysSchema), async (req: Au }); } + // Re-arm the low-prekey signal once the device is back at or above the + // threshold, so a future crossing fires again. Recounted rather than derived + // from `currentCount + trimmedBatch.length` because duplicate keyIds are + // dropped by ON CONFLICT DO NOTHING and would inflate the derived figure. + const replenishedCount = await countAvailableOneTimePreKeys(deviceId); + if (replenishedCount >= PREKEY_LOW_THRESHOLD) { + await releasePrekeysLowLatch(deviceId); + } + res.status(200).json({ uploadedSignedPreKey: true, uploadedOneTimePreKeys: trimmedBatch.length, capped: trimmedBatch.length < otpBatch.length, + oneTimePreKeysRemaining: replenishedCount, }); }); diff --git a/apps/backend/src/routes/users.ts b/apps/backend/src/routes/users.ts index 30c9fe3..a77430e 100644 --- a/apps/backend/src/routes/users.ts +++ b/apps/backend/src/routes/users.ts @@ -1,7 +1,6 @@ import { createHash } from 'node:crypto'; import { Router, type Router as RouterType } from 'express'; -import rateLimit, { type RateLimitRequestHandler } from 'express-rate-limit'; -import { eq, and, or, ilike, exists, sql, isNull } from 'drizzle-orm'; +import { eq, and, or, ilike, exists, sql, isNull, count } from 'drizzle-orm'; import { db } from '../db/index.js'; import { users, wallets, devices, devicePrekeys, conversationMembers, deviceKeyHistory } from '../db/schema.js'; import { requireAuth, type AuthRequest } from '../middleware/auth.js'; @@ -10,6 +9,7 @@ import { redis } from '../lib/redis.js'; import { isOnline, deriveDevicePresence } from '../services/presence.js'; import { getSocketServer } from '../lib/socket.js'; import { conversationRoom } from '../services/roomManager.js'; +import { signalPrekeysLowIfNeeded } from '../services/prekeyLowSignal.js'; import { actorFromRequest, recordAuditEvent } from '../services/auditLog.js'; import { prekeyConsumedTotal } from '../lib/metrics.js'; @@ -285,7 +285,7 @@ usersRouter.get( return { keyId: candidate.keyId, publicKey: candidate.publicKey }; }); - const claimedOneTimePreKey = await db.transaction(async (tx) => { + const claimed = await db.transaction(async (tx) => { const [candidate] = await tx .select({ id: devicePrekeys.id, @@ -311,9 +311,31 @@ usersRouter.get( .set({ consumed: true }) .where(eq(devicePrekeys.id, candidate.id)); - return { keyId: candidate.keyId, publicKey: candidate.publicKey }; + // Counted inside the claiming transaction so `remaining` reflects this + // consumption and cannot race a concurrent fetch into a wrong value. + const [remainingRow] = await tx + .select({ total: count() }) + .from(devicePrekeys) + .where( + and( + eq(devicePrekeys.deviceId, deviceId), + eq(devicePrekeys.keyType, 'one_time'), + eq(devicePrekeys.consumed, false), + ), + ); + + return { + oneTimePreKey: { keyId: candidate.keyId, publicKey: candidate.publicKey }, + remaining: remainingRow?.total ?? 0, + }; }); + const claimedOneTimePreKey = claimed?.oneTimePreKey ?? null; + + // Fire-and-forget: the device that owns this bundle is told to replenish + // once per threshold crossing. Never blocks or fails the bundle response. + if (claimed) { + void signalPrekeysLowIfNeeded(deviceId, claimed.remaining); // A one-time prekey was consumed and cannot be handed out again (#376). // Draining a device's supply forces every later session with it down from // 4-DH to 3-DH, and it happens quietly, so the count left is the signal an diff --git a/apps/backend/src/services/prekeyLowSignal.ts b/apps/backend/src/services/prekeyLowSignal.ts new file mode 100644 index 0000000..df4cffe --- /dev/null +++ b/apps/backend/src/services/prekeyLowSignal.ts @@ -0,0 +1,140 @@ +/** + * Low one-time-prekey signalling. + * + * A device's one-time prekeys are consumed one per X3DH bundle fetch and are + * never regenerated server-side — once they run out, initiators fall back to a + * 3-DH bundle with no forward-secrecy contribution from a one-time key. So the + * device needs to be told to replenish *before* it hits zero. + * + * The signal fires on the bundle-fetch path (the only place the count drops), + * and is latched so a device that stays below the threshold across many + * fetches is told once, not once per fetch. The latch clears when the device + * replenishes back to or above the threshold, so the next crossing signals + * again. + */ + +import type { Redis } from 'ioredis'; +import { and, count, eq } from 'drizzle-orm'; +import { db } from '../db/index.js'; +import { devicePrekeys } from '../db/schema.js'; +import { getSocketServer } from '../lib/socket.js'; +import { redis } from '../lib/redis.js'; + +/** + * Signal when a device's unconsumed one-time prekey count drops below this. + * Overridable via `PREKEY_LOW_THRESHOLD` for deployments that burn through + * prekeys faster (large group fan-out) and want more headroom. + */ +export const PREKEY_LOW_THRESHOLD = (() => { + const raw = Number(process.env['PREKEY_LOW_THRESHOLD']); + return Number.isInteger(raw) && raw > 0 ? raw : 20; +})(); + +/** + * Safety net so a latch for a device that never replenishes (and never gets + * revoked) cannot linger in Redis forever. Expiry re-arms the signal, which is + * the desired behaviour anyway for a device that has ignored it for this long. + */ +const LATCH_TTL_SECONDS = 30 * 24 * 60 * 60; + +const latchKey = (deviceId: string): string => `prekeys:low:${deviceId}`; + +/** + * In-process fallback used when Redis is not configured (single-node dev, and + * the unit test suite). Same latch semantics, just not shared across nodes. + */ +const localLatch = new Set(); + +/** + * Claim the right to emit for this device. Returns true exactly once per + * threshold crossing — subsequent calls return false until the latch is + * released by {@link releasePrekeysLowLatch}. + * + * On Redis this is a single atomic `SET NX`, so two gateways racing on + * concurrent bundle fetches still produce only one emit. + */ +async function acquirePrekeysLowLatch(client: Redis | null, deviceId: string): Promise { + if (!client) { + if (localLatch.has(deviceId)) return false; + localLatch.add(deviceId); + return true; + } + + try { + const result = await client.set(latchKey(deviceId), '1', 'EX', LATCH_TTL_SECONDS, 'NX'); + return result === 'OK'; + } catch { + // Redis unavailable — fall back to the local latch rather than either + // spamming the device or going silent. + if (localLatch.has(deviceId)) return false; + localLatch.add(deviceId); + return true; + } +} + +/** + * Re-arm the signal for a device. Called when the device replenishes back to + * or above the threshold, and on revocation (where prekeys are deleted). + */ +export async function releasePrekeysLowLatch(deviceId: string): Promise { + localLatch.delete(deviceId); + if (!redis) return; + try { + await redis.del(latchKey(deviceId)); + } catch { + // Best effort — a stale latch only costs one missed signal, and expires. + } +} + +/** Count of a device's unconsumed one-time prekeys. */ +export async function countAvailableOneTimePreKeys(deviceId: string): Promise { + const [row] = await db + .select({ total: count() }) + .from(devicePrekeys) + .where( + and( + eq(devicePrekeys.deviceId, deviceId), + eq(devicePrekeys.keyType, 'one_time'), + eq(devicePrekeys.consumed, false), + ), + ); + return row?.total ?? 0; +} + +/** + * Emit `prekeys_low` to the device's own sockets if `remaining` has crossed + * below the threshold and the signal is not already latched. + * + * Targets the `device:{id}` room, so the event reaches that device on whichever + * gateway holds its socket (via the Socket.IO Redis adapter) and reaches no + * other device on the account — only the owner can replenish its own prekeys. + * + * Safe to call fire-and-forget: it never throws. + */ +export async function signalPrekeysLowIfNeeded(deviceId: string, remaining: number): Promise { + try { + if (remaining >= PREKEY_LOW_THRESHOLD) { + // Back in healthy territory — re-arm so the next crossing signals. + await releasePrekeysLowLatch(deviceId); + return; + } + + if (!(await acquirePrekeysLowLatch(redis, deviceId))) return; + + const io = getSocketServer(); + if (!io) return; + + io.to(`device:${deviceId}`).emit('prekeys_low', { + deviceId, + oneTimePreKeysRemaining: remaining, + threshold: PREKEY_LOW_THRESHOLD, + }); + } catch (err) { + console.warn('[prekeyLowSignal] failed for device', deviceId, (err as Error).message); + } +} + +/** Test-only: drop in-process latch state between cases. */ +export function __resetPrekeyLowLatches(): void { + localLatch.clear(); +}