feat: bring back referrals PUT-1463 (#3599)

extensions/whoami.ts

@@ -160,6 +160,10 @@ export const handleWhoami = async (
         desktop_bg_fit: user.desktop_bg_fit,
         is_temp: user.password === null && user.email === null,
         is_user_token: true,
+        // Present only once the account has actually asked for a code (see the
+        // referral extension) — null until then, and never minted from here:
+        // this endpoint is polled, and a mint is a write.
+        referral_code: user.referral_code,
         oidc_only: oidcOnly,
         taskbar_items: isUser
             ? await getTaskbarItems(
@@ -240,6 +244,8 @@ export const handleWhoami = async (
         delete details.created_ts;
         delete details.is_user_token;
         delete details.metadata;
+        // An app has no business reading the code its user earns credit with.
+        delete details.referral_code;
     }
 
     if (actor.app) {
@@ -257,7 +263,8 @@ export const handleWhoami = async (
     }
 
     const subscription = details.subscription as
-        { offering?: Record<string, unknown> } | undefined;
+        | { offering?: Record<string, unknown> }
+        | undefined;
     if (subscription?.offering) {
         delete subscription.offering.group;
         delete subscription.offering.benefits;

src/backend/clients/database/SqliteDatabaseClient.test.ts

@@ -27,7 +27,7 @@ import { DatabaseClientFactory } from './index.js';
 import { SqliteDatabaseClient } from './SqliteDatabaseClient.js';
 
 /** Highest schema version the migration table can reach. */
-const CURRENT_SCHEMA_VERSION = 63;
+const CURRENT_SCHEMA_VERSION = 64;
 const SYSTEM_USER_UUID = '5d4adce0-a381-4982-9c02-6e2540026238';
 
 const sqliteConfig = (

src/backend/clients/database/SqliteDatabaseClient.ts

@@ -97,6 +97,7 @@ const AVAILABLE_MIGRATIONS: [number, string[]][] = [
     [60, ['0065_app-feedback.sql']],
     [61, ['0066_owned-email-unique.sql']],
     [62, ['0067_share_entries.sql']],
+    [63, ['0068_referral-code-unique.sql']],
 ];
 
 export class SqliteDatabaseClient extends AbstractDatabaseClient {

src/backend/clients/database/migrations/sqlite/0068_referral-code-unique.sql

@@ -0,0 +1,32 @@
+-- Copyright (C) 2024-present Puter Technologies Inc.
+--
+-- This file is part of Puter.
+--
+-- Puter is free software: you can redistribute it and/or modify
+-- it under the terms of the GNU Affero General Public License as published
+-- by the Free Software Foundation, either version 3 of the License, or
+-- (at your option) any later version.
+--
+-- This program is distributed in the hope that it will be useful,
+-- but WITHOUT ANY WARRANTY; without even the implied warranty of
+-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+-- GNU Affero General Public License for more details.
+--
+-- You should have received a copy of the GNU Affero General Public License
+-- along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+-- One referral code, one account. The mysql and postgres baselines have
+-- declared `user.referral_code` UNIQUE since the column was introduced; the
+-- sqlite baseline only declared the column, so the invariant held on hosted
+-- deployments and not on self-hosted ones.
+--
+-- Codes are minted by picking at random and writing, so this index is not a
+-- nicety: it is what makes "did anyone already take this code" a decision the
+-- database makes once, instead of a read the next mint can race. A caller that
+-- loses the race sees a unique violation and picks again.
+--
+-- NULL is distinct from NULL in a sqlite unique index, so accounts that have
+-- never asked for a code (the majority) are unaffected.
+CREATE UNIQUE INDEX IF NOT EXISTS idx_user_referral_code
+    ON user(referral_code)
+    WHERE referral_code IS NOT NULL;

src/backend/stores/user/UserStore.test.ts

@@ -532,6 +532,100 @@ describe('UserStore batched and uncached lookups', () => {
         ).resolves.toBeUndefined();
     });
 
+    describe('referral codes', () => {
+        it('claims a code once and resolves it back to the account', async () => {
+            const user = await makeUser();
+            expect(user.referral_code ?? null).toBeNull();
+
+            expect(
+                await server.stores.user.claimReferralCode(user.id, 'ABCD1234'),
+            ).toBe(true);
+            expect(
+                (await server.stores.user.getById(user.id))?.referral_code,
+            ).toBe('ABCD1234');
+
+            const resolved =
+                await server.stores.user.getByReferralCode('ABCD1234');
+            expect(resolved?.id).toBe(user.id);
+        });
+
+        it('refuses a second claim, so a shared code cannot be replaced', async () => {
+            const user = await makeUser();
+            await server.stores.user.claimReferralCode(user.id, 'FIRST123');
+
+            expect(
+                await server.stores.user.claimReferralCode(user.id, 'SECOND12'),
+            ).toBe(false);
+            expect(
+                (await server.stores.user.getById(user.id))?.referral_code,
+            ).toBe('FIRST123');
+            expect(
+                await server.stores.user.getByReferralCode('SECOND12'),
+            ).toBeNull();
+        });
+
+        it('lets the database reject a code another account already holds', async () => {
+            const first = await makeUser();
+            const second = await makeUser();
+            await server.stores.user.claimReferralCode(first.id, 'TAKEN123');
+
+            await expect(
+                server.stores.user.claimReferralCode(second.id, 'TAKEN123'),
+            ).rejects.toMatchObject({
+                code: expect.stringContaining('SQLITE_CONSTRAINT'),
+            });
+            expect(
+                (await server.stores.user.getByReferralCode('TAKEN123'))?.id,
+            ).toBe(first.id);
+        });
+
+        it('resolves a code however it was typed, and rejects non-codes', async () => {
+            const user = await makeUser();
+            await server.stores.user.claimReferralCode(user.id, 'MIXED123');
+
+            for (const typed of ['mixed123', 'MiXeD123', '  mixed123  ']) {
+                expect(
+                    (await server.stores.user.getByReferralCode(typed))?.id,
+                ).toBe(user.id);
+            }
+            // Never reaches the DB, so a code-shaped injection can't either.
+            expect(
+                await server.stores.user.getByReferralCode("' OR 1=1 --"),
+            ).toBeNull();
+        });
+
+        it('stops resolving a code the account no longer holds', async () => {
+            const user = await makeUser();
+            await server.stores.user.claimReferralCode(user.id, 'OLDCODE1');
+            // Warm the cache under the code's key.
+            expect(
+                (await server.stores.user.getByReferralCode('OLDCODE1'))?.id,
+            ).toBe(user.id);
+
+            await server.stores.user.update(user.id, {
+                referral_code: 'NEWCODE1',
+            });
+
+            expect(
+                await server.stores.user.getByReferralCode('OLDCODE1'),
+            ).toBeNull();
+            expect(
+                (await server.stores.user.getByReferralCode('NEWCODE1'))?.id,
+            ).toBe(user.id);
+        });
+
+        it('allows any number of accounts without a code', async () => {
+            const first = await makeUser();
+            const second = await makeUser();
+            expect(
+                (await server.stores.user.getById(first.id))?.referral_code,
+            ).toBeNull();
+            expect(
+                (await server.stores.user.getById(second.id))?.referral_code,
+            ).toBeNull();
+        });
+    });
+
     it('normalizes a corrupt metadata column to an empty object', async () => {
         const user = await makeUser();
         await server.clients.db.write(

src/backend/stores/user/UserStore.ts

@@ -18,6 +18,7 @@
  */
 
 import { cleanEmail } from '../../util/email.js';
+import { normalizeReferralCode } from '../../util/referralCode.js';
 import { PuterStore } from '../types';
 
 // -- Types ------------------------------------------------------------
@@ -53,6 +54,14 @@ export interface UserRow {
     reputation?: number;
     /** E.164 phone number collected during SMS verification. */
     phone?: string | null;
+    /**
+     * The account's own referral code — unique across accounts, shareable, and
+     * minted on demand rather than at signup, so accounts that never look at
+     * theirs never get one. Null for every account that hasn't asked.
+     */
+    referral_code?: string | null;
+    /** `user.id` of the account whose referral code this one signed up with. */
+    referred_by?: number | null;
     /** True while the account must complete SMS phone verification before use. */
     requires_phone_verification?: boolean;
     /** True while the account must complete credit-card verification before use. */
@@ -73,7 +82,13 @@ export interface UserRow {
  * is as simple as adding a key here — lookups + cache fan-out follow
  * automatically.
  */
-export const USER_ID_PROPERTIES = ['id', 'uuid', 'username', 'email'] as const;
+export const USER_ID_PROPERTIES = [
+    'id',
+    'uuid',
+    'username',
+    'email',
+    'referral_code',
+] as const;
 export type UserIdProperty = (typeof USER_ID_PROPERTIES)[number];
 
 // -- Constants --------------------------------------------------------
@@ -207,6 +222,24 @@ export class UserStore extends PuterStore {
         return this.getByProperty('email', email, opts);
     }
 
+    /**
+     * Resolve a referral code to the account that owns it.
+     *
+     * Codes are stored canonically (upper case), so the lookup normalizes
+     * first: on a case-sensitive collation a typed-in lower-case code would
+     * otherwise miss the row, and a case-varying code would resolve to the same
+     * account under two different cache keys. A value outside the stored shape
+     * can't match any row, so it never reaches the DB.
+     */
+    async getByReferralCode(
+        code: string,
+        opts: { cached?: boolean; force?: boolean } = {},
+    ): Promise<UserRow | null> {
+        const normalized = normalizeReferralCode(code);
+        if (!normalized) return null;
+        return this.getByProperty('referral_code', normalized, opts);
+    }
+
     /**
      * Batched lookup by id. Dedupes input ids, reads cache via a pipelined
      * MGET, and resolves remaining misses with a single `SELECT … WHERE id IN
@@ -565,9 +598,33 @@ export class UserStore extends PuterStore {
     }
 
     /**
-     * Shared write path for `update` / `claimPlaceholder`. `guard` is extra SQL
-     * ANDed into the WHERE clause; the write is reported as lost when it
-     * matches no row.
+     * Assign a referral code to an account that doesn't have one yet.
+     *
+     * Two guards, for two different races:
+     *
+     * - `referral_code IS NULL` — two concurrent requests for one account's code
+     *   each mint a candidate, and an unguarded write would let the second
+     *   replace the first. The account's code is a value it shares with other
+     *   people, so it has to be write-once: a code that changes after being
+     *   shared silently stops crediting anyone. Returns false to the loser,
+     *   which then reads the winner's code.
+     * - The unique index on the column, which surfaces as a thrown unique
+     *   violation when two accounts happen to mint the same code. Callers
+     *   distinguish the two: false means "someone already set yours", a throw
+     *   means "pick a different code and retry".
+     */
+    async claimReferralCode(userId: number, code: string): Promise<boolean> {
+        return this.#write(
+            userId,
+            { referral_code: code },
+            '`referral_code` IS NULL',
+        );
+    }
+
+    /**
+     * Shared write path for `update` / `claimPlaceholder` /
+     * `claimReferralCode`. `guard` is extra SQL ANDed into the WHERE clause;
+     * the write is reported as lost when it matches no row.
      */
     async #write(
         userId: number,
@@ -823,7 +880,6 @@ export class UserStore extends PuterStore {
      * strings on SQLite, parsed objects on MySQL.
      */
     #normalizeRow(row: Record<string, unknown>): UserRow {
-        const { referral_code: _referralCode, ...rest } = row;
         const asBool = (v: unknown): boolean | undefined => {
             if (v === null || v === undefined) return undefined;
             if (typeof v === 'boolean') return v;
@@ -848,23 +904,27 @@ export class UserStore extends PuterStore {
         })();
 
         return {
-            ...rest,
-            id: Number(rest.id),
-            uuid: String(rest.uuid),
-            username: String(rest.username),
-            email: rest.email == null ? null : String(rest.email),
-            suspended: asBool(rest.suspended),
-            email_confirmed: asBool(rest.email_confirmed),
+            ...row,
+            id: Number(row.id),
+            uuid: String(row.uuid),
+            username: String(row.username),
+            email: row.email == null ? null : String(row.email),
+            suspended: asBool(row.suspended),
+            email_confirmed: asBool(row.email_confirmed),
             requires_email_confirmation: asBool(
-                rest.requires_email_confirmation,
+                row.requires_email_confirmation,
             ),
             requires_phone_verification: asBool(
-                rest.requires_phone_verification,
+                row.requires_phone_verification,
             ),
-            requires_card_verification: asBool(rest.requires_card_verification),
-            phone: rest.phone == null ? null : String(rest.phone),
+            requires_card_verification: asBool(row.requires_card_verification),
+            phone: row.phone == null ? null : String(row.phone),
             reputation:
-                rest.reputation == null ? undefined : Number(rest.reputation),
+                row.reputation == null ? undefined : Number(row.reputation),
+            // Normalized on the way out so a code stored before the canonical
+            // form was enforced still reads back as the canonical one, and the
+            // cache key fanned out for it matches what a lookup asks for.
+            referral_code: normalizeReferralCode(row.referral_code),
             metadata,
         };
     }

src/backend/util/referralCode.test.ts

@@ -0,0 +1,70 @@
+import { describe, expect, test } from 'vitest';
+
+import {
+    REFERRAL_CODE_LENGTH,
+    REFERRAL_CODE_MAX_LENGTH,
+    generateReferralCode,
+    normalizeReferralCode,
+} from './referralCode.js';
+
+describe('normalizeReferralCode', () => {
+    test.for([
+        ['already canonical', 'ABCD1234', 'ABCD1234'],
+        ['lower case', 'abcd1234', 'ABCD1234'],
+        ['mixed case', 'AbCd1234', 'ABCD1234'],
+        ['surrounding whitespace', '  abcd1234\n', 'ABCD1234'],
+        ['shortest accepted', 'ab12', 'AB12'],
+        ['longest accepted', 'a'.repeat(16), 'A'.repeat(16)],
+    ] as const)('%s → %s', ([, input, expected]) => {
+        expect(normalizeReferralCode(input)).toBe(expected);
+    });
+
+    test.for([
+        ['too short', 'abc'],
+        ['longer than the column', 'a'.repeat(REFERRAL_CODE_MAX_LENGTH + 1)],
+        ['inner whitespace', 'ab cd'],
+        ['a dot', 'ab.cd'],
+        ['a colon', 'abuse:referral:x'],
+        ['a hyphen', 'ab-cd'],
+        ['non-ASCII', 'абвг1234'],
+        ['empty', ''],
+        ['only whitespace', '   '],
+    ] as const)('%s is rejected', ([, input]) => {
+        expect(normalizeReferralCode(input)).toBeNull();
+    });
+
+    test('non-strings are rejected', () => {
+        expect(normalizeReferralCode(undefined)).toBeNull();
+        expect(normalizeReferralCode(null)).toBeNull();
+        expect(normalizeReferralCode(12345678)).toBeNull();
+        expect(normalizeReferralCode({ code: 'abcd1234' })).toBeNull();
+    });
+});
+
+describe('generateReferralCode', () => {
+    test('mints canonical codes of the expected length', () => {
+        for (let i = 0; i < 200; i++) {
+            const code = generateReferralCode();
+            expect(code).toHaveLength(REFERRAL_CODE_LENGTH);
+            expect(normalizeReferralCode(code)).toBe(code);
+        }
+    });
+
+    test('leaves out the characters that read as each other', () => {
+        // 200 8-char codes is 1,600 draws; a 32-letter alphabet would show any
+        // included letter long before that.
+        const drawn = new Set(
+            Array.from({ length: 200 }, () => generateReferralCode()).join(''),
+        );
+        for (const ambiguous of ['I', 'L', 'O', 'U']) {
+            expect(drawn.has(ambiguous)).toBe(false);
+        }
+    });
+
+    test('is not derivable — successive codes differ', () => {
+        const codes = new Set(
+            Array.from({ length: 500 }, () => generateReferralCode()),
+        );
+        expect(codes.size).toBe(500);
+    });
+});

src/backend/util/referralCode.ts

@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2024-present Puter Technologies Inc.
+ *
+ * This file is part of Puter.
+ *
+ * Puter is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+/**
+ * Shape and minting of `user.referral_code`.
+ *
+ * A referral code is a public, shareable handle for one account, so the only
+ * mechanism here is "what a code may look like" and "produce a fresh one". What
+ * a redemption is worth, and when a pattern of redemptions is abuse, is decided
+ * elsewhere.
+ *
+ * Codes have exactly ONE canonical form — upper case — and every boundary
+ * normalizes to it before touching a database or a counter. Two reasons:
+ *
+ * - Collations disagree. MySQL's `latin1_swedish_ci` matches a code
+ *   case-insensitively, SQLite and Postgres do not, so `ab12cd34` resolving to
+ *   an account would depend on the engine.
+ * - Anything that counts redemptions per code keys off the code string. Left
+ *   un-normalized, `AB12CD34` and `Ab12cd34` are one code to the lookup and two
+ *   to the counters, which is a free way around any per-code cap.
+ */
+import crypto from 'crypto';
+
+/** `user.referral_code` is `varchar(16)`; a longer value can't be stored. */
+export const REFERRAL_CODE_MAX_LENGTH = 16;
+
+/** Length of a newly minted code. 32^8 ≈ 1.1e12 possibilities. */
+export const REFERRAL_CODE_LENGTH = 8;
+
+/**
+ * Accepted canonical shape. Deliberately wider than the mint alphabet below:
+ * codes minted by earlier versions of the program drew from the full
+ * alphanumeric set, and they must keep resolving.
+ */
+export const REFERRAL_CODE_SHAPE = new RegExp(
+    `^[A-Z0-9]{4,${REFERRAL_CODE_MAX_LENGTH}}$`,
+);
+
+/**
+ * Crockford base32 — the alphanumerics minus `I`, `L`, `O` and `U`. Dropping
+ * the letters that read as digits keeps a code transcribable from a screen or
+ * over the phone, and dropping `U` keeps most accidental profanity out of a
+ * code users are asked to share.
+ */
+const MINT_ALPHABET = '0123456789ABCDEFGHJKMNPQRSTVWXYZ';
+
+/**
+ * The canonical form of a client- or DB-supplied code, or null when the value
+ * can't be a code at all. Everything that looks a code up, stores one, or
+ * counts against one goes through here first.
+ */
+export const normalizeReferralCode = (value: unknown): string | null => {
+    if (typeof value !== 'string') return null;
+    const normalized = value.trim().toUpperCase();
+    return REFERRAL_CODE_SHAPE.test(normalized) ? normalized : null;
+};
+
+/**
+ * Mint a fresh code. Uniqueness is not this function's job — it is held by the
+ * unique index on `user.referral_code`, and callers retry on violation.
+ *
+ * `crypto.randomInt` rather than `Math.random`: a code that can be derived from
+ * an account id (as an earlier seeded-RNG version could) lets anyone
+ * reconstruct another user's code and spend their referral allowance for them.
+ */
+export const generateReferralCode = (length = REFERRAL_CODE_LENGTH): string => {
+    let code = '';
+    for (let i = 0; i < length; i++) {
+        code += MINT_ALPHABET[crypto.randomInt(0, MINT_ALPHABET.length)];
+    }
+    return code;
+};