feat(auth): add user referral system with invite codes

Generate unique referral codes for users Allow new users to register with a referral code Display referral stats and invite link in user profile
2026-05-03 10:27:45 +08:00
parent 3d99f00c75
commit 05898f9441
10 changed files with 422 additions and 14 deletions
--- a/backend/db/schema.sql
+++ b/backend/db/schema.sql
@@ -78,13 +78,28 @@ CREATE TABLE IF NOT EXISTS users (
  email text NOT NULL UNIQUE,
  display_name text NOT NULL,
  password_hash text NOT NULL,
+  referral_code text,
+  referred_by_user_id integer REFERENCES users(id) ON DELETE SET NULL,
  email_verified_at timestamptz,
  created_at timestamptz NOT NULL DEFAULT now(),
  updated_at timestamptz NOT NULL DEFAULT now(),
  CHECK (email = lower(email)),
-  CHECK (length(display_name) BETWEEN 1 AND 40)
+  CHECK (length(display_name) BETWEEN 1 AND 40),
+  CHECK (referral_code IS NULL OR referral_code ~ '^[A-Z0-9]{8,16}$')
 );

+ALTER TABLE users ADD COLUMN IF NOT EXISTS referral_code text;
+ALTER TABLE users ADD COLUMN IF NOT EXISTS referred_by_user_id integer REFERENCES users(id) ON DELETE SET NULL;
+ALTER TABLE users DROP CONSTRAINT IF EXISTS users_referral_code_check;
+ALTER TABLE users ADD CONSTRAINT users_referral_code_check CHECK (referral_code IS NULL OR referral_code ~ '^[A-Z0-9]{8,16}$');
+
+CREATE UNIQUE INDEX IF NOT EXISTS users_referral_code_idx
+  ON users(referral_code)
+  WHERE referral_code IS NOT NULL;
+
+CREATE INDEX IF NOT EXISTS users_referred_by_user_id_idx
+  ON users(referred_by_user_id);
+
 CREATE TABLE IF NOT EXISTS system_wording_keys (
  key text PRIMARY KEY,
  module text NOT NULL,
--- a/backend/src/auth.ts
+++ b/backend/src/auth.ts
@@ -11,6 +11,9 @@ const passwordResetTokenHours = 1;
 const rememberedSessionDays = 30;
 const sessionOnlySessionDays = 1;
 const defaultLocale = 'en';
+const referralCodeLength = 8;
+const referralAlphabet = 'ABCDEFGHJKLMNPQRSTUVWXYZ23456789';
+const referralCodePattern = /^[A-Z0-9]{8,16}$/;

 type DbClient = PoolClient;

@@ -27,6 +30,15 @@ type LoginUserRow = UserRow & {
  password_hash: string;
 };

+type RegistrationUserRow = UserRow & {
+  referral_code: string | null;
+  referred_by_user_id: number | null;
+};
+
+type ReferralCodeRow = QueryResultRow & {
+  referral_code: string | null;
+};
+
 type AuthMessageKey =
  | 'emailRequired'
  | 'invalidEmail'
@@ -42,6 +54,7 @@ type AuthMessageKey =
  | 'invalidCredentials'
  | 'verifyEmailFirst'
  | 'invalidResetToken'
+  | 'invalidReferralCode'
  | 'emailSubject'
  | 'emailHtml'
  | 'emailText'
@@ -58,6 +71,12 @@ export type AuthUser = {
  emailVerified: boolean;
 };

+export type ReferralSummary = {
+  code: string;
+  url: string;
+  verifiedReferralCount: number;
+};
+
 function statusError(message: string, statusCode: number): StatusError {
  const error = new Error(message) as StatusError;
  error.statusCode = statusCode;
@@ -80,6 +99,7 @@ function authMessage(locale: string, key: AuthMessageKey, params: Record<string,
    invalidCredentials: 'server.auth.invalidCredentials',
    verifyEmailFirst: 'server.auth.verifyEmailFirst',
    invalidResetToken: 'server.auth.invalidResetToken',
+    invalidReferralCode: 'server.auth.invalidReferralCode',
    emailSubject: 'email.auth.verificationSubject',
    emailHtml: 'email.auth.verificationHtml',
    emailText: 'email.auth.verificationText',
@@ -137,6 +157,27 @@ async function cleanToken(
  return value.trim();
 }

+async function cleanReferralCode(value: unknown, locale: string): Promise<string | null> {
+  if (value === undefined || value === null) {
+    return null;
+  }
+
+  if (typeof value !== 'string') {
+    throw statusError(await authMessage(locale, 'invalidReferralCode'), 400);
+  }
+
+  const referralCode = value.trim().toUpperCase();
+  if (!referralCode) {
+    return null;
+  }
+
+  if (!referralCodePattern.test(referralCode)) {
+    throw statusError(await authMessage(locale, 'invalidReferralCode'), 400);
+  }
+
+  return referralCode;
+}
+
 function toPublicUser(user: UserRow): AuthUser {
  return {
    id: user.id,
@@ -196,6 +237,80 @@ function hashToken(token: string): string {
  return createHash('sha256').update(token).digest('hex');
 }

+function createReferralCode(): string {
+  const bytes = randomBytes(referralCodeLength);
+  return [...bytes].map((byte) => referralAlphabet[byte % referralAlphabet.length]).join('');
+}
+
+function isUniqueViolation(error: unknown): boolean {
+  return typeof error === 'object' && error !== null && 'code' in error && (error as { code?: unknown }).code === '23505';
+}
+
+async function ensureReferralCode(client: DbClient, userId: number): Promise<string> {
+  const existing = await clientQueryOne<ReferralCodeRow>(client, 'SELECT referral_code FROM users WHERE id = $1', [userId]);
+  if (existing?.referral_code) {
+    return existing.referral_code;
+  }
+
+  for (let attempt = 0; attempt < 10; attempt += 1) {
+    const referralCode = createReferralCode();
+    try {
+      const updated = await clientQueryOne<ReferralCodeRow>(
+        client,
+        `
+          UPDATE users
+          SET referral_code = $1, updated_at = now()
+          WHERE id = $2
+            AND referral_code IS NULL
+          RETURNING referral_code
+        `,
+        [referralCode, userId]
+      );
+
+      if (updated?.referral_code) {
+        return updated.referral_code;
+      }
+
+      const current = await clientQueryOne<ReferralCodeRow>(client, 'SELECT referral_code FROM users WHERE id = $1', [
+        userId
+      ]);
+      if (current?.referral_code) {
+        return current.referral_code;
+      }
+    } catch (error) {
+      if (!isUniqueViolation(error)) {
+        throw error;
+      }
+    }
+  }
+
+  throw new Error('Failed to assign referral code');
+}
+
+async function referralUserId(
+  client: DbClient,
+  referralCode: string,
+  currentUserId: number | null,
+  locale: string
+): Promise<number> {
+  const row = await clientQueryOne<QueryResultRow & { id: number }>(client, 'SELECT id FROM users WHERE referral_code = $1', [
+    referralCode
+  ]);
+
+  if (!row || (currentUserId !== null && row.id === currentUserId)) {
+    throw statusError(await authMessage(locale, 'invalidReferralCode'), 400);
+  }
+
+  return row.id;
+}
+
+function buildReferralUrl(code: string): string {
+  const origin = process.env.APP_ORIGIN ?? process.env.FRONTEND_ORIGIN ?? 'http://localhost:3000';
+  const url = new URL('/register', origin);
+  url.searchParams.set('ref', code);
+  return url.toString();
+}
+
 function getEmailConfig() {
  const apiKey = process.env.RESEND_API_KEY;
  const from = process.env.EMAIL_FROM;
@@ -280,14 +395,15 @@ export async function registerUser(payload: Record<string, unknown>, locale = de
  const email = await cleanEmail(payload.email, locale);
  const displayName = await cleanDisplayName(payload.displayName, locale);
  const password = await cleanPassword(payload.password, locale);
+  const referralCode = await cleanReferralCode(payload.referralCode, locale);
  const passwordHash = await hashPassword(password);
  const verificationToken = createPlainToken();
  const verificationTokenHash = hashToken(verificationToken);

  await withTransaction(async (client) => {
-    const existingUser = await clientQueryOne<UserRow>(
+    const existingUser = await clientQueryOne<RegistrationUserRow>(
      client,
-      'SELECT id, email, display_name, email_verified_at FROM users WHERE email = $1',
+      'SELECT id, email, display_name, referral_code, referred_by_user_id, email_verified_at FROM users WHERE email = $1',
      [email]
    );

@@ -295,23 +411,24 @@ export async function registerUser(payload: Record<string, unknown>, locale = de
      throw statusError(await authMessage(locale, 'emailAlreadyRegistered'), 409);
    }

+    const referrerUserId = referralCode ? await referralUserId(client, referralCode, existingUser?.id ?? null, locale) : null;
    const user = existingUser
-      ? await clientQueryOne<UserRow>(
+      ? await clientQueryOne<RegistrationUserRow>(
          client,
          `
            UPDATE users
            SET display_name = $1, password_hash = $2, updated_at = now()
            WHERE id = $3
-            RETURNING id, email, display_name, email_verified_at
+            RETURNING id, email, display_name, referral_code, referred_by_user_id, email_verified_at
          `,
          [displayName, passwordHash, existingUser.id]
        )
-      : await clientQueryOne<UserRow>(
+      : await clientQueryOne<RegistrationUserRow>(
          client,
          `
            INSERT INTO users (email, display_name, password_hash)
            VALUES ($1, $2, $3)
-            RETURNING id, email, display_name, email_verified_at
+            RETURNING id, email, display_name, referral_code, referred_by_user_id, email_verified_at
          `,
          [email, displayName, passwordHash]
        );
@@ -320,6 +437,14 @@ export async function registerUser(payload: Record<string, unknown>, locale = de
      throw new Error('Failed to save user');
    }

+    await ensureReferralCode(client, user.id);
+    if (referrerUserId && user.referred_by_user_id === null) {
+      await client.query('UPDATE users SET referred_by_user_id = $1, updated_at = now() WHERE id = $2', [
+        referrerUserId,
+        user.id
+      ]);
+    }
+
    await client.query('DELETE FROM email_verification_tokens WHERE user_id = $1 AND used_at IS NULL', [user.id]);
    await client.query(
      `
@@ -530,6 +655,23 @@ export async function updateCurrentUser(
  return toPublicUser(user);
 }

+export async function getReferralSummary(userId: number): Promise<ReferralSummary> {
+  return withTransaction(async (client) => {
+    const code = await ensureReferralCode(client, userId);
+    const countRow = await clientQueryOne<QueryResultRow & { count: string }>(
+      client,
+      'SELECT COUNT(*)::text AS count FROM users WHERE referred_by_user_id = $1 AND email_verified_at IS NOT NULL',
+      [userId]
+    );
+
+    return {
+      code,
+      url: buildReferralUrl(code),
+      verifiedReferralCount: Number(countRow?.count ?? 0)
+    };
+  });
+}
+
 export async function logoutSession(token: string): Promise<void> {
  if (token.length < 32) {
    return;
--- a/backend/src/server.ts
+++ b/backend/src/server.ts
@@ -5,6 +5,7 @@ import Fastify from 'fastify';
 import type { FastifyReply, FastifyRequest } from 'fastify';
 import { mkdir } from 'node:fs/promises';
 import {
+  getReferralSummary,
  getUserBySessionToken,
  loginUser,
  logoutSession,
@@ -239,6 +240,17 @@ app.patch('/api/auth/me', async (request, reply) => {
  return { user: await updateCurrentUser(user.id, payload, requestLocale(request)) };
 });

+app.get('/api/auth/referral', async (request, reply) => {
+  const token = getBearerToken(request.headers.authorization);
+  const user = token ? await getUserBySessionToken(token) : null;
+
+  if (!user) {
+    return reply.code(401).send({ message: await serverMessage(requestLocale(request), 'loginRequired') });
+  }
+
+  return { referral: await getReferralSummary(user.id) };
+});
+
 app.post('/api/auth/logout', async (request, reply) => {
  const token = getBearerToken(request.headers.authorization);
  if (token) {