AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,203 @@
|
||||
package apikey
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
type AdminHTTPHandler struct {
|
||||
repository *Repository
|
||||
authenticator *Authenticator
|
||||
usage *UsageStore
|
||||
identity *identity.Service
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
type createRequest struct {
|
||||
Name string `json:"name"`
|
||||
Scopes []string `json:"scopes"`
|
||||
RequestsPerMinute int `json:"requests_per_minute"`
|
||||
MonthlyRequestQuota int64 `json:"monthly_request_quota"`
|
||||
MonthlyTokenQuota int64 `json:"monthly_token_quota"`
|
||||
ExpiresAt *time.Time `json:"expires_at"`
|
||||
}
|
||||
|
||||
type limitsRequest struct {
|
||||
RequestsPerMinute int `json:"requests_per_minute"`
|
||||
MonthlyRequestQuota int64 `json:"monthly_request_quota"`
|
||||
MonthlyTokenQuota int64 `json:"monthly_token_quota"`
|
||||
}
|
||||
|
||||
func NewAdminHTTPHandler(repository *Repository, authenticator *Authenticator, identityService *identity.Service) *AdminHTTPHandler {
|
||||
h := &AdminHTTPHandler{repository: repository, authenticator: authenticator, identity: identityService, mux: http.NewServeMux()}
|
||||
h.mux.HandleFunc("GET /api/v1/admin/api-keys", h.list)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/api-keys", h.create)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/api-keys/{api_key_id}/limits", h.updateLimits)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/api-keys/{api_key_id}", h.revoke)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) SetUsageStore(store *UsageStore) { h.usage = store }
|
||||
|
||||
func (h *AdminHTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
|
||||
h.mux.ServeHTTP(writer, request)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) list(writer http.ResponseWriter, request *http.Request) {
|
||||
if _, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyRead); !ok {
|
||||
return
|
||||
}
|
||||
records, err := h.repository.List(request.Context())
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
ids := make([]string, 0, len(records))
|
||||
for _, record := range records {
|
||||
ids = append(ids, record.ID)
|
||||
}
|
||||
usage, err := h.usage.MonthlyTokens(request.Context(), ids, time.Now())
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
items := make([]map[string]any, 0, len(records))
|
||||
for _, record := range records {
|
||||
item := publicRecord(record)
|
||||
item["monthly_token_usage"] = usage[record.ID]
|
||||
items = append(items, item)
|
||||
}
|
||||
apiresponse.OK(writer, items)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) updateLimits(writer http.ResponseWriter, request *http.Request) {
|
||||
account, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input limitsRequest
|
||||
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
|
||||
decoder.DisallowUnknownFields()
|
||||
if decoder.Decode(&input) != nil || !validLimits(input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota) {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "API Key 限流或月配额无效")
|
||||
return
|
||||
}
|
||||
record, hash, err := h.repository.UpdateLimits(request.Context(), request.PathValue("api_key_id"), input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota, account.ID)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
if err := h.authenticator.Invalidate(request.Context(), hash); err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
response := publicRecord(record)
|
||||
if h.usage != nil {
|
||||
usage, usageErr := h.usage.MonthlyTokens(request.Context(), []string{record.ID}, time.Now())
|
||||
if usageErr == nil {
|
||||
response["monthly_token_usage"] = usage[record.ID]
|
||||
}
|
||||
}
|
||||
apiresponse.OK(writer, response)
|
||||
}
|
||||
|
||||
func validLimits(requestsPerMinute int, monthlyRequestQuota, monthlyTokenQuota int64) bool {
|
||||
return requestsPerMinute >= 0 && requestsPerMinute <= 1_000_000 &&
|
||||
monthlyRequestQuota >= 0 && monthlyRequestQuota <= 1_000_000_000_000 &&
|
||||
monthlyTokenQuota >= 0 && monthlyTokenQuota <= 1_000_000_000_000_000
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) create(writer http.ResponseWriter, request *http.Request) {
|
||||
account, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input createRequest
|
||||
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
|
||||
decoder.DisallowUnknownFields()
|
||||
if decoder.Decode(&input) != nil {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "请求格式无效")
|
||||
return
|
||||
}
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
if input.Name == "" || len(input.Name) > 128 || len(input.Scopes) == 0 || input.ExpiresAt != nil && !input.ExpiresAt.After(time.Now()) {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "名称、权限范围或过期时间无效")
|
||||
return
|
||||
}
|
||||
if !validLimits(input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota) {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "API Key 限流或月配额无效")
|
||||
return
|
||||
}
|
||||
for _, scope := range input.Scopes {
|
||||
if scope != "gateway:invoke" && scope != "*" {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "包含不支持的权限范围")
|
||||
return
|
||||
}
|
||||
}
|
||||
record, secret, err := h.repository.Create(request.Context(), input.Name, input.Scopes, input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota, input.ExpiresAt, account.ID)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
response := publicRecord(record)
|
||||
response["key"] = secret
|
||||
response["warning"] = "密钥只显示一次，请立即安全保存"
|
||||
apiresponse.OK(writer, response)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) revoke(writer http.ResponseWriter, request *http.Request) {
|
||||
account, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
hash, err := h.repository.Revoke(request.Context(), request.PathValue("api_key_id"), account.ID)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
if err := h.authenticator.Invalidate(request.Context(), hash); err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(writer, map[string]bool{"revoked": true})
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) requirePermission(writer http.ResponseWriter, request *http.Request, permission string) (identity.Account, bool) {
|
||||
account, err := h.identity.Authenticate(request.Context(), identity.KindAdmin, request.Header.Get("Authorization"))
|
||||
if err != nil {
|
||||
apiresponse.Error(writer, http.StatusUnauthorized, "登录状态无效或已过期")
|
||||
return identity.Account{}, false
|
||||
}
|
||||
if !identity.HasPermission(account, permission) {
|
||||
apiresponse.Error(writer, http.StatusForbidden, "缺少 API Key 操作权限")
|
||||
return identity.Account{}, false
|
||||
}
|
||||
return account, true
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) writeError(writer http.ResponseWriter, err error) {
|
||||
switch {
|
||||
case errors.Is(err, ErrInvalid):
|
||||
apiresponse.Error(writer, http.StatusNotFound, "API Key 不存在或已撤销")
|
||||
case errors.Is(err, ErrStore):
|
||||
apiresponse.Error(writer, http.StatusServiceUnavailable, "API Key 服务暂不可用")
|
||||
default:
|
||||
apiresponse.Error(writer, http.StatusInternalServerError, "API Key 处理失败")
|
||||
}
|
||||
}
|
||||
|
||||
func publicRecord(record Record) map[string]any {
|
||||
return map[string]any{
|
||||
"id": record.ID, "name": record.Name, "key_prefix": record.KeyPrefix,
|
||||
"scopes": record.Scopes, "enabled": record.Enabled, "expires_at": record.ExpiresAt,
|
||||
"requests_per_minute": record.RequestsPerMinute, "monthly_request_quota": record.MonthlyRequestQuota,
|
||||
"monthly_token_quota": record.MonthlyTokenQuota,
|
||||
"last_used_at": record.LastUsedAt, "created_at": record.CreatedAt,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package apikey
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrInvalid = errors.New("invalid API key")
|
||||
ErrStore = errors.New("API key store unavailable")
|
||||
)
|
||||
|
||||
type Record struct {
|
||||
ID string
|
||||
TenantID *string
|
||||
Name string
|
||||
KeyPrefix string
|
||||
KeyHash []byte
|
||||
Scopes []string
|
||||
Enabled bool
|
||||
RequestsPerMinute int
|
||||
MonthlyRequestQuota int64
|
||||
MonthlyTokenQuota int64
|
||||
ExpiresAt *time.Time
|
||||
LastUsedAt *time.Time
|
||||
CreatedAt time.Time
|
||||
}
|
||||
|
||||
type Principal struct {
|
||||
APIKeyID string `json:"api_key_id"`
|
||||
TenantID *string `json:"tenant_id,omitempty"`
|
||||
Scopes []string `json:"scopes"`
|
||||
RequestsPerMinute int `json:"requests_per_minute"`
|
||||
MonthlyRequestQuota int64 `json:"monthly_request_quota"`
|
||||
MonthlyTokenQuota int64 `json:"monthly_token_quota"`
|
||||
}
|
||||
|
||||
func Generate() (secret, prefix string, hash []byte, err error) {
|
||||
random := make([]byte, 32)
|
||||
if _, err = rand.Read(random); err != nil {
|
||||
return "", "", nil, err
|
||||
}
|
||||
secret = "gw_" + base64.RawURLEncoding.EncodeToString(random)
|
||||
prefix = secret[:16]
|
||||
digest := sha256.Sum256([]byte(secret))
|
||||
return secret, prefix, digest[:], nil
|
||||
}
|
||||
|
||||
func Digest(secret string) ([]byte, string) {
|
||||
digest := sha256.Sum256([]byte(strings.TrimSpace(secret)))
|
||||
return digest[:], hex.EncodeToString(digest[:])
|
||||
}
|
||||
|
||||
func HasScope(scopes []string, required string) bool {
|
||||
for _, scope := range scopes {
|
||||
if scope == "*" || scope == required {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
type KeyAuthenticator interface {
|
||||
Authenticate(context.Context, string) error
|
||||
}
|
||||
|
||||
type PrincipalAuthenticator interface {
|
||||
AuthenticatePrincipal(context.Context, string) (Principal, error)
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package apikey
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBootstrapCompatibilityUsageCounter(t *testing.T) {
|
||||
authenticator := NewAuthenticator(nil, nil, "temporary-bootstrap-key")
|
||||
if err := authenticator.Authenticate(context.Background(), "temporary-bootstrap-key"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if authenticator.BootstrapUses() != 1 {
|
||||
t.Fatalf("unexpected bootstrap usage count %d", authenticator.BootstrapUses())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateAndDigest(t *testing.T) {
|
||||
first, prefix, hash, err := Generate()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(first) < 40 || len(prefix) != 16 || prefix != first[:16] {
|
||||
t.Fatal("invalid API key format")
|
||||
}
|
||||
computed, _ := Digest(first)
|
||||
if !bytes.Equal(hash, computed) {
|
||||
t.Fatal("stored digest differs")
|
||||
}
|
||||
second, _, _, err := Generate()
|
||||
if err != nil || second == first {
|
||||
t.Fatal("API keys must be independently random")
|
||||
}
|
||||
}
|
||||
|
||||
func TestScopes(t *testing.T) {
|
||||
if !HasScope([]string{"gateway:invoke"}, "gateway:invoke") || !HasScope([]string{"*"}, "gateway:invoke") {
|
||||
t.Fatal("expected scope is missing")
|
||||
}
|
||||
if HasScope([]string{"gateway:read"}, "gateway:invoke") {
|
||||
t.Fatal("unexpected scope accepted")
|
||||
}
|
||||
}
|
||||
|
||||
// TestInvokeScope pins the regression where portal application runtime
|
||||
// credentials ("application:run") could never reach the gateway: the scope
|
||||
// check only admitted "gateway:invoke", so every hosted application chat
|
||||
// returned 401 even though the credential is legitimate.
|
||||
func TestInvokeScope(t *testing.T) {
|
||||
accept := map[string][]string{
|
||||
"gateway key": {"gateway:invoke"},
|
||||
"admin wildcard": {"*"},
|
||||
"application runtime key": {"application:run"},
|
||||
"runtime + read": {"application:run", "gateway:read"},
|
||||
}
|
||||
for name, scopes := range accept {
|
||||
if !invokeScope(scopes) {
|
||||
t.Fatalf("invokeScope(%v) = false, want true for %s", scopes, name)
|
||||
}
|
||||
}
|
||||
reject := map[string][]string{
|
||||
"read-only": {"gateway:read"},
|
||||
"unrelated scope": {"workbench:run"},
|
||||
"empty": {},
|
||||
"runtime scope missing": {"gateway:read", "workbench:run"},
|
||||
}
|
||||
for name, scopes := range reject {
|
||||
if invokeScope(scopes) {
|
||||
t.Fatalf("invokeScope(%v) = true, want false for %s", scopes, name)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package apikey
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type Authenticator struct {
|
||||
repository *Repository
|
||||
redis *redis.Client
|
||||
bootstrap string
|
||||
cacheTTL time.Duration
|
||||
bootstrapUses atomic.Uint64
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
func NewAuthenticator(repository *Repository, client *redis.Client, bootstrap string) *Authenticator {
|
||||
return &Authenticator{repository: repository, redis: client, bootstrap: bootstrap, cacheTTL: 30 * time.Second}
|
||||
}
|
||||
|
||||
// SetLogger wires an optional logger used for best-effort cache diagnostics.
|
||||
func (a *Authenticator) SetLogger(logger *slog.Logger) { a.logger = logger }
|
||||
|
||||
func (a *Authenticator) Authenticate(ctx context.Context, secret string) error {
|
||||
_, err := a.AuthenticatePrincipal(ctx, secret)
|
||||
return err
|
||||
}
|
||||
|
||||
// invokeScope reports whether the key carries a scope that may reach the
|
||||
// gateway. Besides regular gateway keys it admits the server-side
|
||||
// application runtime credential ("application:run"), which is created only
|
||||
// by the portal runtime (encrypted in PostgreSQL, tenant-bound, never shown
|
||||
// to browsers) so that hosted application conversations can call the gateway
|
||||
// on behalf of the app owner. Both call paths in the workbench runtime and
|
||||
// the main proxy use this single method, so the runtime credential works
|
||||
// end-to-end without widening any other surface.
|
||||
func invokeScope(scopes []string) bool {
|
||||
return HasScope(scopes, "gateway:invoke") || HasScope(scopes, "application:run")
|
||||
}
|
||||
|
||||
func (a *Authenticator) AuthenticatePrincipal(ctx context.Context, secret string) (Principal, error) {
|
||||
secret = strings.TrimSpace(secret)
|
||||
if secret == "" {
|
||||
return Principal{}, ErrInvalid
|
||||
}
|
||||
if a.bootstrap != "" && len(secret) == len(a.bootstrap) && subtle.ConstantTimeCompare([]byte(secret), []byte(a.bootstrap)) == 1 {
|
||||
a.bootstrapUses.Add(1)
|
||||
// The bootstrap credential is a real gateway key, not an anonymous
|
||||
// pass. Return a stable, well-known identity so the admission and
|
||||
// token-quota paths run (they allow it outright: RPM/quota are 0 by
|
||||
// design for a migration key) and audit records attribute usage to
|
||||
// "bootstrap" instead of an empty principal that silently skips every
|
||||
// policy stage. An empty APIKeyID previously bypassed rate limiting,
|
||||
// quota and audit attribution entirely.
|
||||
return Principal{APIKeyID: "bootstrap", Scopes: []string{"gateway:invoke"}}, nil
|
||||
}
|
||||
hash, hexHash := Digest(secret)
|
||||
if a.redis != nil {
|
||||
payload, err := a.redis.Get(ctx, cacheKey(hexHash)).Bytes()
|
||||
if err == nil {
|
||||
var principal Principal
|
||||
if json.Unmarshal(payload, &principal) == nil && invokeScope(principal.Scopes) {
|
||||
return principal, nil
|
||||
}
|
||||
return Principal{}, ErrInvalid
|
||||
}
|
||||
if err != nil && !errors.Is(err, redis.Nil) {
|
||||
return Principal{}, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
}
|
||||
principal, err := a.repository.Validate(ctx, hash)
|
||||
if err != nil {
|
||||
return Principal{}, err
|
||||
}
|
||||
if !invokeScope(principal.Scopes) {
|
||||
return Principal{}, ErrInvalid
|
||||
}
|
||||
// Best-effort cache: the database is the source of truth for validation.
|
||||
// A cache write failure must not turn a key the database just accepted
|
||||
// into a 503, or a transient Redis blip would take the whole gateway down.
|
||||
if a.redis != nil {
|
||||
payload, _ := json.Marshal(principal)
|
||||
if err := a.redis.Set(ctx, cacheKey(hexHash), payload, a.cacheTTL).Err(); err != nil && a.logger != nil {
|
||||
a.logger.Warn("api key cache write failed; continuing with database result", "error", err)
|
||||
}
|
||||
}
|
||||
return principal, nil
|
||||
}
|
||||
|
||||
func (a *Authenticator) BootstrapUses() uint64 {
|
||||
if a == nil {
|
||||
return 0
|
||||
}
|
||||
return a.bootstrapUses.Load()
|
||||
}
|
||||
|
||||
func (a *Authenticator) Invalidate(ctx context.Context, hash []byte) error {
|
||||
if a.redis == nil {
|
||||
return nil
|
||||
}
|
||||
_, hexHash := DigestFromHash(hash)
|
||||
return a.redis.Del(ctx, cacheKey(hexHash)).Err()
|
||||
}
|
||||
|
||||
func DigestFromHash(hash []byte) ([]byte, string) {
|
||||
const hex = "0123456789abcdef"
|
||||
encoded := make([]byte, len(hash)*2)
|
||||
for i, value := range hash {
|
||||
encoded[i*2] = hex[value>>4]
|
||||
encoded[i*2+1] = hex[value&15]
|
||||
}
|
||||
return hash, string(encoded)
|
||||
}
|
||||
|
||||
func cacheKey(hexHash string) string { return "gateway:api-key:v1:" + hexHash }
|
||||
@@ -0,0 +1,179 @@
|
||||
package apikey
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type Repository struct{ pool *pgxpool.Pool }
|
||||
|
||||
func NewRepository(pool *pgxpool.Pool) *Repository { return &Repository{pool: pool} }
|
||||
|
||||
func (r *Repository) List(ctx context.Context) ([]Record, error) {
|
||||
if r.pool == nil {
|
||||
return nil, ErrStore
|
||||
}
|
||||
rows, err := r.pool.Query(ctx, `
|
||||
SELECT id::text, tenant_id::text, name, key_prefix, scopes, enabled,
|
||||
requests_per_minute, monthly_request_quota, monthly_token_quota, expires_at, last_used_at, created_at
|
||||
FROM gateway.api_keys ORDER BY created_at DESC`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
defer rows.Close()
|
||||
var records []Record
|
||||
for rows.Next() {
|
||||
var record Record
|
||||
if err := rows.Scan(&record.ID, &record.TenantID, &record.Name, &record.KeyPrefix, &record.Scopes, &record.Enabled, &record.RequestsPerMinute, &record.MonthlyRequestQuota, &record.MonthlyTokenQuota, &record.ExpiresAt, &record.LastUsedAt, &record.CreatedAt); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
records = append(records, record)
|
||||
}
|
||||
return records, rows.Err()
|
||||
}
|
||||
|
||||
func (r *Repository) Create(ctx context.Context, name string, scopes []string, requestsPerMinute int, monthlyRequestQuota, monthlyTokenQuota int64, expiresAt *time.Time, actorID string) (Record, string, error) {
|
||||
if r.pool == nil {
|
||||
return Record{}, "", ErrStore
|
||||
}
|
||||
secret, prefix, hash, err := Generate()
|
||||
if err != nil {
|
||||
return Record{}, "", err
|
||||
}
|
||||
id, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Record{}, "", err
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Record{}, "", err
|
||||
}
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return Record{}, "", fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
record := Record{ID: id, Name: name, KeyPrefix: prefix, KeyHash: hash, Scopes: scopes, Enabled: true, RequestsPerMinute: requestsPerMinute, MonthlyRequestQuota: monthlyRequestQuota, MonthlyTokenQuota: monthlyTokenQuota, ExpiresAt: expiresAt}
|
||||
err = tx.QueryRow(ctx, `
|
||||
INSERT INTO gateway.api_keys (id, name, key_prefix, key_hash, scopes, requests_per_minute, monthly_request_quota, monthly_token_quota, expires_at, created_by)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, nullif($10,'')::uuid)
|
||||
RETURNING created_at`, id, name, prefix, hash, scopes, requestsPerMinute, monthlyRequestQuota, monthlyTokenQuota, expiresAt, actorID).Scan(&record.CreatedAt)
|
||||
if err != nil {
|
||||
return Record{}, "", fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"api_key_id": id, "name": name, "key_prefix": prefix, "requests_per_minute": requestsPerMinute, "monthly_request_quota": monthlyRequestQuota, "monthly_token_quota": monthlyTokenQuota})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway.outbox_events (event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
|
||||
VALUES ($1, 'api_key.created', 1, 'api_key', $2, $3)`, eventID, id, payload); err != nil {
|
||||
return Record{}, "", fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return Record{}, "", fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
return record, secret, nil
|
||||
}
|
||||
|
||||
func (r *Repository) Validate(ctx context.Context, hash []byte) (Principal, error) {
|
||||
if r.pool == nil {
|
||||
return Principal{}, ErrStore
|
||||
}
|
||||
var principal Principal
|
||||
err := r.pool.QueryRow(ctx, `
|
||||
UPDATE gateway.api_keys
|
||||
SET last_used_at = CASE WHEN last_used_at IS NULL OR last_used_at < clock_timestamp() - interval '5 minutes' THEN clock_timestamp() ELSE last_used_at END
|
||||
WHERE key_hash = $1 AND enabled AND (expires_at IS NULL OR expires_at > clock_timestamp())
|
||||
RETURNING id::text, tenant_id::text, scopes, requests_per_minute, monthly_request_quota, monthly_token_quota`, hash).Scan(&principal.APIKeyID, &principal.TenantID, &principal.Scopes, &principal.RequestsPerMinute, &principal.MonthlyRequestQuota, &principal.MonthlyTokenQuota)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return Principal{}, ErrInvalid
|
||||
}
|
||||
if err != nil {
|
||||
return Principal{}, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
return principal, nil
|
||||
}
|
||||
|
||||
func (r *Repository) Revoke(ctx context.Context, id, actorID string) ([]byte, error) {
|
||||
if r.pool == nil {
|
||||
return nil, ErrStore
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
var hash []byte
|
||||
err = tx.QueryRow(ctx, `
|
||||
UPDATE gateway.api_keys SET enabled = false, updated_at = clock_timestamp()
|
||||
WHERE id = $1 AND enabled RETURNING key_hash`, id).Scan(&hash)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, ErrInvalid
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"api_key_id": id, "actor_id": actorID})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway.outbox_events (event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
|
||||
VALUES ($1, 'api_key.revoked', 1, 'api_key', $2, $3)`, eventID, id, payload); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
return hash, nil
|
||||
}
|
||||
|
||||
func (r *Repository) UpdateLimits(ctx context.Context, id string, requestsPerMinute int, monthlyRequestQuota, monthlyTokenQuota int64, actorID string) (Record, []byte, error) {
|
||||
if r.pool == nil {
|
||||
return Record{}, nil, ErrStore
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Record{}, nil, err
|
||||
}
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return Record{}, nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
var record Record
|
||||
err = tx.QueryRow(ctx, `
|
||||
UPDATE gateway.api_keys
|
||||
SET requests_per_minute = $2, monthly_request_quota = $3, monthly_token_quota = $4, updated_at = clock_timestamp()
|
||||
WHERE id = $1 AND enabled
|
||||
RETURNING id::text, tenant_id::text, name, key_prefix, key_hash, scopes, enabled,
|
||||
requests_per_minute, monthly_request_quota, monthly_token_quota, expires_at, last_used_at, created_at`,
|
||||
id, requestsPerMinute, monthlyRequestQuota, monthlyTokenQuota,
|
||||
).Scan(&record.ID, &record.TenantID, &record.Name, &record.KeyPrefix, &record.KeyHash, &record.Scopes, &record.Enabled,
|
||||
&record.RequestsPerMinute, &record.MonthlyRequestQuota, &record.MonthlyTokenQuota, &record.ExpiresAt, &record.LastUsedAt, &record.CreatedAt)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return Record{}, nil, ErrInvalid
|
||||
}
|
||||
if err != nil {
|
||||
return Record{}, nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"api_key_id": id, "actor_id": actorID, "requests_per_minute": requestsPerMinute,
|
||||
"monthly_request_quota": monthlyRequestQuota, "monthly_token_quota": monthlyTokenQuota,
|
||||
})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway.outbox_events (event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
|
||||
VALUES ($1, 'api_key.limits_updated', 1, 'api_key', $2, $3)`, eventID, id, payload); err != nil {
|
||||
return Record{}, nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return Record{}, nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
return record, record.KeyHash, nil
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package apikey
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type UsageStore struct{ redis *redis.Client }
|
||||
|
||||
func NewUsageStore(client *redis.Client) *UsageStore { return &UsageStore{redis: client} }
|
||||
|
||||
func MonthlyTokenUsageKey(apiKeyID string, now time.Time) string {
|
||||
return fmt.Sprintf("gateway:usage:api-key:%s:tokens:%s", apiKeyID, now.UTC().Format("200601"))
|
||||
}
|
||||
|
||||
func (s *UsageStore) MonthlyTokens(ctx context.Context, apiKeyIDs []string, now time.Time) (map[string]int64, error) {
|
||||
usage := make(map[string]int64, len(apiKeyIDs))
|
||||
if len(apiKeyIDs) == 0 || s == nil || s.redis == nil {
|
||||
return usage, nil
|
||||
}
|
||||
pipe := s.redis.Pipeline()
|
||||
commands := make(map[string]*redis.StringCmd, len(apiKeyIDs))
|
||||
for _, id := range apiKeyIDs {
|
||||
commands[id] = pipe.Get(ctx, MonthlyTokenUsageKey(id, now))
|
||||
}
|
||||
_, err := pipe.Exec(ctx)
|
||||
if err != nil && err != redis.Nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
for id, command := range commands {
|
||||
value, commandErr := command.Int64()
|
||||
if commandErr == nil {
|
||||
usage[id] = max(value, 0)
|
||||
} else if commandErr != redis.Nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, commandErr)
|
||||
}
|
||||
}
|
||||
return usage, nil
|
||||
}
|
||||
Reference in New Issue
Block a user