ai-drama-platform/server/saml.mjs

import { createHash, randomBytes } from "node:crypto";
import { SAML } from "@node-saml/node-saml";
import { dbGet, dbRun } from "./db.mjs";
import { resolveOidcUser } from "./oidc.mjs";

const LOGIN_STATE_TTL_MS = 10 * 60 * 1000;
const CLOCK_SKEW_MS = 5000;
const DEFAULT_NAME_ID_FORMAT = "urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress";

function samlError(status, code, message, details = {}) {
  const error = new Error(message);
  error.status = status;
  error.code = code;
  error.details = details;
  return error;
}

function nowIso() {
  return new Date().toISOString();
}

function hashValue(value) {
  return createHash("sha256").update(String(value || "")).digest("hex");
}

function base64url(buffer) {
  return Buffer.from(buffer).toString("base64url");
}

function parseJson(value, fallback) {
  try { return JSON.parse(value); } catch { return fallback; }
}

function ensureSsoEnabled() {
  const policy = dbGet("SELECT sso_enabled FROM identity_policies WHERE id = 'default'");
  if (!policy?.sso_enabled) throw samlError(403, "sso_disabled", "平台当前未启用企业 SSO");
}

function providerRow(providerId, { requireSso = true, allowDisabled = false } = {}) {
  if (requireSso) ensureSsoEnabled();
  const provider = dbGet(`SELECT * FROM identity_providers WHERE id = ?${allowDisabled ? "" : " AND enabled = 1"}`, [providerId]);
  if (!provider) throw samlError(404, "sso_provider_not_found", "企业身份提供商不存在或未启用");
  if (provider.kind !== "saml") throw samlError(400, "sso_provider_kind_unsupported", "当前登录链路不是 SAML 提供商");
  if (requireSso && (!provider.entry_point || !provider.idp_cert_ref || !provider.sp_issuer)) {
    throw samlError(400, "saml_provider_incomplete", "SAML 提供商缺少 IdP Entry Point、证书环境变量名或 SP Issuer");
  }
  if (!["never", "ifPresent", "always"].includes(provider.validate_in_response_to || "ifPresent")) {
    throw samlError(400, "saml_validate_in_response_to_invalid", "SAML InResponseTo 校验策略无效");
  }
  if (!allowDisabled && !["configured", "ready"].includes(provider.status)) {
    throw samlError(400, "saml_provider_not_ready", "SAML 提供商尚未完成探测或已停用");
  }
  if (allowDisabled && !["disabled", "configured", "ready"].includes(provider.status)) {
    throw samlError(400, "saml_provider_not_ready", "SAML 提供商尚未完成探测或已停用");
  }
  return provider;
}

function envCertificate(ref) {
  const key = String(ref || "").trim();
  if (!/^[A-Za-z_][A-Za-z0-9_]*$/.test(key)) {
    throw samlError(400, "saml_certificate_ref_invalid", "SAML IdP 证书环境变量名无效");
  }
  const value = String(process.env[key] || "").trim().replace(/\\n/g, "\n");
  if (!value) throw samlError(503, "saml_certificate_missing", `服务端环境变量 ${key} 未配置 IdP 证书`);
  return value;
}

function normalizeValidateInResponseTo(value) {
  return ["never", "ifPresent", "always"].includes(value) ? value : "ifPresent";
}

function providerOptions(provider, { callbackUrl, cacheProvider, requireIdpCert = true }) {
  const issuer = String(provider.sp_issuer || process.env.AI_DRAMA_SAML_SP_ISSUER || `${process.env.AI_DRAMA_API_ORIGIN || "http://127.0.0.1:8787"}/api/auth/sso/saml/metadata`).trim();
  return {
    entryPoint: provider.entry_point,
    idpCert: requireIdpCert ? envCertificate(provider.idp_cert_ref) : "AA==",
    issuer,
    callbackUrl,
    audience: provider.audience || issuer,
    idpIssuer: provider.issuer_url || undefined,
    identifierFormat: provider.saml_name_id_format || DEFAULT_NAME_ID_FORMAT,
    wantAssertionsSigned: provider.want_assertions_signed !== 0,
    wantAuthnResponseSigned: provider.want_authn_response_signed !== 0,
    validateInResponseTo: normalizeValidateInResponseTo(provider.validate_in_response_to),
    acceptedClockSkewMs: CLOCK_SKEW_MS,
    requestIdExpirationPeriodMs: LOGIN_STATE_TTL_MS,
    cacheProvider,
    disableRequestedAuthnContext: true,
    signatureAlgorithm: "sha256",
    providerName: "AI 短剧生产平台"
  };
}

class DatabaseSamlCacheProvider {
  constructor(relayStateHash) {
    this.relayStateHash = relayStateHash;
  }

  async saveAsync(key, value) {
    const result = dbRun(
      "UPDATE saml_login_states SET request_id = ?, request_issue_instant = ? WHERE relay_state_hash = ? AND consumed_at IS NULL AND expires_at > ?",
      [String(key), String(value), this.relayStateHash, nowIso()]
    );
    if (!Number(result.changes || 0)) throw samlError(401, "saml_state_invalid", "SAML 登录状态不存在或已过期");
    return { value: String(value), createdAt: Date.now() };
  }

  async getAsync(key) {
    const row = dbGet(
      "SELECT request_issue_instant FROM saml_login_states WHERE relay_state_hash = ? AND request_id = ? AND consumed_at IS NULL AND expires_at > ?",
      [this.relayStateHash, String(key), nowIso()]
    );
    return row?.request_issue_instant || null;
  }

