ai-gateway-go/internal/workbench/envvars.go

package workbench

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"regexp"
	"strings"

	"aigateway.local/core/internal/identity"
	"aigateway.local/core/internal/platform/apiresponse"
	"aigateway.local/core/internal/platform/cryptox"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// EnvVarService 管理门户用户个人环境变量(值加密存储)。
type EnvVarService struct {
	pool   *pgxpool.Pool
	cipher cryptox.Cipher
}

func NewEnvVarService(pool *pgxpool.Pool, cipher cryptox.Cipher) *EnvVarService {
	return &EnvVarService{pool: pool, cipher: cipher}
}

// List 返回用户环境变量(键列表,不含值)。
func (s *EnvVarService) List(ctx context.Context, userID string) ([]map[string]any, error) {
	if s == nil || s.pool == nil {
		return nil, errors.New("环境变量服务不可用")
	}
	rows, err := s.pool.Query(ctx, `SELECT key,octet_length(encrypted_value)>0,updated_at FROM gateway.user_env_vars WHERE portal_user_id=$1 ORDER BY key`, userID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []map[string]any{}
	for rows.Next() {
		var key string
		var hasValue bool
		var updatedAt any
		if err := rows.Scan(&key, &hasValue, &updatedAt); err != nil {
			return nil, err
		}
		items = append(items, map[string]any{"key": key, "configured": hasValue, "updated_at": updatedAt})
	}
	return items, rows.Err()
}

// Upsert 设置一个环境变量;value 为空时删除。
func (s *EnvVarService) Upsert(ctx context.Context, userID, key, value string) error {
	if s == nil || s.pool == nil || s.cipher == nil {
		return errors.New("环境变量服务不可用")
	}
	key = strings.TrimSpace(key)
	if key == "" || len(key) > 128 || !envKeyPattern.MatchString(key) {
		return errors.New("变量名必须以字母开头,可含字母/数字/下划线,最长 128 字符")
	}
	if len(value) > 4096 {
		return errors.New("变量值过长")
	}
	value = strings.TrimSpace(value)
	if value == "" {
		tag, err := s.pool.Exec(ctx, `DELETE FROM gateway.user_env_vars WHERE portal_user_id=$1 AND key=$2`, userID, key)
		if err != nil {
			return err
		}
		if tag.RowsAffected() == 0 {
			return errors.New("变量不存在")
		}
		return nil
	}
	encrypted, version, err := s.cipher.Encrypt([]byte(value))
	if err != nil {
		return err
	}
	_, err = s.pool.Exec(ctx, `INSERT INTO gateway.user_env_vars(portal_user_id,key,encrypted_value,value_kek_version) VALUES($1,$2,$3,$4)
		ON CONFLICT(portal_user_id,key) DO UPDATE SET encrypted_value=$3,value_kek_version=$4,updated_at=clock_timestamp()`,
		userID, key, encrypted, version)
	return err
}

// Decrypt 解密单个变量(运行时合并用);不存在返回 ok=false。
func (s *EnvVarService) Decrypt(ctx context.Context, userID, key string) (string, bool, error) {
	if s == nil || s.pool == nil || s.cipher == nil {
		return "", false, nil
	}
	var encrypted []byte
	var version int
	err := s.pool.QueryRow(ctx, `SELECT encrypted_value,value_kek_version FROM gateway.user_env_vars WHERE portal_user_id=$1 AND key=$2`, userID, key).Scan(&encrypted, &version)
	if errors.Is(err, pgx.ErrNoRows) {
		return "", false, nil
	}
	if err != nil {
		return "", false, err
	}
	plaintext, err := s.cipher.Decrypt(encrypted, version)
	if err != nil {
		return "", false, err
	}
	return string(plaintext), true, nil
}

// MergeVariables 把平台变量与个人变量合并进请求变量(请求已提供的键保持优先,
// 个人变量覆盖平台默认值)。
func (s *EnvVarService) MergeVariables(ctx context.Context, userID string, variables map[string]any) error {
	if s == nil || s.pool == nil || s.cipher == nil {
		return nil
	}
	if len(variables) >= 100 {
		return nil
	}
	merged, err := s.mergeAll(ctx, userID, variables)
	if err != nil {
		return err
	}
	for key, value := range merged {
		variables[key] = value
	}
	return nil
}

func (s *EnvVarService) mergeAll(ctx context.Context, userID string, variables map[string]any) (map[string]any, error) {
	out := map[string]any{}
	type pair struct {
		key     string
		value   []byte
		version int
	}
	collect := func(query string, args ...any) ([]pair, error) {
		rows, err := s.pool.Query(ctx, query, args...)
		if err != nil {
			return nil, err
		}
		defer rows.Close()
		pairs := []pair{}
		for rows.Next() {
			var p pair
			if err := rows.Scan(&p.key, &p.value, &p.version); err != nil {
				return nil, err
			}
			pairs = append(pairs, p)
		}
		return pairs, rows.Err()
	}
	// 平台变量(全部,最多 200)。
	platform, err := collect(`SELECT key,encrypted_value,value_kek_version FROM gateway.platform_env_vars ORDER BY key LIMIT 200`)
	if err != nil {
		return nil, err
	}
	for _, p := range platform {
		if _, exists := variables[p.key]; exists {
			continue
		}
		plaintext, decryptErr := s.cipher.Decrypt(p.value, p.version)
		if decryptErr != nil {
			continue
		}
		out[p.key] = string(plaintext)
	}
	// 个人变量覆盖平台默认值。
	if userID != "" {
		personal, err := collect(`SELECT key,encrypted_value,value_kek_version FROM gateway.user_env_vars WHERE portal_user_id=$1 ORDER BY key LIMIT 200`, userID)
		if err != nil {
			return nil, err
		}
		for _, p := range personal {
			if _, exists := variables[p.key]; exists {
				continue
			}
			plaintext, decryptErr := s.cipher.Decrypt(p.value, p.version)
			if decryptErr != nil {
				continue
			}
			out[p.key] = string(plaintext)
		}
	}
	return out, nil
}

// PlatformList 返回平台环境变量(不含值)。
func (s *EnvVarService) PlatformList(ctx context.Context) ([]map[string]any, error) {
	if s == nil || s.pool == nil {
		return nil, errors.New("环境变量服务不可用")
	}
	rows, err := s.pool.Query(ctx, `SELECT key,octet_length(encrypted_value)>0,description,updated_at FROM gateway.platform_env_vars ORDER BY key`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []map[string]any{}
	for rows.Next() {
		var key, description string
		var hasValue bool
		var updatedAt any
		if err := rows.Scan(&key, &hasValue, &description, &updatedAt); err != nil {
			return nil, err
		}
		items = append(items, map[string]any{"key": key, "configured": hasValue, "description": description, "updated_at": updatedAt})
	}
	return items, rows.Err()
}

// PlatformUpsert 设置平台环境变量;value 为空时删除。
func (s *EnvVarService) PlatformUpsert(ctx context.Context, actorID, key, value, description string) error {
	if s == nil || s.pool == nil || s.cipher == nil {
		return errors.New("环境变量服务不可用")
	}
	key = strings.TrimSpace(key)
	if key == "" || len(key) > 128 || !envKeyPattern.MatchString(key) {
		return errors.New("变量名必须以字母开头,可含字母/数字/下划线,最长 128 字符")
	}
	if len(description) > 512 {
		return errors.New("描述过长")
	}
	value = strings.TrimSpace(value)
	if value == "" {
		tag, err := s.pool.Exec(ctx, `DELETE FROM gateway.platform_env_vars WHERE key=$1`, key)
		if err != nil {
			return err
		}
		if tag.RowsAffected() == 0 {
			return errors.New("变量不存在")
		}
		return nil
	}
	if len(value) > 4096 {
		return errors.New("变量值过长")
	}
	encrypted, version, err := s.cipher.Encrypt([]byte(value))
	if err != nil {
		return err
	}
	_, err = s.pool.Exec(ctx, `INSERT INTO gateway.platform_env_vars(key,encrypted_value,value_kek_version,description,updated_by) VALUES($1,$2,$3,$4,$5)
		ON CONFLICT(key) DO UPDATE SET encrypted_value=$2,value_kek_version=$3,description=$4,updated_by=$5,updated_at=clock_timestamp()`,
		key, encrypted, version, description, actorID)
	return err
}

var envKeyPattern = regexp.MustCompile(`^[A-Za-z][A-Za-z0-9_]{0,127}$`)

// EnvVarHTTPHandler 门户环境变量 CRUD。
type EnvVarHTTPHandler struct {
	service  *EnvVarService
	identity *identity.Service
	mux      *http.ServeMux
}

func NewEnvVarHTTPHandler(service *EnvVarService, identityService *identity.Service) *EnvVarHTTPHandler {
	h := &EnvVarHTTPHandler{service: service, identity: identityService, mux: http.NewServeMux()}
	h.mux.HandleFunc("GET /api/v1/portal/env-vars", h.list)
	h.mux.HandleFunc("PUT /api/v1/portal/env-vars/{key}", h.upsert)
	h.mux.HandleFunc("DELETE /api/v1/portal/env-vars/{key}", h.delete)
	return h
}

func (h *EnvVarHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }

func (h *EnvVarHTTPHandler) account(w http.ResponseWriter, r *http.Request) (identity.Account, bool) {
	account, err := h.identity.Authenticate(r.Context(), identity.KindPortal, r.Header.Get("Authorization"))
	if err != nil {
		apiresponse.Error(w, http.StatusUnauthorized, "登录状态无效或已过期")
		return identity.Account{}, false
	}
	return account, true
}

func (h *EnvVarHTTPHandler) list(w http.ResponseWriter, r *http.Request) {
	a, ok := h.account(w, r)
	if !ok {
		return
	}
	items, err := h.service.List(r.Context(), a.ID)
	if err != nil {
		apiresponse.Error(w, http.StatusServiceUnavailable, "环境变量查询失败")
		return
	}
	apiresponse.OK(w, items)
}

func (h *EnvVarHTTPHandler) upsert(w http.ResponseWriter, r *http.Request) {
	a, ok := h.account(w, r)
	if !ok {
		return
	}
	var input struct {
		Value string `json:"value"`
	}
	decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
	decoder.DisallowUnknownFields()
	if decoder.Decode(&input) != nil {
		apiresponse.Error(w, http.StatusBadRequest, "请求格式无效")
		return
	}
	if err := h.service.Upsert(r.Context(), a.ID, r.PathValue("key"), input.Value); err != nil {
		apiresponse.Error(w, http.StatusBadRequest, err.Error())
		return
	}
	apiresponse.OK(w, map[string]bool{"saved": true})
}

func (h *EnvVarHTTPHandler) delete(w http.ResponseWriter, r *http.Request) {
	a, ok := h.account(w, r)
	if !ok {
		return
	}
	if err := h.service.Upsert(r.Context(), a.ID, r.PathValue("key"), ""); err != nil {
		apiresponse.Error(w, http.StatusBadRequest, err.Error())
		return
	}
	apiresponse.OK(w, map[string]bool{"deleted": true})
}

// AdminEnvVarHTTPHandler 平台环境变量管理(系统管理员)。
type AdminEnvVarHTTPHandler struct {
	service  *EnvVarService
	identity *identity.Service
	mux      *http.ServeMux
}

func NewAdminEnvVarHTTPHandler(service *EnvVarService, identityService *identity.Service) *AdminEnvVarHTTPHandler {
	h := &AdminEnvVarHTTPHandler{service: service, identity: identityService, mux: http.NewServeMux()}
	h.mux.HandleFunc("GET /api/v1/admin/env-vars", h.list)
	h.mux.HandleFunc("PUT /api/v1/admin/env-vars/{key}", h.upsert)
	h.mux.HandleFunc("DELETE /api/v1/admin/env-vars/{key}", h.delete)
	return h
}

func (h *AdminEnvVarHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.mux.ServeHTTP(w, r)
}

func (h *AdminEnvVarHTTPHandler) admin(w http.ResponseWriter, r *http.Request) (identity.Account, bool) {
	account, err := h.identity.Authenticate(r.Context(), identity.KindAdmin, r.Header.Get("Authorization"))
	if err != nil {
		apiresponse.Error(w, http.StatusUnauthorized, "登录状态无效或已过期")
		return identity.Account{}, false
	}
	if !identity.HasPermission(account, identity.PermissionSystemManage) {
		apiresponse.Error(w, http.StatusForbidden, "无系统管理权限")
		return identity.Account{}, false
	}
	return account, true
}

func (h *AdminEnvVarHTTPHandler) list(w http.ResponseWriter, r *http.Request) {
	if _, ok := h.admin(w, r); !ok {
		return
	}
	items, err := h.service.PlatformList(r.Context())
	if err != nil {
		apiresponse.Error(w, http.StatusServiceUnavailable, "环境变量查询失败")
		return
	}
	apiresponse.OK(w, items)
}

func (h *AdminEnvVarHTTPHandler) upsert(w http.ResponseWriter, r *http.Request) {
	admin, ok := h.admin(w, r)
	if !ok {
		return
	}
	var input struct {
		Value       string `json:"value"`
		Description string `json:"description"`
	}
	decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
	decoder.DisallowUnknownFields()
	if decoder.Decode(&input) != nil {
		apiresponse.Error(w, http.StatusBadRequest, "请求格式无效")
		return
	}
	if err := h.service.PlatformUpsert(r.Context(), admin.ID, r.PathValue("key"), input.Value, input.Description); err != nil {
		apiresponse.Error(w, http.StatusBadRequest, err.Error())
		return
	}
	apiresponse.OK(w, map[string]bool{"saved": true})
}

func (h *AdminEnvVarHTTPHandler) delete(w http.ResponseWriter, r *http.Request) {
	if _, ok := h.admin(w, r); !ok {
		return
	}
	if err := h.service.PlatformUpsert(r.Context(), "", r.PathValue("key"), "", ""); err != nil {
		apiresponse.Error(w, http.StatusBadRequest, err.Error())
		return
	}
	apiresponse.OK(w, map[string]bool{"deleted": true})
}