sub2api-cn-relay-manager/internal/app/key_self_service.go

package app

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"net/http"
	"strings"
)

func generatePlaintextKey() (string, string) {
	buf := make([]byte, 32)
	_, _ = rand.Read(buf)
	plaintext := "sk-" + hex.EncodeToString(buf)
	hash := sha256.Sum256([]byte(plaintext))
	return plaintext, "sha256:" + hex.EncodeToString(hash[:])
}

type UserKeyHandler struct {
	TrustedSubjectHeader     string
	TrustedProxySecretHeader string
	TrustedProxySecret       string
	createFn                 func(ctx context.Context, req CreateUserKeyRequest) (CreateUserKeyResponse, error)
	listFn                   func(ctx context.Context, subjectID string) ([]UserKeyMeta, error)
	getFn                    func(ctx context.Context, keyID, subjectID string) (UserKeyMeta, error)
	resetFn                  func(ctx context.Context, keyID, subjectID string) (ResetUserKeyResponse, error)
	pauseFn                  func(ctx context.Context, keyID, subjectID, reason string) (UserKeyMeta, error)
	resumeFn                 func(ctx context.Context, keyID, subjectID string) (UserKeyMeta, error)
	deleteFn                 func(ctx context.Context, keyID, subjectID string) error
}

type CreateUserKeyRequest struct {
	LogicalGroupID string   `json:"logical_group_id"`
	DisplayName    string   `json:"display_name"`
	AllowedModels  []string `json:"allowed_models"`
	SubjectID      string   `json:"-"`
}

type CreateUserKeyResponse struct {
	Key          UserKeyMeta `json:"key"`
	PlaintextKey string      `json:"plaintext_key,omitempty"`
}

type ResetUserKeyResponse struct {
	PlaintextKey  string `json:"plaintext_key,omitempty"`
	MaskedPreview string `json:"masked_preview"`
	AdminStatus   string `json:"admin_status"`
}

type UserKeyMeta struct {
	KeyID          string   `json:"key_id"`
	MaskedPreview  string   `json:"masked_preview"`
	DisplayName    string   `json:"display_name"`
	LogicalGroupID string   `json:"logical_group_id"`
	AllowedModels  []string `json:"allowed_models"`
	AdminStatus    string   `json:"admin_status"`
	QuotaStatus    string   `json:"quota_status"`
	LastUsedAt     string   `json:"last_used_at,omitempty"`
	CreatedAt      string   `json:"created_at"`
	ExpiresAt      string   `json:"expires_at,omitempty"`
}

func (h *UserKeyHandler) extractSubjectID(r *http.Request) (string, *httpError) {
	if h == nil {
		return "", &httpError{StatusCode: http.StatusUnauthorized, Code: "unauthorized", Message: "user credentials required"}
	}
	subjectHeader := strings.TrimSpace(h.TrustedSubjectHeader)
	secretHeader := strings.TrimSpace(h.TrustedProxySecretHeader)
	secret := strings.TrimSpace(h.TrustedProxySecret)
	if subjectHeader == "" || secretHeader == "" || secret == "" {
		return "", &httpError{StatusCode: http.StatusUnauthorized, Code: "unauthorized", Message: "trusted user identity proxy not configured"}
	}
	if got := strings.TrimSpace(r.Header.Get(secretHeader)); got != secret {
		return "", &httpError{StatusCode: http.StatusUnauthorized, Code: "unauthorized", Message: "trusted proxy authentication required"}
	}
	if subjectID := strings.TrimSpace(r.Header.Get(subjectHeader)); subjectID != "" {
		return subjectID, nil
	}
	return "", &httpError{StatusCode: http.StatusUnauthorized, Code: "unauthorized", Message: "trusted subject header required"}
}

func writeSvcNotImplError(w http.ResponseWriter) {
	writeHTTPError(w, &httpError{StatusCode: http.StatusNotImplemented, Code: "not_implemented", Message: "user key service not configured"})
}

func handleCreateUserKey(w http.ResponseWriter, r *http.Request, h *UserKeyHandler) {
	if h == nil || h.createFn == nil {
		writeSvcNotImplError(w)
		return
	}
	subjectID, httpErr := h.extractSubjectID(r)
	if httpErr != nil {
		writeHTTPError(w, httpErr)
		return
	}
	var req CreateUserKeyRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeHTTPError(w, &httpError{StatusCode: http.StatusBadRequest, Code: "invalid_json", Message: err.Error()})
		return
	}
	req.SubjectID = subjectID
	resp, svcErr := h.createFn(r.Context(), req)
	if svcErr != nil {
		writeHTTPError(w, classifyError(svcErr))
		return
	}
	writeJSON(w, http.StatusCreated, resp)
}

