cipher/src/routes/keys.js

// The key directory. This is what makes X3DH work: a sender fetches the
// recipient's "prekey bundle" (public keys) to establish a session without the
// recipient being online. The server hands out one-time prekeys and deletes them.

import { Hono } from "hono";
import { requireAuth } from "../lib/auth.js";
import { now } from "../lib/util.js";

const app = new Hono();

// GET /api/users/:username  -> directory lookup
app.get("/users/:username", requireAuth(), async (c) => {
  const uname = c.req.param("username").toLowerCase();
  const user = await c.env.DB.prepare(
    "SELECT id, username, display_name FROM users WHERE username = ?"
  )
    .bind(uname)
    .first();
  if (!user) return c.json({ error: "not found" }, 404);
  return c.json({ userId: user.id, username: user.username, displayName: user.display_name });
});

// POST /api/keys/prekeys  — replenish one-time prekeys (client tops up when low)
// body: { oneTimePreKeys: [{id,pub}, ...] }
app.post("/keys/prekeys", requireAuth(), async (c) => {
  const deviceId = c.get("deviceId");
  const body = await c.req.json().catch(() => null);
  const otks = body?.oneTimePreKeys || [];
  if (!Array.isArray(otks) || otks.length === 0)
    return c.json({ error: "no prekeys" }, 400);
  const stmts = otks.map((o) =>
    c.env.DB.prepare(
      "INSERT OR IGNORE INTO one_time_prekeys (device_id, prekey_id, prekey_pub) VALUES (?,?,?)"
    ).bind(deviceId, o.id, o.pub)
  );
  await c.env.DB.batch(stmts);
  return c.json({ ok: true, added: otks.length });
});

// POST /api/keys/signed  — rotate the signed prekey
// body: { signedPreKey: {id,pub,sig} }
app.post("/keys/signed", requireAuth(), async (c) => {
  const deviceId = c.get("deviceId");
  const body = await c.req.json().catch(() => null);
  const spk = body?.signedPreKey;
  if (!spk) return c.json({ error: "missing signedPreKey" }, 400);
  await c.env.DB.prepare(
    "UPDATE devices SET signed_prekey_id=?, signed_prekey_pub=?, signed_prekey_sig=? WHERE id=?"
  )
    .bind(spk.id, spk.pub, spk.sig, deviceId)
    .run();
  return c.json({ ok: true });
});

// GET /api/keys/count  — how many one-time prekeys remain for my device
app.get("/keys/count", requireAuth(), async (c) => {
  const deviceId = c.get("deviceId");
  const row = await c.env.DB.prepare(
    "SELECT COUNT(*) AS n FROM one_time_prekeys WHERE device_id = ?"
  )
    .bind(deviceId)
    .first();
  return c.json({ count: row?.n ?? 0 });
});

// GET /api/keys/bundle/:userId  — fetch prekey bundle(s) for every device of a user.
// Consumes one one-time prekey per device (atomic delete-and-return).
app.get("/keys/bundle/:userId", requireAuth(), async (c) => {
  const userId = c.req.param("userId");
  const devices = await c.env.DB.prepare(
    `SELECT id, registration_id, identity_key, signed_prekey_id,
            signed_prekey_pub, signed_prekey_sig FROM devices WHERE user_id = ?`
  )
    .bind(userId)
    .all();

  if (!devices.results?.length) return c.json({ error: "no devices" }, 404);

  const bundles = [];
  for (const d of devices.results) {
    // Grab and consume one one-time prekey (best-effort; bundle still valid without one).
    const otk = await c.env.DB.prepare(
      "SELECT prekey_id, prekey_pub FROM one_time_prekeys WHERE device_id = ? LIMIT 1"
    )
      .bind(d.id)
      .first();
    if (otk) {
      await c.env.DB.prepare(
        "DELETE FROM one_time_prekeys WHERE device_id = ? AND prekey_id = ?"
      )
        .bind(d.id, otk.prekey_id)
        .run();
    }
    bundles.push({
      deviceId: d.id,
      registrationId: d.registration_id,
      identityKey: d.identity_key,
      signedPreKey: { id: d.signed_prekey_id, pub: d.signed_prekey_pub, sig: d.signed_prekey_sig },
      oneTimePreKey: otk ? { id: otk.prekey_id, pub: otk.prekey_pub } : null,
    });
  }
  return c.json({ userId, bundles });
});

// POST /api/push/subscribe  — register a Web Push subscription
app.post("/push/subscribe", requireAuth(), async (c) => {
  const deviceId = c.get("deviceId");
  const sub = await c.req.json().catch(() => null);
  if (!sub?.endpoint || !sub?.keys) return c.json({ error: "bad subscription" }, 400);
  await c.env.DB.prepare(
    `INSERT INTO push_subs (device_id, endpoint, p256dh, auth, created_at)
     VALUES (?,?,?,?,?)
     ON CONFLICT(device_id) DO UPDATE SET endpoint=excluded.endpoint,
       p256dh=excluded.p256dh, auth=excluded.auth`
  )
    .bind(deviceId, sub.endpoint, sub.keys.p256dh, sub.keys.auth, now())
    .run();
  return c.json({ ok: true });
});

export default app;