omniai-server/src/betaInviteCodes.js

const fs = require("node:fs");
const path = require("node:path");
const {
  ENTERPRISE_BETA_ACCOUNTS,
  findEnterpriseBetaAccountByInviteCode,
  isEnterpriseBetaInviteCode,
} = require("./enterpriseBetaAccounts");

const DEFAULT_BETA_INVITE_CODES_FILE = path.resolve(__dirname, "../config/internal-beta-codes.md");

function normalizeBetaInviteCode(value) {
  return String(value || "")
    .trim()
    .replace(/[\s-]/g, "")
    .toUpperCase();
}

function loadBetaInviteCodesFromText(text) {
  const codes = new Set();
  const content = String(text || "");
  const pattern = /\b(?:[0-9A-Fa-f]{16}|[0-7]{16,24})\b/g;
  for (const match of content.matchAll(pattern)) {
    const code = normalizeBetaInviteCode(match[0]);
    if (code) codes.add(code);
  }
  return codes;
}

function loadBetaInviteCodes(filePath = process.env.BETA_INVITE_CODES_FILE || DEFAULT_BETA_INVITE_CODES_FILE) {
  const codes = new Set();

  const envCodes = String(process.env.BETA_INVITE_CODES || "")
    .split(/[,\s]+/)
    .map(normalizeBetaInviteCode)
    .filter(Boolean);
  envCodes.forEach((code) => codes.add(code));

  try {
    const markdown = fs.readFileSync(filePath, "utf8");
    loadBetaInviteCodesFromText(markdown).forEach((code) => codes.add(code));
  } catch (error) {
    if (error?.code !== "ENOENT") {
      console.warn("[beta-invite] failed to read invite code file", error);
    }
  }

  for (const account of ENTERPRISE_BETA_ACCOUNTS) {
    codes.add(normalizeBetaInviteCode(account.inviteCode));
  }

  return codes;
}

function getBetaInviteCodeFromBody(body) {
  return normalizeBetaInviteCode(
    body?.betaCode ??
      body?.beta_code ??
      body?.inviteCode ??
      body?.invite_code ??
      body?.internalBetaCode ??
      body?.internal_beta_code,
  );
}

function validateBetaInviteCode(value, allowedCodes = loadBetaInviteCodes()) {
  const code = normalizeBetaInviteCode(value);
  return Boolean(code && allowedCodes.has(code));
}

function validateBetaInviteCodeFromBody(body) {
  return validateBetaInviteCode(getBetaInviteCodeFromBody(body));
}

async function isBetaInviteCodeUsed(value, client) {
  const code = normalizeBetaInviteCode(value);
  if (!code || !client?.query) return false;
  const { rows } = await client.query(
    "SELECT 1 FROM beta_invite_code_uses WHERE code = $1 LIMIT 1",
    [code],
  );
  return rows.length > 0;
}

async function checkBetaInviteCodeForRegistration(value, client, allowedCodes = loadBetaInviteCodes()) {
  const code = normalizeBetaInviteCode(value);
  const enterpriseAccount = findEnterpriseBetaAccountByInviteCode(code);
  if (enterpriseAccount) {
    return {
      ok: true,
      code,
      enterpriseBeta: true,
      account: enterpriseAccount,
    };
  }
  if (!code || !allowedCodes.has(code)) {
    return { ok: false, status: 403, error: "内测码无效或缺失", code };
  }
  if (await isBetaInviteCodeUsed(code, client)) {
    return { ok: false, status: 409, error: "内测码已被使用", code };
  }
  return { ok: true, code };
}

async function consumeBetaInviteCode(value, userId, client, allowedCodes = loadBetaInviteCodes()) {
  const check = await checkBetaInviteCodeForRegistration(value, client, allowedCodes);
  if (!check.ok) return check;
  if (check.enterpriseBeta) return check;

  const { rows } = await client.query(
    `
    INSERT INTO beta_invite_code_uses (code, user_id)
    VALUES ($1, $2)
    ON CONFLICT (code) DO NOTHING
    RETURNING code
  `,
    [check.code, userId || null],
  );

  if (rows.length === 0) {
    return { ok: false, status: 409, error: "内测码已被使用", code: check.code };
  }
  return { ok: true, code: check.code };
}

module.exports = {
  DEFAULT_BETA_INVITE_CODES_FILE,
  normalizeBetaInviteCode,
  loadBetaInviteCodesFromText,
  loadBetaInviteCodes,
  getBetaInviteCodeFromBody,
  findEnterpriseBetaAccountByInviteCode,
  isEnterpriseBetaInviteCode,
  validateBetaInviteCode,
  validateBetaInviteCodeFromBody,
  isBetaInviteCodeUsed,
  checkBetaInviteCodeForRegistration,
  consumeBetaInviteCode,
};