  // node-saml calls removeAsync during assertion processing. Final consumption is
  // handled atomically after the whole signed response has passed validation.
  async removeAsync(key) {
    const row = dbGet("SELECT request_issue_instant FROM saml_login_states WHERE relay_state_hash = ? AND request_id = ?", [this.relayStateHash, String(key)]);
    return row?.request_issue_instant || null;
  }
}

function cleanupExpiredSamlStates() {
  const timestamp = nowIso();
  dbRun("DELETE FROM saml_login_states WHERE expires_at <= ? OR consumed_at IS NOT NULL", [timestamp]);
}

function profileToClaims(profile, provider) {
  const claims = { ...(profile?.attributes || {}), ...(profile || {}) };
  const first = (...values) => values.flatMap((value) => Array.isArray(value) ? value : [value]).find((value) => typeof value === "string" && value.trim()) || "";
  const nameId = first(profile?.nameID, profile?.nameId, profile?.uid, profile?.subject);
  const email = first(profile?.email, profile?.mail, profile?.userPrincipalName, profile?.preferred_username, claims.email, claims.mail, claims["urn:oid:0.9.2342.19200300.100.1.3"]);
  const displayName = first(profile?.displayName, profile?.name, profile?.cn, profile?.givenName, claims.displayName, claims.name, email, nameId);
  claims.sub = first(claims.sub, nameId);
  claims.email = first(claims.email, email);
  claims.name = first(claims.name, displayName);
  claims.iss = first(claims.iss, profile?.issuer, provider.issuer_url);
  return claims;
}

export function isSamlProvider(providerId) {
  return Boolean(dbGet("SELECT id FROM identity_providers WHERE id = ? AND kind = 'saml'", [providerId]));
}

export async function startSamlLogin(providerId, { callbackUrl, apiOrigin, returnTo = "/", selection = {}, ipAddress = "", userAgent = "", host = "" } = {}) {
  cleanupExpiredSamlStates();
  const provider = providerRow(providerId);
  const relayState = base64url(randomBytes(32));
  const timestamp = nowIso();
  const expiresAt = new Date(Date.now() + LOGIN_STATE_TTL_MS).toISOString();
  const id = `saml-state-${Date.now()}-${randomBytes(4).toString("hex")}`;
  dbRun(
    "INSERT INTO saml_login_states(id, relay_state_hash, provider_id, return_to, selection_json, ip_address, user_agent, expires_at, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
    [id, hashValue(relayState), provider.id, String(returnTo || "/"), JSON.stringify(selection || {}), String(ipAddress || ""), String(userAgent || ""), expiresAt, timestamp]
  );
  try {
    const cacheProvider = new DatabaseSamlCacheProvider(hashValue(relayState));
    const saml = new SAML(providerOptions(provider, { callbackUrl, cacheProvider }));
    const authorizationUrl = await saml.getAuthorizeUrlAsync(relayState, host || apiOrigin, {});
    return { authorizationUrl, relayState, expiresAt, provider };
  } catch (error) {
    dbRun("DELETE FROM saml_login_states WHERE relay_state_hash = ?", [hashValue(relayState)]);
    if (error.status) throw error;
    throw samlError(502, "saml_request_failed", `SAML 登录请求生成失败：${error.message}`);
  }
}

export async function handleSamlCallback({ samlResponse, relayState, callbackUrl } = {}) {
  cleanupExpiredSamlStates();
  if (!samlResponse || !relayState) throw samlError(401, "saml_response_state_required", "SAML ACS 必须包含断言和 RelayState");
  const stateRow = dbGet(
    "SELECT * FROM saml_login_states WHERE relay_state_hash = ? AND consumed_at IS NULL AND expires_at > ?",
    [hashValue(relayState), nowIso()]
  );
  if (!stateRow) throw samlError(401, "saml_state_invalid", "SAML 登录状态无效、已使用或已过期");
  const provider = providerRow(stateRow.provider_id);
  const cacheProvider = new DatabaseSamlCacheProvider(hashValue(relayState));
  const saml = new SAML(providerOptions(provider, { callbackUrl, cacheProvider }));
  let result;
  try {
    result = await saml.validatePostResponseAsync({ SAMLResponse: String(samlResponse) });
  } catch (error) {
    throw samlError(401, "saml_response_invalid", `SAML 断言校验失败：${error.message}`);
  }
  if (result.loggedOut || !result.profile) throw samlError(401, "saml_profile_missing", "SAML 响应没有可用的登录身份");
  const consumed = dbRun("UPDATE saml_login_states SET consumed_at = ? WHERE id = ? AND consumed_at IS NULL AND expires_at > ?", [nowIso(), stateRow.id, nowIso()]);
  if (!Number(consumed.changes || 0)) throw samlError(401, "saml_state_replayed", "SAML 登录状态已被使用");
  const claims = profileToClaims(result.profile, provider);
  return {
    ...resolveOidcUser(provider, claims),
    selection: parseJson(stateRow.selection_json, {}),
    returnTo: stateRow.return_to,
    provider,
    profile: result.profile
  };
}

export function samlServiceProviderMetadata(providerId, { callbackUrl } = {}) {
  const provider = providerRow(providerId, { requireSso: false, allowDisabled: true });
  const cacheProvider = new DatabaseSamlCacheProvider("metadata");
  const saml = new SAML(providerOptions(provider, { callbackUrl, cacheProvider, requireIdpCert: false }));
  return saml.generateServiceProviderMetadata(null, null);
}