ruochen/ez-api
1
0
Fork 0
mirror of https://github.com/EZ-Api/ez-api.git synced 2026-01-13 09:37:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(api): enhance whoami endpoint with realtime stats and extended key info

Browse Source 
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
This commit is contained in:
RC-CHN
2026-01-06 09:15:49 +08:00
parent f99e4a15ab
commit 1ee6bea413
3 changed files with 188 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
cmd/server/main.go
View File

@@ -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)

162
internal/api/auth_handler.go
View File

@@ -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
}
}

56
internal/api/auth_handler_test.go
View File

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