From 339dd5e0fc231d63a37cf509fa838d52fa847168 Mon Sep 17 00:00:00 2001 From: Andreschuks101 Date: Fri, 26 Jun 2026 00:00:08 +0100 Subject: [PATCH] feat: fetch and consume device prekey bundle --- .../drizzle/0007_device_key_bundles.sql | 29 ++++ apps/backend/drizzle/meta/_journal.json | 7 + apps/backend/src/__tests__/keyBundle.test.ts | 136 ++++++++++++++++++ .../__tests__/users.keyBundle.routes.test.ts | 101 +++++++++++++ apps/backend/src/db/schema.ts | 77 +++++++++- apps/backend/src/routes/users.ts | 26 ++++ apps/backend/src/services/keyBundle.ts | 86 +++++++++++ 7 files changed, 461 insertions(+), 1 deletion(-) create mode 100644 apps/backend/drizzle/0007_device_key_bundles.sql create mode 100644 apps/backend/src/__tests__/keyBundle.test.ts create mode 100644 apps/backend/src/__tests__/users.keyBundle.routes.test.ts create mode 100644 apps/backend/src/services/keyBundle.ts diff --git a/apps/backend/drizzle/0007_device_key_bundles.sql b/apps/backend/drizzle/0007_device_key_bundles.sql new file mode 100644 index 0000000..3a7eec8 --- /dev/null +++ b/apps/backend/drizzle/0007_device_key_bundles.sql @@ -0,0 +1,29 @@ +CREATE TABLE "devices" ( + "id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL, + "user_id" uuid NOT NULL, + "identity_public_key" text NOT NULL, + "registration_id" integer NOT NULL, + "signed_pre_key_id" integer NOT NULL, + "signed_pre_key_public" text NOT NULL, + "signed_pre_key_signature" text NOT NULL, + "revoked_at" timestamp, + "created_at" timestamp DEFAULT now() NOT NULL +); +--> statement-breakpoint +CREATE TABLE "one_time_pre_keys" ( + "id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL, + "device_id" uuid NOT NULL, + "key_id" integer NOT NULL, + "public_key" text NOT NULL, + "consumed" boolean DEFAULT false NOT NULL, + "created_at" timestamp DEFAULT now() NOT NULL, + CONSTRAINT "one_time_pre_keys_device_key_unique" UNIQUE("device_id","key_id") +); +--> statement-breakpoint +ALTER TABLE "devices" ADD CONSTRAINT "devices_user_id_users_id_fk" FOREIGN KEY ("user_id") REFERENCES "public"."users"("id") ON DELETE cascade ON UPDATE no action; +--> statement-breakpoint +ALTER TABLE "one_time_pre_keys" ADD CONSTRAINT "one_time_pre_keys_device_id_devices_id_fk" FOREIGN KEY ("device_id") REFERENCES "public"."devices"("id") ON DELETE cascade ON UPDATE no action; +--> statement-breakpoint +CREATE INDEX "devices_user_id_idx" ON "devices" USING btree ("user_id"); +--> statement-breakpoint +CREATE INDEX "one_time_pre_keys_device_consumed_idx" ON "one_time_pre_keys" USING btree ("device_id","consumed"); diff --git a/apps/backend/drizzle/meta/_journal.json b/apps/backend/drizzle/meta/_journal.json index 8fea85e..1baaa94 100644 --- a/apps/backend/drizzle/meta/_journal.json +++ b/apps/backend/drizzle/meta/_journal.json @@ -50,6 +50,13 @@ "when": 1780560000000, "tag": "0006_add_conversation_avatar_url", "breakpoints": true + }, + { + "idx": 7, + "version": "7", + "when": 1780646400000, + "tag": "0007_device_key_bundles", + "breakpoints": true } ] } \ No newline at end of file diff --git a/apps/backend/src/__tests__/keyBundle.test.ts b/apps/backend/src/__tests__/keyBundle.test.ts new file mode 100644 index 0000000..ecf253d --- /dev/null +++ b/apps/backend/src/__tests__/keyBundle.test.ts @@ -0,0 +1,136 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +const mockFindDevice = vi.fn(); +const mockExecute = vi.fn(); + +vi.mock('../db/index.js', () => ({ + db: { + query: { devices: { findFirst: mockFindDevice } }, + execute: mockExecute, + }, +})); + +vi.mock('../db/schema.js', () => ({ + devices: { id: 'id', userId: 'userId' }, + oneTimePreKeys: { deviceId: 'deviceId', keyId: 'keyId', consumed: 'consumed' }, +})); + +vi.mock('drizzle-orm', () => ({ + and: (...args: unknown[]) => args.filter(Boolean), + eq: (col: unknown, val: unknown) => ({ col, val }), + sql: vi.fn(), +})); + +const { fetchAndConsumeKeyBundle } = await import('../services/keyBundle.js'); + +const DEVICE = { + id: 'dev-1', + userId: 'user-1', + identityPublicKey: 'identity-pub', + registrationId: 4242, + signedPreKeyId: 7, + signedPreKeyPublic: 'spk-pub', + signedPreKeySignature: 'spk-sig', + revokedAt: null, + createdAt: new Date('2026-01-01'), +}; + +beforeEach(() => { + vi.clearAllMocks(); +}); + +describe('fetchAndConsumeKeyBundle', () => { + it('returns 404 for an unknown device and does not touch prekeys', async () => { + mockFindDevice.mockResolvedValue(undefined); + + const result = await fetchAndConsumeKeyBundle('user-1', 'dev-1'); + + expect(result).toEqual({ ok: false, status: 404, error: 'Device not found' }); + expect(mockExecute).not.toHaveBeenCalled(); + }); + + it('returns 404 for a revoked device and does not consume a prekey', async () => { + mockFindDevice.mockResolvedValue({ ...DEVICE, revokedAt: new Date('2026-06-01') }); + + const result = await fetchAndConsumeKeyBundle('user-1', 'dev-1'); + + expect(result).toMatchObject({ ok: false, status: 404 }); + expect(mockExecute).not.toHaveBeenCalled(); + }); + + it('returns the bundle and consumes one one-time prekey', async () => { + mockFindDevice.mockResolvedValue(DEVICE); + mockExecute.mockResolvedValue([{ keyId: 100, publicKey: 'otp-pub' }]); + + const result = await fetchAndConsumeKeyBundle('user-1', 'dev-1'); + + expect(result).toEqual({ + ok: true, + bundle: { + identityPublicKey: 'identity-pub', + registrationId: 4242, + signedPreKey: { keyId: 7, publicKey: 'spk-pub', signature: 'spk-sig' }, + oneTimePreKey: { keyId: 100, publicKey: 'otp-pub' }, + }, + }); + expect(mockExecute).toHaveBeenCalledTimes(1); + }); + + it('returns oneTimePreKey: null when the pool is exhausted', async () => { + mockFindDevice.mockResolvedValue(DEVICE); + mockExecute.mockResolvedValue([]); + + const result = await fetchAndConsumeKeyBundle('user-1', 'dev-1'); + + expect(result.ok).toBe(true); + if (result.ok) { + expect(result.bundle.oneTimePreKey).toBeNull(); + // The signed prekey is still served so a session can be established. + expect(result.bundle.signedPreKey.keyId).toBe(7); + } + }); + + it('never exposes private key material', async () => { + mockFindDevice.mockResolvedValue(DEVICE); + mockExecute.mockResolvedValue([{ keyId: 100, publicKey: 'otp-pub' }]); + + const result = await fetchAndConsumeKeyBundle('user-1', 'dev-1'); + + if (!result.ok) throw new Error('expected a bundle'); + const serialized = JSON.stringify(result.bundle).toLowerCase(); + expect(serialized).not.toContain('private'); + expect(serialized).not.toContain('secret'); + expect(Object.keys(result.bundle).sort()).toEqual([ + 'identityPublicKey', + 'oneTimePreKey', + 'registrationId', + 'signedPreKey', + ]); + }); + + it('hands out distinct prekeys to concurrent fetches, then null', async () => { + // Emulate the DB-side atomic claim: each UPDATE pops one key from the pool. + const pool = [ + { keyId: 1, publicKey: 'a' }, + { keyId: 2, publicKey: 'b' }, + ]; + mockFindDevice.mockResolvedValue(DEVICE); + mockExecute.mockImplementation(() => { + const claimed = pool.shift(); + return Promise.resolve(claimed ? [claimed] : []); + }); + + const [first, second, third] = await Promise.all([ + fetchAndConsumeKeyBundle('user-1', 'dev-1'), + fetchAndConsumeKeyBundle('user-1', 'dev-1'), + fetchAndConsumeKeyBundle('user-1', 'dev-1'), + ]); + + const otps = [first, second, third].map((r) => (r.ok ? r.bundle.oneTimePreKey : undefined)); + const issued = otps.filter((o): o is { keyId: number; publicKey: string } => o != null); + + expect(issued).toHaveLength(2); + expect(new Set(issued.map((o) => o.keyId)).size).toBe(2); // no key handed out twice + expect(otps.filter((o) => o === null)).toHaveLength(1); // exhausted fetch gets null + }); +}); diff --git a/apps/backend/src/__tests__/users.keyBundle.routes.test.ts b/apps/backend/src/__tests__/users.keyBundle.routes.test.ts new file mode 100644 index 0000000..df16718 --- /dev/null +++ b/apps/backend/src/__tests__/users.keyBundle.routes.test.ts @@ -0,0 +1,101 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import request from 'supertest'; +import express from 'express'; + +const mockFetchBundle = vi.fn(); + +vi.mock('../services/keyBundle.js', () => ({ + fetchAndConsumeKeyBundle: mockFetchBundle, +})); + +vi.mock('../db/index.js', () => ({ + db: { + query: { users: { findFirst: vi.fn(), findMany: vi.fn() } }, + update: vi.fn(), + select: vi.fn(), + }, +})); + +vi.mock('../lib/redis.js', () => ({ + get redis() { + return null; + }, +})); + +vi.mock('../services/presence.js', () => ({ + isOnline: vi.fn(), +})); + +vi.mock('../middleware/auth.js', () => ({ + requireAuth: (req: express.Request, _res: express.Response, next: express.NextFunction) => { + (req as express.Request & { auth: { userId: string } }).auth = { userId: 'caller-1' }; + next(); + }, +})); + +const { usersRouter } = await import('../routes/users.js'); + +function makeApp() { + const app = express(); + app.use(express.json()); + app.use('/users', usersRouter); + return app; +} + +beforeEach(() => { + vi.clearAllMocks(); +}); + +describe('GET /users/:userId/devices/:deviceId/key-bundle', () => { + it('returns the prekey bundle and forwards the path params', async () => { + const bundle = { + identityPublicKey: 'identity-pub', + registrationId: 4242, + signedPreKey: { keyId: 7, publicKey: 'spk-pub', signature: 'spk-sig' }, + oneTimePreKey: { keyId: 100, publicKey: 'otp-pub' }, + }; + mockFetchBundle.mockResolvedValue({ ok: true, bundle }); + + const res = await request(makeApp()).get('/users/user-1/devices/dev-1/key-bundle'); + + expect(res.status).toBe(200); + expect(res.body).toEqual(bundle); + expect(mockFetchBundle).toHaveBeenCalledWith('user-1', 'dev-1'); + }); + + it('returns a bundle with oneTimePreKey: null when the pool is exhausted', async () => { + mockFetchBundle.mockResolvedValue({ + ok: true, + bundle: { + identityPublicKey: 'identity-pub', + registrationId: 4242, + signedPreKey: { keyId: 7, publicKey: 'spk-pub', signature: 'spk-sig' }, + oneTimePreKey: null, + }, + }); + + const res = await request(makeApp()).get('/users/user-1/devices/dev-1/key-bundle'); + + expect(res.status).toBe(200); + expect(res.body.oneTimePreKey).toBeNull(); + }); + + it('returns 404 for an unknown or revoked device', async () => { + mockFetchBundle.mockResolvedValue({ ok: false, status: 404, error: 'Device not found' }); + + const res = await request(makeApp()).get('/users/user-1/devices/dev-1/key-bundle'); + + expect(res.status).toBe(404); + expect(res.body).toEqual({ error: 'Device not found' }); + }); + + it('requires authentication', async () => { + // requireAuth is mocked to always authenticate; assert the route sits behind it + // by confirming the handler runs only after auth injected req.auth. + mockFetchBundle.mockResolvedValue({ ok: false, status: 404, error: 'Device not found' }); + + await request(makeApp()).get('/users/user-1/devices/dev-1/key-bundle'); + + expect(mockFetchBundle).toHaveBeenCalled(); + }); +}); diff --git a/apps/backend/src/db/schema.ts b/apps/backend/src/db/schema.ts index 419c59f..ccc8e0e 100644 --- a/apps/backend/src/db/schema.ts +++ b/apps/backend/src/db/schema.ts @@ -1,4 +1,14 @@ -import { pgTable, text, timestamp, uuid, boolean, pgEnum, index } from 'drizzle-orm/pg-core'; +import { + pgTable, + text, + timestamp, + uuid, + boolean, + integer, + pgEnum, + index, + unique, +} from 'drizzle-orm/pg-core'; import { relations, sql } from 'drizzle-orm'; export const users = pgTable('users', { @@ -19,6 +29,57 @@ export const wallets = pgTable('wallets', { createdAt: timestamp('created_at').notNull().defaultNow(), }); +// ─── Devices & E2E prekey bundles (#160) ─────────────────────────────────────── +// +// Every device advertises an X3DH/Signal-style key bundle so other users can +// open an end-to-end encrypted session with it: +// - a long-term `identityPublicKey` + numeric `registrationId` +// - one medium-term signed prekey (`signedPreKey*`), and +// - a pool of single-use one-time prekeys (`one_time_pre_keys`). +// +// Only PUBLIC key material and signatures are stored here — private keys never +// leave the owning client. A one-time prekey is handed out at most once: it is +// claimed with a single atomic `UPDATE ... WHERE consumed = false ... RETURNING` +// so concurrent senders can never receive the same key. `revokedAt` soft-revokes +// a device, after which its bundle is no longer served. + +export const devices = pgTable( + 'devices', + { + id: uuid('id').primaryKey().defaultRandom(), + userId: uuid('user_id') + .notNull() + .references(() => users.id, { onDelete: 'cascade' }), + identityPublicKey: text('identity_public_key').notNull(), + registrationId: integer('registration_id').notNull(), + signedPreKeyId: integer('signed_pre_key_id').notNull(), + signedPreKeyPublic: text('signed_pre_key_public').notNull(), + signedPreKeySignature: text('signed_pre_key_signature').notNull(), + revokedAt: timestamp('revoked_at'), + createdAt: timestamp('created_at').notNull().defaultNow(), + }, + (table) => [index('devices_user_id_idx').on(table.userId)], +); + +export const oneTimePreKeys = pgTable( + 'one_time_pre_keys', + { + id: uuid('id').primaryKey().defaultRandom(), + deviceId: uuid('device_id') + .notNull() + .references(() => devices.id, { onDelete: 'cascade' }), + keyId: integer('key_id').notNull(), + publicKey: text('public_key').notNull(), + consumed: boolean('consumed').notNull().default(false), + createdAt: timestamp('created_at').notNull().defaultNow(), + }, + (table) => [ + // Partial-friendly lookup of the next unconsumed key for a device. + index('one_time_pre_keys_device_consumed_idx').on(table.deviceId, table.consumed), + unique('one_time_pre_keys_device_key_unique').on(table.deviceId, table.keyId), + ], +); + // ─── Conversations ──────────────────────────────────────────────────────────── export const conversationTypeEnum = pgEnum('conversation_type', ['dm', 'group']); @@ -98,12 +159,22 @@ export const usersRelations = relations(users, ({ many }) => ({ memberships: many(conversationMembers), messages: many(messages), transfers: many(tokenTransfers), + devices: many(devices), })); export const walletsRelations = relations(wallets, ({ one }) => ({ user: one(users, { fields: [wallets.userId], references: [users.id] }), })); +export const devicesRelations = relations(devices, ({ one, many }) => ({ + user: one(users, { fields: [devices.userId], references: [users.id] }), + oneTimePreKeys: many(oneTimePreKeys), +})); + +export const oneTimePreKeysRelations = relations(oneTimePreKeys, ({ one }) => ({ + device: one(devices, { fields: [oneTimePreKeys.deviceId], references: [devices.id] }), +})); + export const conversationsRelations = relations(conversations, ({ many }) => ({ members: many(conversationMembers), messages: many(messages), @@ -150,3 +221,7 @@ export type Message = typeof messages.$inferSelect; export type NewMessage = typeof messages.$inferInsert; export type TokenTransfer = typeof tokenTransfers.$inferSelect; export type NewTokenTransfer = typeof tokenTransfers.$inferInsert; +export type Device = typeof devices.$inferSelect; +export type NewDevice = typeof devices.$inferInsert; +export type OneTimePreKey = typeof oneTimePreKeys.$inferSelect; +export type NewOneTimePreKey = typeof oneTimePreKeys.$inferInsert; diff --git a/apps/backend/src/routes/users.ts b/apps/backend/src/routes/users.ts index 4498551..3883079 100644 --- a/apps/backend/src/routes/users.ts +++ b/apps/backend/src/routes/users.ts @@ -5,6 +5,7 @@ import { users, wallets } from '../db/schema.js'; import { requireAuth, type AuthRequest } from '../middleware/auth.js'; import { redis } from '../lib/redis.js'; import { isOnline } from '../services/presence.js'; +import { fetchAndConsumeKeyBundle } from '../services/keyBundle.js'; export const usersRouter: RouterType = Router(); @@ -141,6 +142,31 @@ usersRouter.get('/:id', async (req: AuthRequest, res) => { } }); +// GET /users/:userId/devices/:deviceId/key-bundle — fetch + consume a prekey bundle. +// +// Returns the recipient device's public prekey bundle and atomically consumes one +// one-time prekey so no two senders are handed the same one. When the one-time +// pool is exhausted the bundle is returned with `oneTimePreKey: null`. Unknown or +// revoked devices return 404. Only public key material is ever returned. +usersRouter.get('/:userId/devices/:deviceId/key-bundle', async (req: AuthRequest, res) => { + const userId = req.params['userId'] as string | undefined; + const deviceId = req.params['deviceId'] as string | undefined; + + if (!userId || !deviceId) { + res.status(400).json({ error: 'userId and deviceId are required' }); + return; + } + + const result = await fetchAndConsumeKeyBundle(userId, deviceId); + + if (!result.ok) { + res.status(result.status).json({ error: result.error }); + return; + } + + res.json(result.bundle); +}); + usersRouter.get('/:id/presence', async (req: AuthRequest, res) => { const id = req.params['id'] as string; if (!redis) { diff --git a/apps/backend/src/services/keyBundle.ts b/apps/backend/src/services/keyBundle.ts new file mode 100644 index 0000000..df732cd --- /dev/null +++ b/apps/backend/src/services/keyBundle.ts @@ -0,0 +1,86 @@ +/** + * Prekey bundle fetch + one-time prekey consumption (#160). + * + * Builds the X3DH/Signal-style key bundle a sender needs to start an encrypted + * session with a recipient device. The single one-time prekey (OTP) in the + * bundle is *consumed* as it is handed out: it is claimed with one atomic + * `UPDATE ... WHERE consumed = false ... RETURNING` guarded by `FOR UPDATE SKIP + * LOCKED`, so two senders fetching concurrently can never receive the same OTP. + * When the pool is exhausted the bundle is still returned with `oneTimePreKey: + * null` — sessions can be established from the signed prekey alone. + * + * Only public key material and signatures are exposed; private keys never reach + * the server, so nothing private can ever be returned. + */ +import { and, eq, sql } from 'drizzle-orm'; +import { db } from '../db/index.js'; +import { devices, oneTimePreKeys } from '../db/schema.js'; + +export interface PreKeyBundle { + identityPublicKey: string; + registrationId: number; + signedPreKey: { keyId: number; publicKey: string; signature: string }; + oneTimePreKey: { keyId: number; publicKey: string } | null; +} + +export type KeyBundleResult = + | { ok: true; bundle: PreKeyBundle } + | { ok: false; status: 404; error: string }; + +/** + * Atomically claim the next unconsumed one-time prekey for a device. + * + * The whole select-and-mark is a single statement, so it is race-free under + * concurrent fetches: `FOR UPDATE SKIP LOCKED` makes parallel callers skip a row + * another transaction is already claiming rather than block on or re-read it. + * Returns `null` when no unconsumed prekey remains. + */ +async function consumeOneTimePreKey( + deviceId: string, +): Promise<{ keyId: number; publicKey: string } | null> { + const rows = await db.execute<{ keyId: number; publicKey: string }>(sql` + UPDATE ${oneTimePreKeys} + SET consumed = true + WHERE id = ( + SELECT id + FROM ${oneTimePreKeys} + WHERE device_id = ${deviceId} AND consumed = false + ORDER BY key_id + FOR UPDATE SKIP LOCKED + LIMIT 1 + ) + RETURNING key_id AS "keyId", public_key AS "publicKey" + `); + + return rows[0] ?? null; +} + +export async function fetchAndConsumeKeyBundle( + userId: string, + deviceId: string, +): Promise { + const device = await db.query.devices.findFirst({ + where: and(eq(devices.id, deviceId), eq(devices.userId, userId)), + }); + + // Unknown or revoked devices are indistinguishable to callers — both 404. + if (!device || device.revokedAt) { + return { ok: false, status: 404, error: 'Device not found' }; + } + + const oneTimePreKey = await consumeOneTimePreKey(deviceId); + + return { + ok: true, + bundle: { + identityPublicKey: device.identityPublicKey, + registrationId: device.registrationId, + signedPreKey: { + keyId: device.signedPreKeyId, + publicKey: device.signedPreKeyPublic, + signature: device.signedPreKeySignature, + }, + oneTimePreKey, + }, + }; +}