ai-gateway-go/internal/portal/credentials.go

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
	}
	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
}

// 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
	}
	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()
}