DiscordClone/convex/auth.ts

import { query, mutation, internalQuery, internalMutation } from "./_generated/server";
import { v } from "convex/values";
import { getPublicStorageUrl } from "./storageUrl";
import { getRolesForUser } from "./roles";

async function sha256Hex(input: string): Promise<string> {
  const buffer = await crypto.subtle.digest(
    "SHA-256",
    new TextEncoder().encode(input)
  );
  return Array.from(new Uint8Array(buffer))
    .map((b) => b.toString(16).padStart(2, "0"))
    .join("");
}

// Get salt for a username (returns fake salt for non-existent users)
export const getSalt = query({
  args: { username: v.string() },
  returns: v.object({ salt: v.string() }),
  handler: async (ctx, args) => {
    const user = await ctx.db
      .query("userProfiles")
      .withIndex("by_username", (q) => q.eq("username", args.username))
      .unique();

    if (user) {
      return { salt: user.clientSalt };
    }

    // Generate deterministic fake salt for non-existent users (privacy)
    const fakeSalt = await sha256Hex("SERVER_SECRET_KEY" + args.username);
    return { salt: fakeSalt };
  },
});

// Verify user credentials (DAK comparison)
export const verifyUser = mutation({
  args: {
    username: v.string(),
    dak: v.string(),
  },
  returns: v.union(
    v.object({
      success: v.boolean(),
      userId: v.string(),
      encryptedMK: v.string(),
      encryptedPrivateKeys: v.string(),
      publicKey: v.string(),
    }),
    v.object({ error: v.string() })
  ),
  handler: async (ctx, args) => {
    const user = await ctx.db
      .query("userProfiles")
      .withIndex("by_username", (q) => q.eq("username", args.username))
      .unique();

    if (!user) {
      return { error: "Invalid credentials" };
    }

    const hashedDAK = await sha256Hex(args.dak);

    if (hashedDAK === user.hashedAuthKey) {
      return {
        success: true,
        userId: user._id,
        encryptedMK: user.encryptedMasterKey,
        encryptedPrivateKeys: user.encryptedPrivateKeys,
        publicKey: user.publicIdentityKey,
      };
    }

    return { error: "Invalid credentials" };
  },
});

// Register new user with crypto keys
export const createUserWithProfile = mutation({
  args: {
    username: v.string(),
    salt: v.string(),
    encryptedMK: v.string(),
    hak: v.string(),
    publicKey: v.string(),
    signingKey: v.string(),
    encryptedPrivateKeys: v.string(),
    inviteCode: v.optional(v.string()),
  },
  returns: v.union(
    v.object({ success: v.boolean(), userId: v.string() }),
    v.object({ error: v.string() })
  ),
  handler: async (ctx, args) => {
    const existing = await ctx.db
      .query("userProfiles")
      .withIndex("by_username", (q) => q.eq("username", args.username))
      .unique();

    if (existing) {
      return { error: "Username taken" };
    }

    const isFirstUser =
      (await ctx.db.query("userProfiles").first()) === null;

    if (!isFirstUser) {
      if (!args.inviteCode) {
        return { error: "Invite code required" };
      }

      const inviteCode = args.inviteCode!;
      const invite = await ctx.db
        .query("invites")
        .withIndex("by_code", (q) => q.eq("code", inviteCode))
        .unique();

      if (!invite) {
        return { error: "Invalid invite code" };
      }

      if (invite.expiresAt && Date.now() > invite.expiresAt) {
        return { error: "Invite expired" };
      }

      if (
        invite.maxUses !== undefined &&
        invite.maxUses !== null &&
        invite.uses >= invite.maxUses
      ) {
        return { error: "Invite max uses reached" };
      }

      await ctx.db.patch(invite._id, { uses: invite.uses + 1 });
    }

    const userId = await ctx.db.insert("userProfiles", {
      username: args.username,
      clientSalt: args.salt,
      encryptedMasterKey: args.encryptedMK,
      hashedAuthKey: args.hak,
      publicIdentityKey: args.publicKey,
      publicSigningKey: args.signingKey,
      encryptedPrivateKeys: args.encryptedPrivateKeys,
      isAdmin: isFirstUser,
    });

    if (isFirstUser) {
      const everyoneRoleId = await ctx.db.insert("roles", {
        name: "@everyone",
        color: "#99aab5",
        position: 0,
        permissions: {
          create_invite: true,
          embed_links: true,
          attach_files: true,
        },
        isHoist: false,
      });

      const ownerRoleId = await ctx.db.insert("roles", {
        name: "Owner",
        color: "#e91e63",
        position: 100,
        permissions: {
          manage_channels: true,
          manage_roles: true,
          manage_messages: true,
          manage_nicknames: true,
          create_invite: true,
          embed_links: true,
          attach_files: true,
        },
        isHoist: true,
      });

      await ctx.db.insert("userRoles", { userId, roleId: everyoneRoleId });
      await ctx.db.insert("userRoles", { userId, roleId: ownerRoleId });
    } else {
      const everyoneRole = await ctx.db
        .query("roles")
        .filter((q) => q.eq(q.field("name"), "@everyone"))
        .first();

      if (everyoneRole) {
        await ctx.db.insert("userRoles", {
          userId,
          roleId: everyoneRole._id,
        });
      }
    }

    return { success: true, userId };
  },
});