func handleListUserKeys(w http.ResponseWriter, r *http.Request, h *UserKeyHandler) {
	if h == nil || h.listFn == nil {
		writeSvcNotImplError(w)
		return
	}
	subjectID, httpErr := h.extractSubjectID(r)
	if httpErr != nil {
		writeHTTPError(w, httpErr)
		return
	}
	keys, svcErr := h.listFn(r.Context(), subjectID)
	if svcErr != nil {
		writeHTTPError(w, classifyError(svcErr))
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"keys": keys})
}

func handleGetUserKey(w http.ResponseWriter, r *http.Request, h *UserKeyHandler) {
	if h == nil || h.getFn == nil {
		writeSvcNotImplError(w)
		return
	}
	subjectID, httpErr := h.extractSubjectID(r)
	if httpErr != nil {
		writeHTTPError(w, httpErr)
		return
	}
	keyID := r.PathValue("key_id")
	if keyID == "" {
		writeHTTPError(w, &httpError{StatusCode: http.StatusBadRequest, Code: "missing_key_id", Message: "key_id required"})
		return
	}
	key, svcErr := h.getFn(r.Context(), keyID, subjectID)
	if svcErr != nil {
		writeHTTPError(w, classifyError(svcErr))
		return
	}
	writeJSON(w, http.StatusOK, key)
}

func handleResetUserKey(w http.ResponseWriter, r *http.Request, h *UserKeyHandler) {
	if h == nil || h.resetFn == nil {
		writeSvcNotImplError(w)
		return
	}
	subjectID, httpErr := h.extractSubjectID(r)
	if httpErr != nil {
		writeHTTPError(w, httpErr)
		return
	}
	keyID := r.PathValue("key_id")
	if keyID == "" {
		writeHTTPError(w, &httpError{StatusCode: http.StatusBadRequest, Code: "missing_key_id", Message: "key_id required"})
		return
	}
	resp, svcErr := h.resetFn(r.Context(), keyID, subjectID)
	if svcErr != nil {
		writeHTTPError(w, classifyError(svcErr))
		return
	}
	writeJSON(w, http.StatusOK, resp)
}

func handlePauseUserKey(w http.ResponseWriter, r *http.Request, h *UserKeyHandler) {
	if h == nil || h.pauseFn == nil {
		writeSvcNotImplError(w)
		return
	}
	subjectID, httpErr := h.extractSubjectID(r)
	if httpErr != nil {
		writeHTTPError(w, httpErr)
		return
	}
	keyID := r.PathValue("key_id")
	var req struct {
		Reason string `json:"reason"`
	}
	if r.Body != nil && r.ContentLength != 0 {
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeHTTPError(w, &httpError{StatusCode: http.StatusBadRequest, Code: "invalid_json", Message: err.Error()})
			return
		}
	}
	key, svcErr := h.pauseFn(r.Context(), keyID, subjectID, strings.TrimSpace(req.Reason))
	if svcErr != nil {
		writeHTTPError(w, classifyError(svcErr))
		return
	}
	writeJSON(w, http.StatusOK, key)
}

func handleResumeUserKey(w http.ResponseWriter, r *http.Request, h *UserKeyHandler) {
	if h == nil || h.resumeFn == nil {
		writeSvcNotImplError(w)
		return
	}
	subjectID, httpErr := h.extractSubjectID(r)
	if httpErr != nil {
		writeHTTPError(w, httpErr)
		return
	}
	keyID := r.PathValue("key_id")
	key, svcErr := h.resumeFn(r.Context(), keyID, subjectID)
	if svcErr != nil {
		writeHTTPError(w, classifyError(svcErr))
		return
	}
	writeJSON(w, http.StatusOK, key)
}

func handleDeleteUserKey(w http.ResponseWriter, r *http.Request, h *UserKeyHandler) {
	if h == nil || h.deleteFn == nil {
		writeSvcNotImplError(w)
		return
	}
	subjectID, httpErr := h.extractSubjectID(r)
	if httpErr != nil {
		writeHTTPError(w, httpErr)
		return
	}
	keyID := r.PathValue("key_id")
	if svcErr := h.deleteFn(r.Context(), keyID, subjectID); svcErr != nil {
		writeHTTPError(w, classifyError(svcErr))
		return
	}
	writeJSON(w, http.StatusOK, map[string]string{"status": "deleted"})
}