package portal import ( "context" "errors" "fmt" "aigateway.local/core/internal/apikey" "aigateway.local/core/internal/platform/cryptox" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" ) type RuntimeCredentials struct { pool *pgxpool.Pool repository *apikey.Repository cipher cryptox.Cipher } func NewRuntimeCredentials(pool *pgxpool.Pool, repository *apikey.Repository, cipher cryptox.Cipher) *RuntimeCredentials { return &RuntimeCredentials{pool: pool, repository: repository, cipher: cipher} } // Ensure returns an internal application credential. The plaintext is only // held for the duration of the request and is never returned to a browser. func (s *RuntimeCredentials) Ensure(ctx context.Context, applicationID string, departmentID *string) (string, string, error) { var encrypted []byte var version int var keyID string err := s.pool.QueryRow(ctx, `SELECT encrypted_key,key_kek_version,api_key_id::text FROM gateway.application_runtime_credentials WHERE application_id=$1 AND department_id IS NOT DISTINCT FROM $2::uuid`, applicationID, departmentID).Scan(&encrypted, &version, &keyID) if err == nil { plain, decryptErr := s.cipher.Decrypt(encrypted, version) return string(plain), keyID, decryptErr } if !errors.Is(err, pgx.ErrNoRows) { return "", "", err } if err := s.checkTenantKeyQuota(ctx, departmentID); err != nil { return "", "", err } record, secret, err := s.repository.Create(ctx, "application-runtime", []string{"application:run"}, 120, 0, 0, nil, "") if err != nil { return "", "", err } encrypted, version, err = s.cipher.Encrypt([]byte(secret)) if err != nil { _, _ = s.repository.Revoke(ctx, record.ID, "") return "", "", err } tx, err := s.pool.Begin(ctx) if err != nil { _, _ = s.repository.Revoke(ctx, record.ID, "") return "", "", err } defer func() { _ = tx.Rollback(ctx) }() _, err = tx.Exec(ctx, `UPDATE gateway.api_keys SET tenant_id=$2,application_id=$3 WHERE id=$1`, record.ID, departmentID, applicationID) if err == nil { _, err = tx.Exec(ctx, `INSERT INTO gateway.application_runtime_credentials(application_id,department_id,api_key_id,encrypted_key,key_kek_version) VALUES($1,$2,$3,$4,$5)`, applicationID, departmentID, record.ID, encrypted, version) } if err != nil { _ = tx.Rollback(ctx) _, _ = s.repository.Revoke(ctx, record.ID, "") // A concurrent request may have won the unique-key race. var pgError *pgconn.PgError if errors.As(err, &pgError) && pgError.Code == "23505" { return s.Ensure(ctx, applicationID, departmentID) } return "", "", fmt.Errorf("store application runtime credential: %w", err) } if err = tx.Commit(ctx); err != nil { _, _ = s.repository.Revoke(ctx, record.ID, "") return "", "", err } return secret, record.ID, nil } // checkTenantKeyQuota 校验租户(部门)Key 配额:max_api_keys>0 且已达上限时拒绝 // 新开通运行时凭据。tenant 为空(未分配部门)不限制。 func (s *RuntimeCredentials) checkTenantKeyQuota(ctx context.Context, tenantID *string) error { if s == nil || s.pool == nil || tenantID == nil || *tenantID == "" { return nil } var maxAPIKeys, used int if err := s.pool.QueryRow(ctx, `SELECT COALESCE(max_api_keys,0) FROM gateway.departments WHERE id=$1`, *tenantID).Scan(&maxAPIKeys); err != nil { return err } if maxAPIKeys <= 0 { return nil } if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM gateway.api_keys WHERE tenant_id=$1 AND enabled`, *tenantID).Scan(&used); err != nil { return err } if used >= maxAPIKeys { return fmt.Errorf("租户 Key 配额已达上限(%d),请联系平台管理员提升配额", maxAPIKeys) } return nil } // EnsureUser returns the portal user's personal runtime credential used by the // general chat. It is created lazily on first use with limits derived from the // user's approved model requests, and bound to the user's department tenant so // usage/audit are attributed to the user's own key. Idempotent: the unique // primary key makes concurrent first-use requests converge on one credential. func (s *RuntimeCredentials) EnsureUser(ctx context.Context, userID string, departmentID *string, rpm int, monthlyTokens int64) (string, string, error) { if s == nil || s.pool == nil || s.repository == nil || s.cipher == nil { return "", "", errors.New("runtime credentials unavailable") } var encrypted []byte var version int var keyID string err := s.pool.QueryRow(ctx, `SELECT encrypted_key,key_kek_version,api_key_id::text FROM gateway.portal_user_runtime_credentials WHERE portal_user_id=$1`, userID).Scan(&encrypted, &version, &keyID) if err == nil { plain, decryptErr := s.cipher.Decrypt(encrypted, version) return string(plain), keyID, decryptErr } if !errors.Is(err, pgx.ErrNoRows) { return "", "", err } // 租户(部门)Key 配额:max_api_keys>0 时校验当前已绑定 Key 数。 if err := s.checkTenantKeyQuota(ctx, departmentID); err != nil { return "", "", err } if rpm < 1 { rpm = 120 } record, secret, err := s.repository.Create(ctx, "portal-chat-runtime", []string{"gateway:invoke"}, rpm, 0, monthlyTokens, nil, "") if err != nil { return "", "", err } encrypted, version, err = s.cipher.Encrypt([]byte(secret)) if err != nil { _, _ = s.repository.Revoke(ctx, record.ID, "") return "", "", err } tx, err := s.pool.Begin(ctx) if err != nil { _, _ = s.repository.Revoke(ctx, record.ID, "") return "", "", err } defer func() { _ = tx.Rollback(ctx) }() _, err = tx.Exec(ctx, `UPDATE gateway.api_keys SET portal_user_id=$2,tenant_id=$3 WHERE id=$1`, record.ID, userID, departmentID) if err == nil { _, err = tx.Exec(ctx, `INSERT INTO gateway.portal_user_runtime_credentials(portal_user_id,api_key_id,encrypted_key,key_kek_version) VALUES($1,$2,$3,$4)`, userID, record.ID, encrypted, version) } if err != nil { _ = tx.Rollback(ctx) _, _ = s.repository.Revoke(ctx, record.ID, "") // A concurrent first-use request may have won the unique-key race. var pgError *pgconn.PgError if errors.As(err, &pgError) && pgError.Code == "23505" { return s.EnsureUser(ctx, userID, departmentID, rpm, monthlyTokens) } return "", "", fmt.Errorf("store portal runtime credential: %w", err) } if err = tx.Commit(ctx); err != nil { _, _ = s.repository.Revoke(ctx, record.ID, "") return "", "", err } return secret, record.ID, nil } // UserSecret returns the portal user's runtime credential plaintext for the // duration of the request. Empty when not yet provisioned. func (s *RuntimeCredentials) UserSecret(ctx context.Context, userID string) (string, string, error) { if s == nil || s.pool == nil || s.cipher == nil { return "", "", errors.New("runtime credentials unavailable") } var encrypted []byte var version int var keyID string err := s.pool.QueryRow(ctx, `SELECT encrypted_key,key_kek_version,api_key_id::text FROM gateway.portal_user_runtime_credentials WHERE portal_user_id=$1`, userID).Scan(&encrypted, &version, &keyID) if errors.Is(err, pgx.ErrNoRows) { return "", "", nil } if err != nil { return "", "", err } plain, err := s.cipher.Decrypt(encrypted, version) return string(plain), keyID, err } func (s *RuntimeCredentials) Metadata(ctx context.Context, applicationID string) ([]map[string]any, error) { rows, err := s.pool.Query(ctx, `SELECT c.api_key_id::text,k.key_prefix,c.department_id::text,k.enabled,c.created_at FROM gateway.application_runtime_credentials c JOIN gateway.api_keys k ON k.id=c.api_key_id WHERE c.application_id=$1 ORDER BY c.created_at`, applicationID) if err != nil { return nil, err } defer rows.Close() items := []map[string]any{} for rows.Next() { var id, prefix string var department *string var enabled bool var created any if err := rows.Scan(&id, &prefix, &department, &enabled, &created); err != nil { return nil, err } items = append(items, map[string]any{"api_key_id": id, "key_prefix": prefix, "department_id": department, "enabled": enabled, "created_at": created}) } return items, rows.Err() }