From 03e960d5ee265e46eb1c498f9215ff09c1d40040 Mon Sep 17 00:00:00 2001
From: Nathan Wang <149958172+nwang783@users.noreply.github.com>
Date: Sat, 11 Apr 2026 21:10:13 -0400
Subject: [PATCH] Fix join-code lookup without collection-group dependency

---
 firestore.indexes.json      | 17 +++++++++-
 functions/index.js          | 63 +++++++++++++++++++++++++++++++------
 test/join-flow-api.test.cjs | 48 +++++++++++++++++++++++++++-
 3 files changed, 117 insertions(+), 11 deletions(-)

diff --git a/firestore.indexes.json b/firestore.indexes.json
index 415027e..7c87dc9 100644
--- a/firestore.indexes.json
+++ b/firestore.indexes.json
@@ -1,4 +1,19 @@
 {
   "indexes": [],
-  "fieldOverrides": []
+  "fieldOverrides": [
+    {
+      "collectionGroup": "joinCodes",
+      "fieldPath": "codeHash",
+      "indexes": [
+        {
+          "order": "ASCENDING",
+          "queryScope": "COLLECTION_GROUP"
+        },
+        {
+          "order": "DESCENDING",
+          "queryScope": "COLLECTION_GROUP"
+        }
+      ]
+    }
+  ]
 }
