From 1ee6bea413bb6aae265df23506dd016b90701436 Mon Sep 17 00:00:00 2001 From: RC-CHN <1051989940@qq.com> Date: Tue, 6 Jan 2026 09:15:49 +0800 Subject: [PATCH] feat(api): enhance whoami endpoint with realtime stats and extended key info Add realtime statistics (requests, tokens, QPS, rate limiting) to whoami response for both master and key authentication types. Extend key response with additional fields including master name, model limits, quota tracking, and usage statistics. - Inject StatsService into AuthHandler for realtime stats retrieval - Add WhoamiRealtimeView struct for realtime statistics - Include admin permissions field in admin response - Add comprehensive key metadata (quotas, model limits, usage stats) - Add test for expired key returning 401 Unauthorized --- cmd/server/main.go | 2 +- internal/api/auth_handler.go | 162 ++++++++++++++++++++++++------ internal/api/auth_handler_test.go | 56 ++++++++++- 3 files changed, 188 insertions(+), 32 deletions(-) diff --git a/cmd/server/main.go b/cmd/server/main.go index a3cfc05..310fb5c 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -232,7 +232,7 @@ func main() { alertHandler := api.NewAlertHandler(db) internalHandler := api.NewInternalHandler(db) featureHandler := api.NewFeatureHandler(rdb) - authHandler := api.NewAuthHandler(db, rdb, adminService, masterService) + authHandler := api.NewAuthHandler(db, rdb, adminService, masterService, statsService) ipBanService := service.NewIPBanService(db, rdb) ipBanHandler := api.NewIPBanHandler(ipBanService) ipBanManager := cron.NewIPBanManager(ipBanService) diff --git a/internal/api/auth_handler.go b/internal/api/auth_handler.go index 690756a..94788b9 100644 --- a/internal/api/auth_handler.go +++ b/internal/api/auth_handler.go @@ -20,15 +20,17 @@ type AuthHandler struct { rdb *redis.Client adminService *service.AdminService masterService *service.MasterService + statsService *service.StatsService } // NewAuthHandler creates a new AuthHandler. -func NewAuthHandler(db *gorm.DB, rdb *redis.Client, adminService *service.AdminService, masterService *service.MasterService) *AuthHandler { +func NewAuthHandler(db *gorm.DB, rdb *redis.Client, adminService *service.AdminService, masterService *service.MasterService, statsService *service.StatsService) *AuthHandler { return &AuthHandler{ db: db, rdb: rdb, adminService: adminService, masterService: masterService, + statsService: statsService, } } @@ -72,6 +74,16 @@ type WhoamiKeyResponse struct { UpdatedAt int64 `json:"updated_at" example:"1703505600"` // Unix timestamp } +// WhoamiRealtimeView represents realtime statistics for the authenticated identity. +type WhoamiRealtimeView struct { + Requests int64 `json:"requests" example:"100"` // Total requests + Tokens int64 `json:"tokens" example:"50000"` // Total tokens used + QPS int64 `json:"qps" example:"5"` // Current QPS + QPSLimit int64 `json:"qps_limit" example:"100"` // QPS limit + RateLimited bool `json:"rate_limited" example:"false"` // Whether currently rate limited + UpdatedAt *int64 `json:"updated_at,omitempty" example:"1703505600"` // Last updated timestamp +} + // WhoamiResponse is a union type for all identity responses. // swagger:model WhoamiResponse type WhoamiResponse struct { @@ -79,7 +91,8 @@ type WhoamiResponse struct { Type string `json:"type" example:"master"` // Admin fields (only present when type is "admin") - Role string `json:"role,omitempty" example:"admin"` + Role string `json:"role,omitempty" example:"admin"` + Permissions []string `json:"permissions,omitempty"` // Admin permissions (always ["*"]) // Master fields (only present when type is "master") ID uint `json:"id,omitempty" example:"1"` @@ -95,13 +108,26 @@ type WhoamiResponse struct { UpdatedAt int64 `json:"updated_at,omitempty" example:"1703505600"` // Key fields (only present when type is "key") - MasterID uint `json:"master_id,omitempty" example:"1"` - Scopes string `json:"scopes,omitempty" example:"chat:write"` - IssuedAtEpoch int64 `json:"issued_at_epoch,omitempty" example:"1"` - IssuedBy string `json:"issued_by,omitempty" example:"master"` - AllowIPs string `json:"allow_ips,omitempty" example:""` // IP whitelist (for diagnostics) - DenyIPs string `json:"deny_ips,omitempty" example:""` // IP blacklist (for diagnostics) - ExpiresAt int64 `json:"expires_at,omitempty" example:"0"` // Expiration timestamp (0 = never) + MasterID uint `json:"master_id,omitempty" example:"1"` + MasterName string `json:"master_name,omitempty" example:"tenant-a"` // Parent master name (for display) + Scopes string `json:"scopes,omitempty" example:"chat:write"` + IssuedAtEpoch int64 `json:"issued_at_epoch,omitempty" example:"1"` + IssuedBy string `json:"issued_by,omitempty" example:"master"` + AllowIPs string `json:"allow_ips,omitempty" example:""` // IP whitelist (for diagnostics) + DenyIPs string `json:"deny_ips,omitempty" example:""` // IP blacklist (for diagnostics) + ExpiresAt int64 `json:"expires_at,omitempty" example:"0"` // Expiration timestamp (0 = never) + ModelLimits string `json:"model_limits,omitempty" example:"gpt-4,claude"` // Model restrictions + ModelLimitsEnabled bool `json:"model_limits_enabled,omitempty" example:"false"` // Whether model limits are active + QuotaLimit int64 `json:"quota_limit,omitempty" example:"-1"` // Token quota limit (-1 = unlimited) + QuotaUsed int64 `json:"quota_used,omitempty" example:"0"` // Token quota used + QuotaResetAt int64 `json:"quota_reset_at,omitempty" example:"0"` // Quota reset timestamp + QuotaResetType string `json:"quota_reset_type,omitempty" example:"monthly"` // Quota reset type + RequestCount int64 `json:"request_count,omitempty" example:"0"` // Total request count (from DB) + UsedTokens int64 `json:"used_tokens,omitempty" example:"0"` // Total tokens used (from DB) + LastAccessedAt int64 `json:"last_accessed_at,omitempty" example:"0"` // Last access timestamp + + // Realtime stats (for master and key types) + Realtime *WhoamiRealtimeView `json:"realtime,omitempty"` } // Whoami godoc @@ -138,8 +164,9 @@ func (h *AuthHandler) Whoami(c *gin.Context) { // 1. Try admin token first if h.adminService.ValidateToken(token) { c.JSON(http.StatusOK, WhoamiResponse{ - Type: "admin", - Role: "admin", + Type: "admin", + Role: "admin", + Permissions: []string{"*"}, // Admin has all permissions }) return } @@ -147,7 +174,7 @@ func (h *AuthHandler) Whoami(c *gin.Context) { // 2. Try master key master, err := h.masterService.ValidateMasterKey(token) if err == nil && master != nil { - c.JSON(http.StatusOK, WhoamiResponse{ + resp := WhoamiResponse{ Type: "master", ID: master.ID, Name: master.Name, @@ -160,7 +187,31 @@ func (h *AuthHandler) Whoami(c *gin.Context) { GlobalQPS: master.GlobalQPS, CreatedAt: master.CreatedAt.UTC().Unix(), UpdatedAt: master.UpdatedAt.UTC().Unix(), - }) + } + + // Get realtime stats for master + if h.statsService != nil { + if stats, err := h.statsService.GetMasterRealtimeSnapshot(c.Request.Context(), master.ID); err == nil { + if stats.QPSLimit == 0 && master.GlobalQPS > 0 { + stats.QPSLimit = int64(master.GlobalQPS) + } + var updatedAt *int64 + if stats.UpdatedAt != nil { + sec := stats.UpdatedAt.Unix() + updatedAt = &sec + } + resp.Realtime = &WhoamiRealtimeView{ + Requests: stats.Requests, + Tokens: stats.Tokens, + QPS: stats.QPS, + QPSLimit: stats.QPSLimit, + RateLimited: stats.RateLimited, + UpdatedAt: updatedAt, + } + } + } + + c.JSON(http.StatusOK, resp) return } @@ -244,23 +295,74 @@ func (h *AuthHandler) Whoami(c *gin.Context) { // Token is valid, now get full data from DB for response var key model.Key if err := h.db.Where("token_hash = ?", tokenHash).First(&key).Error; err == nil { - c.JSON(http.StatusOK, WhoamiResponse{ - Type: "key", - ID: key.ID, - MasterID: uint(masterID), - Group: key.Group, - Scopes: strings.TrimSpace(key.Scopes), - DefaultNamespace: strings.TrimSpace(key.DefaultNamespace), - Namespaces: strings.TrimSpace(key.Namespaces), - Status: key.Status, - IssuedAtEpoch: key.IssuedAtEpoch, - IssuedBy: strings.TrimSpace(key.IssuedBy), - AllowIPs: strings.TrimSpace(key.AllowIPs), - DenyIPs: strings.TrimSpace(key.DenyIPs), - ExpiresAt: expiresAt, - CreatedAt: key.CreatedAt.UTC().Unix(), - UpdatedAt: key.UpdatedAt.UTC().Unix(), - }) + // Get master name for display + var master model.Master + masterName := "" + if err := h.db.First(&master, masterID).Error; err == nil { + masterName = master.Name + } + + // Convert nullable time fields + var quotaResetAt int64 + if key.QuotaResetAt != nil { + quotaResetAt = key.QuotaResetAt.UTC().Unix() + } + var lastAccessedAt int64 + if key.LastAccessedAt != nil { + lastAccessedAt = key.LastAccessedAt.UTC().Unix() + } + + resp := WhoamiResponse{ + Type: "key", + ID: key.ID, + MasterID: uint(masterID), + MasterName: masterName, + Group: key.Group, + Scopes: strings.TrimSpace(key.Scopes), + DefaultNamespace: strings.TrimSpace(key.DefaultNamespace), + Namespaces: strings.TrimSpace(key.Namespaces), + Status: key.Status, + IssuedAtEpoch: key.IssuedAtEpoch, + IssuedBy: strings.TrimSpace(key.IssuedBy), + AllowIPs: strings.TrimSpace(key.AllowIPs), + DenyIPs: strings.TrimSpace(key.DenyIPs), + ExpiresAt: expiresAt, + ModelLimits: strings.TrimSpace(key.ModelLimits), + ModelLimitsEnabled: key.ModelLimitsEnabled, + QuotaLimit: key.QuotaLimit, + QuotaUsed: key.QuotaUsed, + QuotaResetAt: quotaResetAt, + QuotaResetType: strings.TrimSpace(key.QuotaResetType), + RequestCount: key.RequestCount, + UsedTokens: key.UsedTokens, + LastAccessedAt: lastAccessedAt, + CreatedAt: key.CreatedAt.UTC().Unix(), + UpdatedAt: key.UpdatedAt.UTC().Unix(), + } + + // Get realtime stats for the key's master + if h.statsService != nil { + if stats, err := h.statsService.GetMasterRealtimeSnapshot(c.Request.Context(), uint(masterID)); err == nil { + if stats.QPSLimit == 0 && master.GlobalQPS > 0 { + stats.QPSLimit = int64(master.GlobalQPS) + } + var updatedAt *int64 + if stats.UpdatedAt != nil { + sec := stats.UpdatedAt.Unix() + updatedAt = &sec + } + resp.Realtime = &WhoamiRealtimeView{ + Requests: stats.Requests, + Tokens: stats.Tokens, + QPS: stats.QPS, + QPSLimit: stats.QPSLimit, + RateLimited: stats.RateLimited, + UpdatedAt: updatedAt, + } + } + } + + c.JSON(http.StatusOK, resp) return } } diff --git a/internal/api/auth_handler_test.go b/internal/api/auth_handler_test.go index 4cdb5f4..cd46967 100644 --- a/internal/api/auth_handler_test.go +++ b/internal/api/auth_handler_test.go @@ -40,7 +40,8 @@ func newAuthHandler(t *testing.T) (*AuthHandler, *gorm.DB, *miniredis.Miniredis) rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) masterService := service.NewMasterService(db) - handler := NewAuthHandler(db, rdb, adminService, masterService) + statsService := service.NewStatsService(rdb) + handler := NewAuthHandler(db, rdb, adminService, masterService, statsService) return handler, db, mr } @@ -152,3 +153,56 @@ func TestAuthHandler_Whoami_KeyResponseIncludesIPRules(t *testing.T) { t.Fatalf("expected expires_at=%d, got %v", expiresAt, resp["expires_at"]) } } + +func TestAuthHandler_Whoami_ExpiredKey_Returns401(t *testing.T) { + handler, db, mr := newAuthHandler(t) + + token := "sk-live-expired" + hash := tokenhash.HashToken(token) + // Set expiration to 1 hour ago + expiresAt := time.Now().Add(-time.Hour).Unix() + + mr.HSet("auth:token:"+hash, "master_id", "1") + mr.HSet("auth:token:"+hash, "issued_at_epoch", "1") + mr.HSet("auth:token:"+hash, "status", "active") + mr.HSet("auth:token:"+hash, "group", "default") + mr.HSet("auth:token:"+hash, "expires_at", fmt.Sprintf("%d", expiresAt)) + mr.HSet("auth:master:1", "epoch", "1") + mr.HSet("auth:master:1", "status", "active") + + // Create key in DB + key := &model.Key{ + MasterID: 1, + TokenHash: hash, + Group: "default", + Scopes: "chat:write", + DefaultNamespace: "default", + Namespaces: "default", + Status: "active", + IssuedAtEpoch: 1, + IssuedBy: "master", + } + if err := db.Create(key).Error; err != nil { + t.Fatalf("create key: %v", err) + } + + r := gin.New() + r.GET("/auth/whoami", handler.Whoami) + + req := httptest.NewRequest(http.MethodGet, "/auth/whoami", nil) + req.Header.Set("Authorization", "Bearer "+token) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 for expired key, got %d body=%s", rr.Code, rr.Body.String()) + } + + var resp map[string]any + if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode response: %v", err) + } + if resp["error"] != "token has expired" { + t.Fatalf("expected 'token has expired' error, got %v", resp["error"]) + } +}