// Get all users' public keys
export const getPublicKeys = query({
  args: {},
  returns: v.array(
    v.object({
      id: v.string(),
      username: v.string(),
      public_identity_key: v.string(),
      status: v.optional(v.string()),
      displayName: v.optional(v.string()),
      avatarUrl: v.optional(v.union(v.string(), v.null())),
      aboutMe: v.optional(v.string()),
      customStatus: v.optional(v.string()),
      joinSoundUrl: v.optional(v.union(v.string(), v.null())),
      accentColor: v.optional(v.string()),
    })
  ),
  handler: async (ctx) => {
    const users = await ctx.db.query("userProfiles").collect();
    const results = [];
    for (const u of users) {
      let avatarUrl: string | null = null;
      if (u.avatarStorageId) {
        avatarUrl = await getPublicStorageUrl(ctx, u.avatarStorageId);
      }
      let joinSoundUrl: string | null = null;
      if (u.joinSoundStorageId) {
        joinSoundUrl = await getPublicStorageUrl(ctx, u.joinSoundStorageId);
      }
      results.push({
        id: u._id,
        username: u.username,
        public_identity_key: u.publicIdentityKey,
        status: u.status || "offline",
        displayName: u.displayName,
        avatarUrl,
        aboutMe: u.aboutMe,
        customStatus: u.customStatus,
        joinSoundUrl,
        accentColor: u.accentColor,
      });
    }
    return results;
  },
});

// Update user profile (aboutMe, avatar, customStatus)
export const updateProfile = mutation({
  args: {
    userId: v.id("userProfiles"),
    displayName: v.optional(v.string()),
    aboutMe: v.optional(v.string()),
    avatarStorageId: v.optional(v.id("_storage")),
    customStatus: v.optional(v.string()),
    joinSoundStorageId: v.optional(v.id("_storage")),
    removeJoinSound: v.optional(v.boolean()),
    accentColor: v.optional(v.string()),
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    const patch: Record<string, unknown> = {};
    if (args.displayName !== undefined) patch.displayName = args.displayName;
    if (args.aboutMe !== undefined) patch.aboutMe = args.aboutMe;
    if (args.avatarStorageId !== undefined) patch.avatarStorageId = args.avatarStorageId;
    if (args.customStatus !== undefined) patch.customStatus = args.customStatus;
    if (args.joinSoundStorageId !== undefined) patch.joinSoundStorageId = args.joinSoundStorageId;
    if (args.removeJoinSound) patch.joinSoundStorageId = undefined;
    if (args.accentColor !== undefined) patch.accentColor = args.accentColor;
    await ctx.db.patch(args.userId, patch);
    return null;
  },
});