diff --git a/functions/index.js b/functions/index.js
index 158ba31..b36dbb0 100644
--- a/functions/index.js
+++ b/functions/index.js
@@ -117,15 +117,25 @@ async function requireScopedTeamAdmin(req, requestedTeamId, deps = {}) {
 
 function createJoinCodeDoc(dbClient, batch, teamId, { seatId, seatName }) {
   const joinCode = generateJoinCode();
+  const codeHash = sha256(joinCode);
   const joinCodeRef = dbClient.doc(`teams/${teamId}/joinCodes/${uuidv4()}`);
-  batch.set(joinCodeRef, {
-    codeHash: sha256(joinCode),
+  const joinCodeDoc = {
+    codeHash,
     role: "member",
     uses: 0,
     createdAt: FieldValue.serverTimestamp(),
     createdBySeatId: seatId,
     createdBySeatName: seatName,
-  });
+  };
+
+  batch.set(joinCodeRef, joinCodeDoc);
+  batch.set(joinCodeLookupRef(dbClient, codeHash), buildJoinCodeLookupDoc({
+    teamId,
+    joinCodeRef,
+    role: joinCodeDoc.role,
+    createdBySeatId: seatId,
+    createdBySeatName: seatName,
+  }));
 
   return {
     joinCode,
@@ -133,6 +143,21 @@ function createJoinCodeDoc(dbClient, batch, teamId, { seatId, seatName }) {
   };
 }
 
+function joinCodeLookupRef(dbClient, codeHash) {
+  return dbClient.doc(`joinCodeLookups/${codeHash}`);
+}
+
+function buildJoinCodeLookupDoc({ teamId, joinCodeRef, role, createdBySeatId, createdBySeatName }) {
+  return {
+    teamId,
+    joinCodePath: joinCodeRef.path,
+    role,
+    createdAt: FieldValue.serverTimestamp(),
+    createdBySeatId,
+    createdBySeatName,
+  };
+}
+
 async function issueJoinCodeForTeam({ dbClient = db, teamId, seatId, seatName }) {
   const batch = dbClient.batch();
   const { joinCode, joinCodeRef } = createJoinCodeDoc(dbClient, batch, teamId, {
@@ -244,14 +269,21 @@ async function createTeamReport({
 
 async function joinTeamWithCode({ dbClient = db, adminClient = admin, joinCode, seatName }) {
   const codeHash = sha256(joinCode);
+  const lookupRef = joinCodeLookupRef(dbClient, codeHash);
+  const lookupSnap = await lookupRef.get();
+  const lookup = lookupSnap.exists ? lookupSnap.data() : null;
 
-  const codesSnap = await dbClient
-    .collectionGroup("joinCodes")
-    .where("codeHash", "==", codeHash)
-    .limit(1)
-    .get();
+  let joinCodeRef = lookup?.joinCodePath ? dbClient.doc(lookup.joinCodePath) : null;
+
+  if (!joinCodeRef) {
+    const codesSnap = await dbClient
+      .collectionGroup("joinCodes")
+      .where("codeHash", "==", codeHash)
+      .limit(1)
+      .get();
+    joinCodeRef = codesSnap.empty ? null : codesSnap.docs[0].ref;
+  }
 
-  const joinCodeRef = codesSnap.empty ? null : codesSnap.docs[0].ref;
   const teamId = joinCodeRef ? joinCodeRef.parent.parent.id : null;
 
   if (!joinCodeRef) {
@@ -262,6 +294,12 @@ async function joinTeamWithCode({ dbClient = db, adminClient = admin, joinCode,
 
   const result = await dbClient.runTransaction(async (tx) => {
     const freshSnap = await tx.get(joinCodeRef);
+    if (!freshSnap.exists) {
+      const err = new Error("Invalid join code");
+      err.status = 404;
+      throw err;
+    }
+
     const data = freshSnap.data();
 
     if (data.expiresAt && data.expiresAt.toDate() < new Date()) {
@@ -281,6 +319,13 @@ async function joinTeamWithCode({ dbClient = db, adminClient = admin, joinCode,
     const now = FieldValue.serverTimestamp();
 
     tx.update(joinCodeRef, { uses: (data.uses || 0) + 1 });
+    tx.set(lookupRef, buildJoinCodeLookupDoc({
+      teamId,
+      joinCodeRef,
+      role: data.role || "member",
+      createdBySeatId: data.createdBySeatId || null,
+      createdBySeatName: data.createdBySeatName || null,
+    }), { merge: true });
 
     tx.set(dbClient.doc(`teams/${teamId}/memberships/${seatId}`), {
       role: data.role || "member",
diff --git a/test/join-flow-api.test.cjs b/test/join-flow-api.test.cjs
index 0786d7f..77dc457 100644
--- a/test/join-flow-api.test.cjs
+++ b/test/join-flow-api.test.cjs
@@ -199,9 +199,13 @@ test('issueJoinCodeForTeam always stores member invites', async () => {
   const stored = db._store.get(`teams/team-1/joinCodes/${result.joinCodeId}`);
   assert.equal(stored.role, 'member');
   assert.equal(stored.createdBySeatName, 'Admin Seat');
+
+  const lookup = db._store.get(`joinCodeLookups/${stored.codeHash}`);
+  assert.equal(lookup.teamId, 'team-1');
+  assert.equal(lookup.joinCodePath, `teams/team-1/joinCodes/${result.joinCodeId}`);
 });
 
-test('joinTeamWithCode lands the new seat on the same team', async () => {
+test('joinTeamWithCode lands the new seat on the same team and backfills lookup docs for legacy codes', async () => {
   const db = createMemoryDb({
     'teams/team-1/joinCodes/code-1': {
       codeHash: sha256('SP7E-YKDH-LPC3'),
@@ -229,6 +233,48 @@ test('joinTeamWithCode lands the new seat on the same team', async () => {
   const seatEntries = [...db._store.entries()].filter(([path]) => path.startsWith('seats/'));
   assert.equal(seatEntries.length, 1);
   assert.equal(seatEntries[0][1].homeTeamId, 'team-1');
+
+  const lookup = db._store.get(`joinCodeLookups/${sha256('SP7E-YKDH-LPC3')}`);
+  assert.equal(lookup.teamId, 'team-1');
+  assert.equal(lookup.joinCodePath, 'teams/team-1/joinCodes/code-1');
+});
+
+test('joinTeamWithCode prefers the top-level lookup over collection-group search', async () => {
+  const joinCode = 'SP7E-YKDH-LPC3';
+  const codeHash = sha256(joinCode);
+  const db = createMemoryDb({
+    'joinCodeLookups/55a0a645-should-not-be-used': {
+      teamId: 'team-other',
+      joinCodePath: 'teams/team-other/joinCodes/code-other',
+    },
+    [`joinCodeLookups/${codeHash}`]: {
+      teamId: 'team-1',
+      joinCodePath: 'teams/team-1/joinCodes/code-1',
+      role: 'member',
+    },
+    'teams/team-1/joinCodes/code-1': {
+      codeHash,
+      role: 'member',
+      uses: 0,
+      createdBySeatId: 'seat-admin',
+      createdBySeatName: 'Admin Seat',
+    },
+  });
+
+  db.collectionGroup = () => ({
+    where() {
+      throw new Error('collectionGroup lookup should not run when a lookup doc exists');
+    },
+  });
+
+  const result = await joinTeamWithCode({
+    dbClient: db,
+    adminClient: createAdminClient({ customTokenPrefix: 'joined' }),
+    joinCode,
+    seatName: 'Lookup Path Seat',
+  });
+
+  assert.equal(result.teamId, 'team-1');
 });
 
 test('createTeamReport stores bug reports with seat metadata', async () => {