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"]) + } +}