DiscordClone/convex/channelKeys.ts

import { query, mutation } from "./_generated/server";
import { v } from "convex/values";

/**
 * Rotate the symmetric key for a DM channel. Inserts a brand-new
 * versioned row for each participant — existing rows are left alone
 * so previously-encrypted messages remain decryptable.
 *
 * The caller proves they're a DM participant by passing their own
 * userId; the server cross-checks against `dmParticipants` for the
 * channel. Every recipient userId in `entries` must also be a
 * participant — no leaking keys to random users.
 *
 * The new rows are tagged with `maxExistingVersion + 1`.
 */
export const rotateDMKey = mutation({
  args: {
    channelId: v.id("channels"),
    initiatorUserId: v.id("userProfiles"),
    entries: v.array(
      v.object({
        userId: v.id("userProfiles"),
        encryptedKeyBundle: v.string(),
      }),
    ),
  },
  returns: v.object({ keyVersion: v.number() }),
  handler: async (ctx, args) => {
    const channel = await ctx.db.get(args.channelId);
    if (!channel) throw new Error("Channel not found");
    if (channel.type !== "dm") {
      throw new Error("rotateDMKey is only supported for DM channels");
    }

    // Verify every (initiator + entries) userId is in dmParticipants.
    const participants = await ctx.db
      .query("dmParticipants")
      .withIndex("by_channel", (q) => q.eq("channelId", args.channelId))
      .collect();
    const participantSet = new Set(participants.map((p) => p.userId as string));
    if (!participantSet.has(args.initiatorUserId as unknown as string)) {
      throw new Error("Not a participant in this DM");
    }
    for (const entry of args.entries) {
      if (!participantSet.has(entry.userId as unknown as string)) {
        throw new Error("Target userId is not a participant in this DM");
      }
    }

    // Find the current max keyVersion for this channel. New rows go
    // one above that. If no rows exist yet, start at 2 so legacy
    // messages tagged version 1 still hit their original key.
    const existing = await ctx.db
      .query("channelKeys")
      .withIndex("by_channel", (q) => q.eq("channelId", args.channelId))
      .collect();
    const maxVersion = existing.reduce(
      (m, k) => (k.keyVersion > m ? k.keyVersion : m),
      0,
    );
    const newVersion = maxVersion + 1;

    for (const entry of args.entries) {
      await ctx.db.insert("channelKeys", {
        channelId: args.channelId,
        userId: entry.userId,
        encryptedKeyBundle: entry.encryptedKeyBundle,
        keyVersion: newVersion,
      });
    }

    return { keyVersion: newVersion };
  },
});

// Batch upsert encrypted key bundles
export const uploadKeys = mutation({
  args: {
    keys: v.array(
      v.object({
        channelId: v.id("channels"),
        userId: v.id("userProfiles"),
        encryptedKeyBundle: v.string(),
        keyVersion: v.number(),
      })
    ),
  },
  returns: v.object({ success: v.boolean(), count: v.number() }),
  handler: async (ctx, args) => {
    for (const keyData of args.keys) {
      if (!keyData.channelId || !keyData.userId || !keyData.encryptedKeyBundle) {
        continue;
      }

      // Check if exists (upsert)
      const existing = await ctx.db
        .query("channelKeys")
        .withIndex("by_channel_and_user", (q) =>
          q.eq("channelId", keyData.channelId).eq("userId", keyData.userId)
        )
        .unique();

      if (existing) {
        await ctx.db.patch(existing._id, {
          encryptedKeyBundle: keyData.encryptedKeyBundle,
          keyVersion: keyData.keyVersion,
        });
      } else {
        await ctx.db.insert("channelKeys", {
          channelId: keyData.channelId,
          userId: keyData.userId,
          encryptedKeyBundle: keyData.encryptedKeyBundle,
          keyVersion: keyData.keyVersion,
        });
      }
    }

    return { success: true, count: args.keys.length };
  },
});

// Get user's encrypted key bundles (reactive!)
export const getKeysForUser = query({
  args: { userId: v.id("userProfiles") },
  returns: v.array(
    v.object({
      channel_id: v.id("channels"),
      encrypted_key_bundle: v.string(),
      key_version: v.number(),
    })
  ),
  handler: async (ctx, args) => {
    const keys = await ctx.db
      .query("channelKeys")
      .withIndex("by_user", (q) => q.eq("userId", args.userId))
      .collect();

    return keys.map((k) => ({
      channel_id: k.channelId,
      encrypted_key_bundle: k.encryptedKeyBundle,
      key_version: k.keyVersion,
    }));
  },
});