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

package app

import (
	"context"
	"crypto/rand"
	"fmt"
	"math/big"

	"sub2api-cn-relay-manager/internal/store/sqlite"
)

const (
	keyIDAlphabet = "abcdefghijklmnopqrstuvwxyz0123456789"
)

func generateKeyID() string {
	n := big.NewInt(int64(len(keyIDAlphabet)))
	b := make([]byte, 12)
	for i := range b {
		idx, _ := rand.Int(rand.Reader, n)
		b[i] = keyIDAlphabet[idx.Int64()]
	}
	return "key_" + string(b)
}

func buildUserKeyHandler(sqliteDSN string) *UserKeyHandler {
	return &UserKeyHandler{
		createFn: func(ctx context.Context, req CreateUserKeyRequest) (CreateUserKeyResponse, error) {
			store, err := sqlite.Open(ctx, sqliteDSN)
			if err != nil {
				return CreateUserKeyResponse{}, fmt.Errorf("open store: %w", err)
			}
			defer store.Close()

			plaintext, fingerprint := generatePlaintextKey()
			keyID := generateKeyID()
			masked := "sk-****" + plaintext[len(plaintext)-4:]

			_, err = store.UserKeys().Create(ctx, sqlite.UserKeyRecord{
				KeyID:          keyID,
				OwnerSubjectID: req.SubjectID,
				KeyFingerprint: fingerprint,
				MaskedPreview:  masked,
				DisplayName:    req.DisplayName,
				LogicalGroupID: req.LogicalGroupID,
				AllowedModels:  req.AllowedModels,
				AdminStatus:    "active",
				QuotaStatus:    "ok",
			})
			if err != nil {
				return CreateUserKeyResponse{}, fmt.Errorf("create key: %w", err)
			}

			return CreateUserKeyResponse{
				Key: UserKeyMeta{
					KeyID:          keyID,
					MaskedPreview:  masked,
					DisplayName:    req.DisplayName,
					LogicalGroupID: req.LogicalGroupID,
					AllowedModels:  req.AllowedModels,
					AdminStatus:    "active",
					QuotaStatus:    "ok",
				},
				PlaintextKey: plaintext,
			}, nil
		},
		listFn: func(ctx context.Context, subjectID string) ([]UserKeyMeta, error) {
			store, err := sqlite.Open(ctx, sqliteDSN)
			if err != nil {
				return nil, fmt.Errorf("open store: %w", err)
			}
			defer store.Close()

			records, err := store.UserKeys().ListByOwner(ctx, subjectID)
			if err != nil {
				return nil, fmt.Errorf("list keys: %w", err)
			}

			metas := make([]UserKeyMeta, len(records))
			for i, r := range records {
				metas[i] = UserKeyMeta{
					KeyID:          r.KeyID,
					MaskedPreview:  r.MaskedPreview,
					DisplayName:    r.DisplayName,
					LogicalGroupID: r.LogicalGroupID,
					AllowedModels:  r.AllowedModels,
					AdminStatus:    r.AdminStatus,
					QuotaStatus:    r.QuotaStatus,
					LastUsedAt:     r.LastUsedAt,
					CreatedAt:      r.CreatedAt,
					ExpiresAt:      r.ExpiresAt,
				}
			}
			return metas, nil
		},
		getFn: func(ctx context.Context, keyID, subjectID string) (UserKeyMeta, error) {
			store, err := sqlite.Open(ctx, sqliteDSN)
			if err != nil {
				return UserKeyMeta{}, fmt.Errorf("open store: %w", err)
			}
			defer store.Close()

			rec, err := store.UserKeys().GetByID(ctx, keyID)
			if err != nil {
				return UserKeyMeta{}, fmt.Errorf("get key: %w", err)
			}
			if rec.OwnerSubjectID != subjectID && subjectID != "admin" {
				return UserKeyMeta{}, fmt.Errorf("not found")
			}
			return UserKeyMeta{
				KeyID:          rec.KeyID,
				MaskedPreview:  rec.MaskedPreview,
				DisplayName:    rec.DisplayName,
				LogicalGroupID: rec.LogicalGroupID,
				AllowedModels:  rec.AllowedModels,
				AdminStatus:    rec.AdminStatus,
				QuotaStatus:    rec.QuotaStatus,
				LastUsedAt:     rec.LastUsedAt,
				CreatedAt:      rec.CreatedAt,
				ExpiresAt:      rec.ExpiresAt,
			}, nil
		},
		pauseFn: func(ctx context.Context, keyID, subjectID, reason string) (UserKeyMeta, error) {
			store, err := sqlite.Open(ctx, sqliteDSN)
			if err != nil {
				return UserKeyMeta{}, fmt.Errorf("open store: %w", err)
			}
			defer store.Close()
			if err := store.UserKeys().UpdateStatus(ctx, keyID, "paused"); err != nil {
				return UserKeyMeta{}, fmt.Errorf("pause key: %w", err)
			}
			rec, _ := store.UserKeys().GetByID(ctx, keyID)
			if rec != nil {
				return UserKeyMeta{
					KeyID:         rec.KeyID,
					MaskedPreview: rec.MaskedPreview,
					AdminStatus:   rec.AdminStatus,
				}, nil
			}
			return UserKeyMeta{KeyID: keyID, AdminStatus: "paused"}, nil
		},
		resumeFn: func(ctx context.Context, keyID, subjectID string) (UserKeyMeta, error) {
			store, err := sqlite.Open(ctx, sqliteDSN)
			if err != nil {
				return UserKeyMeta{}, fmt.Errorf("open store: %w", err)
			}
			defer store.Close()
			if err := store.UserKeys().UpdateStatus(ctx, keyID, "active"); err != nil {
				return UserKeyMeta{}, fmt.Errorf("resume key: %w", err)
			}
			return UserKeyMeta{KeyID: keyID, AdminStatus: "active"}, nil
		},
		deleteFn: func(ctx context.Context, keyID, subjectID string) error {
			store, err := sqlite.Open(ctx, sqliteDSN)
			if err != nil {
				return fmt.Errorf("open store: %w", err)
			}
			defer store.Close()
			return store.UserKeys().UpdateStatus(ctx, keyID, "retired")
		},
	}
}