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 {
	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) {
	for _, header := range []string{"X-Portal-Subject", "X-User-Subject", "X-Forwarded-User"} {
		if subjectID := strings.TrimSpace(r.Header.Get(header)); subjectID != "" {
			return subjectID, nil
		}
	}
	if hdr := r.Header.Get("Authorization"); strings.HasPrefix(hdr, "Bearer ") {
		token := strings.TrimSpace(strings.TrimPrefix(hdr, "Bearer "))
		if token != "" {
			n := 8
			if len(token) < n {
				n = len(token)
			}
			return "skeleton_user_" + token[:n], nil
		}
	}
	return "", &httpError{StatusCode: http.StatusUnauthorized, Code: "unauthorized", Message: "user credentials 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")
	key, svcErr := h.pauseFn(r.Context(), keyID, subjectID, "")
	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"})
}