DiscordClone/convex/auth.ts

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

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 invite = await ctx.db
        .query("invites")
        .withIndex("by_code", (q) => q.eq("code", args.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,
          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()),
    })
  ),
  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);
      }
      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,
      });
    }
    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()),
  },
  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;
    await ctx.db.patch(args.userId, patch);
    return null;
  },
});

// 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;
  },
});