// Get the current user's join sound URL
export const getMyJoinSoundUrl = query({
  args: { userId: v.id("userProfiles") },
  returns: v.union(v.string(), v.null()),
  handler: async (ctx, args) => {
    const user = await ctx.db.get(args.userId);
    if (!user?.joinSoundStorageId) return null;
    return await getPublicStorageUrl(ctx, user.joinSoundStorageId);
  },
});

// Update user status
export const updateStatus = mutation({
  args: {
    userId: v.id("userProfiles"),
    status: v.string(),
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    await ctx.db.patch(args.userId, { status: args.status });
    return null;
  },
});

// Get encrypted private keys + public signing key for recovery
export const getRecoveryData = query({
  args: { username: v.string() },
  returns: v.union(
    v.object({
      encryptedPrivateKeys: v.string(),
      publicSigningKey: v.string(),
    }),
    v.object({ error: v.string() })
  ),
  handler: async (ctx, args) => {
    const user = await ctx.db
      .query("userProfiles")
      .withIndex("by_username", (q) => q.eq("username", args.username))
      .unique();

    if (!user) {
      return { error: "User not found" };
    }

    return {
      encryptedPrivateKeys: user.encryptedPrivateKeys,
      publicSigningKey: user.publicSigningKey,
    };
  },
});

// Internal: get userId + publicSigningKey for recovery action verification
export const getUserForRecovery = internalQuery({
  args: { username: v.string() },
  returns: v.union(
    v.object({
      userId: v.id("userProfiles"),
      publicSigningKey: v.string(),
    }),
    v.null()
  ),
  handler: async (ctx, args) => {
    const user = await ctx.db
      .query("userProfiles")
      .withIndex("by_username", (q) => q.eq("username", args.username))
      .unique();

    if (!user) return null;

    return {
      userId: user._id,
      publicSigningKey: user.publicSigningKey,
    };
  },
});

// Set nickname (displayName) for a user
export const setNickname = mutation({
  args: {
    actorUserId: v.id("userProfiles"),
    targetUserId: v.id("userProfiles"),
    displayName: v.string(),
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    // Self-changes are always allowed
    if (args.actorUserId !== args.targetUserId) {
      const roles = await getRolesForUser(ctx, args.actorUserId);
      const canManage = roles.some(
        (role) => (role.permissions as Record<string, boolean>)?.["manage_nicknames"]
      );
      if (!canManage) {
        throw new Error("You don't have permission to change other users' nicknames");
      }
    }

    const trimmed = args.displayName.trim();
    await ctx.db.patch(args.targetUserId, {
      displayName: trimmed || undefined,
    });
    return null;
  },
});

// Delete a user and all their associated data (admin only)
export const deleteUser = mutation({
  args: {
    requestingUserId: v.id("userProfiles"),
    targetUserId: v.id("userProfiles"),
  },
  returns: v.object({ success: v.boolean(), error: v.optional(v.string()) }),
  handler: async (ctx, args) => {
    // Verify requesting user is admin
    const requester = await ctx.db.get(args.requestingUserId);
    if (!requester || !requester.isAdmin) {
      return { success: false, error: "Only admins can delete users" };
    }

    // Prevent self-deletion
    if (args.requestingUserId === args.targetUserId) {
      return { success: false, error: "Cannot delete your own account" };
    }

    const target = await ctx.db.get(args.targetUserId);
    if (!target) {
      return { success: false, error: "User not found" };
    }

    // Prevent deleting other admins
    if (target.isAdmin) {
      return { success: false, error: "Cannot delete another admin" };
    }

    // Delete reactions made by this user (before messages, using index)
    const userReactions = await ctx.db
      .query("messageReactions")
      .withIndex("by_user", (q) => q.eq("userId", args.targetUserId))
      .collect();
    for (const r of userReactions) {
      await ctx.db.delete(r._id);
    }

    // Delete all messages by this user (using index)
    const messages = await ctx.db
      .query("messages")
      .withIndex("by_sender", (q) => q.eq("senderId", args.targetUserId))
      .collect();
    for (const msg of messages) {
      // Delete reactions on this message
      const reactions = await ctx.db
        .query("messageReactions")
        .withIndex("by_message", (q) => q.eq("messageId", msg._id))
        .collect();
      for (const r of reactions) {
        await ctx.db.delete(r._id);
      }
      await ctx.db.delete(msg._id);
    }

    // Delete channel keys
    const channelKeys = await ctx.db
      .query("channelKeys")
      .withIndex("by_user", (q) => q.eq("userId", args.targetUserId))
      .collect();
    for (const ck of channelKeys) {
      await ctx.db.delete(ck._id);
    }

    // Delete role assignments
    const userRoles = await ctx.db
      .query("userRoles")
      .withIndex("by_user", (q) => q.eq("userId", args.targetUserId))
      .collect();
    for (const ur of userRoles) {
      await ctx.db.delete(ur._id);
    }

    // Delete DM participations
    const dmParts = await ctx.db
      .query("dmParticipants")
      .withIndex("by_user", (q) => q.eq("userId", args.targetUserId))
      .collect();
    for (const dp of dmParts) {
      await ctx.db.delete(dp._id);
    }

    // Delete typing indicators
    const typingIndicators = await ctx.db
      .query("typingIndicators")
      .withIndex("by_user", (q) => q.eq("userId", args.targetUserId))
      .collect();
    for (const ti of typingIndicators) {
      await ctx.db.delete(ti._id);
    }

    // Delete voice states
    const voiceStates = await ctx.db
      .query("voiceStates")
      .withIndex("by_user", (q) => q.eq("userId", args.targetUserId))
      .collect();
    for (const vs of voiceStates) {
      await ctx.db.delete(vs._id);
    }

    // Delete read states
    const readStates = await ctx.db
      .query("channelReadState")
      .withIndex("by_user", (q) => q.eq("userId", args.targetUserId))
      .collect();
    for (const rs of readStates) {
      await ctx.db.delete(rs._id);
    }

    // Delete invites created by this user
    const invites = await ctx.db
      .query("invites")
      .withIndex("by_creator", (q) => q.eq("createdBy", args.targetUserId))
      .collect();
    for (const inv of invites) {
      await ctx.db.delete(inv._id);
    }

    // Delete custom emojis uploaded by this user
    const emojis = await ctx.db
      .query("customEmojis")
      .withIndex("by_uploader", (q) => q.eq("uploadedBy", args.targetUserId))
      .collect();
    for (const emoji of emojis) {
      await ctx.storage.delete(emoji.storageId);
      await ctx.db.delete(emoji._id);
    }

    // Delete avatar from storage if exists
    if (target.avatarStorageId) {
      await ctx.storage.delete(target.avatarStorageId);
    }

    // Delete join sound from storage if exists
    if (target.joinSoundStorageId) {
      await ctx.storage.delete(target.joinSoundStorageId);
    }

    // Delete the user profile
    await ctx.db.delete(args.targetUserId);

    return { success: true };
  },
});

// Internal: update credentials after password reset
export const updateCredentials = internalMutation({
  args: {
    userId: v.id("userProfiles"),
    clientSalt: v.string(),
    encryptedMasterKey: v.string(),
    hashedAuthKey: v.string(),
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    await ctx.db.patch(args.userId, {
      clientSalt: args.clientSalt,
      encryptedMasterKey: args.encryptedMasterKey,
      hashedAuthKey: args.hashedAuthKey,
    });
    return null;
  },
});