feat(auth): implement master key authentication system with child key issuance

Add admin and master authentication layers with JWT support. Replace direct
key creation with hierarchical master/child key system. Update database
schema to support master accounts with configurable limits and epoch-based
key revocation. Add health check endpoint with system status monitoring.

BREAKING CHANGE: Removed direct POST /keys endpoint in favor of master-based
key issuance through /v1/tokens. Database migration requires dropping old User
table and creating Master table with new relationships.
This commit is contained in:
zenfun
2025-12-05 00:16:47 +08:00
parent 5360cc6f1a
commit 8645b22b83
16 changed files with 618 additions and 229 deletions

View File

@@ -0,0 +1,54 @@
package api
import (
"net/http"
"github.com/ez-api/ez-api/internal/service"
"github.com/gin-gonic/gin"
)
type AdminHandler struct {
masterService *service.MasterService
}
func NewAdminHandler(masterService *service.MasterService) *AdminHandler {
return &AdminHandler{masterService: masterService}
}
type CreateMasterRequest struct {
Name string `json:"name" binding:"required"`
Group string `json:"group" binding:"required"`
MaxChildKeys int `json:"max_child_keys"`
GlobalQPS int `json:"global_qps"`
}
func (h *AdminHandler) CreateMaster(c *gin.Context) {
var req CreateMasterRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
// Use defaults if not provided
if req.MaxChildKeys == 0 {
req.MaxChildKeys = 5
}
if req.GlobalQPS == 0 {
req.GlobalQPS = 3
}
master, rawMasterKey, err := h.masterService.CreateMaster(req.Name, req.Group, req.MaxChildKeys, req.GlobalQPS)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create master key", "details": err.Error()})
return
}
c.JSON(http.StatusCreated, gin.H{
"id": master.ID,
"name": master.Name,
"group": master.Group,
"master_key": rawMasterKey, // Only show this on creation
"max_child_keys": master.MaxChildKeys,
"global_qps": master.GlobalQPS,
})
}

View File

@@ -22,39 +22,7 @@ func NewHandler(db *gorm.DB, sync *service.SyncService, logger *service.LogWrite
return &Handler{db: db, sync: sync, logger: logger}
}
func (h *Handler) CreateKey(c *gin.Context) {
var req dto.KeyDTO
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
group := strings.TrimSpace(req.Group)
if group == "" {
group = "default"
}
key := model.Key{
KeySecret: req.KeySecret,
Group: group,
Balance: req.Balance,
Status: req.Status,
Weight: req.Weight,
}
if err := h.db.Create(&key).Error; err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create key", "details": err.Error()})
return
}
// Write auth hash and refresh snapshots
if err := h.sync.SyncKey(&key); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to sync key to Redis", "details": err.Error()})
return
}
c.JSON(http.StatusCreated, key)
}
// CreateKey is now handled by MasterHandler
func (h *Handler) CreateProvider(c *gin.Context) {
var req dto.ProviderDTO

View File

@@ -0,0 +1,64 @@
package api
import (
"net/http"
"strings"
"github.com/ez-api/ez-api/internal/model"
"github.com/ez-api/ez-api/internal/service"
"github.com/gin-gonic/gin"
)
type MasterHandler struct {
masterService *service.MasterService
}
func NewMasterHandler(masterService *service.MasterService) *MasterHandler {
return &MasterHandler{masterService: masterService}
}
type IssueChildKeyRequest struct {
Group string `json:"group"`
Scopes string `json:"scopes"`
}
func (h *MasterHandler) IssueChildKey(c *gin.Context) {
master, exists := c.Get("master")
if !exists {
c.JSON(http.StatusUnauthorized, gin.H{"error": "master key not found in context"})
return
}
masterModel := master.(*model.Master)
var req IssueChildKeyRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
// If group is not specified, inherit from master
group := req.Group
if strings.TrimSpace(group) == "" {
group = masterModel.Group
}
// Security: Ensure the requested group is allowed for this master.
// For now, we'll just enforce it's the same group.
if group != masterModel.Group {
c.JSON(http.StatusForbidden, gin.H{"error": "cannot issue key for a different group"})
return
}
key, rawChildKey, err := h.masterService.IssueChildKey(masterModel.ID, group, req.Scopes)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to issue child key", "details": err.Error()})
return
}
c.JSON(http.StatusCreated, gin.H{
"id": key.ID,
"key_secret": rawChildKey,
"group": key.Group,
"scopes": key.Scopes,
})
}

View File

@@ -11,12 +11,17 @@ type Config struct {
Postgres PostgresConfig
Redis RedisConfig
Log LogConfig
Auth AuthConfig
}
type ServerConfig struct {
Port string
}
type AuthConfig struct {
JWTSecret string
}
type PostgresConfig struct {
DSN string
}
@@ -51,6 +56,9 @@ func Load() (*Config, error) {
FlushInterval: getEnvDuration("EZ_LOG_FLUSH_MS", 1000),
QueueCapacity: getEnvInt("EZ_LOG_QUEUE", 10000),
},
Auth: AuthConfig{
JWTSecret: getEnv("EZ_JWT_SECRET", "change_me_in_production"),
},
}, nil
}

View File

@@ -0,0 +1,63 @@
package middleware
import (
"net/http"
"strings"
"github.com/ez-api/ez-api/internal/service"
"github.com/gin-gonic/gin"
)
func AdminAuthMiddleware(adminService *service.AdminService) gin.HandlerFunc {
return func(c *gin.Context) {
authHeader := c.GetHeader("Authorization")
if authHeader == "" {
c.JSON(http.StatusUnauthorized, gin.H{"error": "authorization header required"})
c.Abort()
return
}
parts := strings.Split(authHeader, " ")
if len(parts) != 2 || parts[0] != "Bearer" {
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid authorization header format"})
c.Abort()
return
}
if !adminService.ValidateToken(parts[1]) {
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid admin token"})
c.Abort()
return
}
c.Next()
}
}
func MasterAuthMiddleware(masterService *service.MasterService) gin.HandlerFunc {
return func(c *gin.Context) {
authHeader := c.GetHeader("Authorization")
if authHeader == "" {
c.JSON(http.StatusUnauthorized, gin.H{"error": "authorization header required"})
c.Abort()
return
}
parts := strings.Split(authHeader, " ")
if len(parts) != 2 || parts[0] != "Bearer" {
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid authorization header format"})
c.Abort()
return
}
master, err := masterService.ValidateMasterKey(parts[1])
if err != nil {
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid master key"})
c.Abort()
return
}
c.Set("master", master)
c.Next()
}
}

View File

@@ -4,13 +4,32 @@ import (
"gorm.io/gorm"
)
type User struct {
// Admin is not a database model. It's configured via environment variables.
// Master represents a tenant account.
type Master struct {
gorm.Model
Username string `gorm:"uniqueIndex;not null" json:"username"`
Quota int64 `gorm:"default:0" json:"quota"`
Role string `gorm:"default:'user'" json:"role"` // admin, user
Name string `gorm:"size:255" json:"name"`
MasterKey string `gorm:"size:255;uniqueIndex" json:"-"` // Hashed master key
Group string `gorm:"size:100;default:'default'" json:"group"`
Epoch int64 `gorm:"default:1" json:"epoch"`
Status string `gorm:"size:50;default:'active'" json:"status"` // active, suspended
MaxChildKeys int `gorm:"default:5" json:"max_child_keys"`
GlobalQPS int `gorm:"default:3" json:"global_qps"`
}
// Key represents a child access token issued by a Master.
type Key struct {
gorm.Model
MasterID uint `gorm:"not null;index" json:"master_id"`
KeySecret string `gorm:"size:255;uniqueIndex" json:"key_secret"`
Group string `gorm:"size:100;default:'default'" json:"group"`
Scopes string `gorm:"size:1024" json:"scopes"` // Comma-separated scopes
IssuedAtEpoch int64 `gorm:"not null" json:"issued_at_epoch"`
Status string `gorm:"size:50;default:'active'" json:"status"` // active, suspended
}
// Provider remains the same.
type Provider struct {
gorm.Model
Name string `gorm:"not null" json:"name"`
@@ -21,15 +40,7 @@ type Provider struct {
Models string `json:"models"` // comma-separated list of supported models (e.g. "gpt-4,gpt-3.5-turbo")
}
type Key struct {
gorm.Model
KeySecret string `gorm:"not null" json:"key_secret"`
Group string `gorm:"default:'default'" json:"group"` // routing group/tier
Balance float64 `json:"balance"`
Status string `gorm:"default:'active'" json:"status"` // active, suspended
Weight int `gorm:"default:10" json:"weight"`
}
// Model remains the same.
type Model struct {
gorm.Model
Name string `gorm:"uniqueIndex;not null" json:"name"`

24
internal/service/admin.go Normal file
View File

@@ -0,0 +1,24 @@
package service
import (
"crypto/subtle"
"errors"
"os"
)
type AdminService struct {
adminToken string
}
func NewAdminService() (*AdminService, error) {
token := os.Getenv("EZ_ADMIN_TOKEN")
if token == "" {
return nil, errors.New("EZ_ADMIN_TOKEN environment variable not set")
}
return &AdminService{adminToken: token}, nil
}
// ValidateToken performs a constant-time comparison to prevent timing attacks.
func (s *AdminService) ValidateToken(token string) bool {
return subtle.ConstantTimeCompare([]byte(s.adminToken), []byte(token)) == 1
}

View File

@@ -0,0 +1,56 @@
package service
import (
"context"
"time"
"github.com/redis/go-redis/v9"
"gorm.io/gorm"
)
type HealthCheckService struct {
db *gorm.DB
rdb *redis.Client
}
func NewHealthCheckService(db *gorm.DB, rdb *redis.Client) *HealthCheckService {
return &HealthCheckService{
db: db,
rdb: rdb,
}
}
type HealthStatus struct {
Status string `json:"status"`
Database string `json:"database"`
Redis string `json:"redis"`
Uptime string `json:"uptime"`
}
var startTime = time.Now()
func (s *HealthCheckService) Check(ctx context.Context) HealthStatus {
status := HealthStatus{
Status: "ok",
Database: "up",
Redis: "up",
Uptime: time.Since(startTime).String(),
}
sqlDB, err := s.db.DB()
if err != nil || sqlDB.Ping() != nil {
status.Database = "down"
status.Status = "degraded"
}
if s.rdb.Ping(ctx).Err() != nil {
status.Redis = "down"
status.Status = "degraded"
}
if status.Database == "down" && status.Redis == "down" {
status.Status = "down"
}
return status
}

106
internal/service/master.go Normal file
View File

@@ -0,0 +1,106 @@
package service
import (
"crypto/rand"
"encoding/hex"
"errors"
"fmt"
"github.com/ez-api/ez-api/internal/model"
"golang.org/x/crypto/bcrypt"
"gorm.io/gorm"
)
type MasterService struct {
db *gorm.DB
}
func NewMasterService(db *gorm.DB) *MasterService {
return &MasterService{db: db}
}
func (s *MasterService) CreateMaster(name, group string, maxChildKeys, globalQPS int) (*model.Master, string, error) {
rawMasterKey, err := generateRandomKey(32)
if err != nil {
return nil, "", fmt.Errorf("failed to generate master key: %w", err)
}
hashedMasterKey, err := bcrypt.GenerateFromPassword([]byte(rawMasterKey), bcrypt.DefaultCost)
if err != nil {
return nil, "", fmt.Errorf("failed to hash master key: %w", err)
}
master := &model.Master{
Name: name,
MasterKey: string(hashedMasterKey),
Group: group,
MaxChildKeys: maxChildKeys,
GlobalQPS: globalQPS,
Status: "active",
Epoch: 1,
}
if err := s.db.Create(master).Error; err != nil {
return nil, "", err
}
return master, rawMasterKey, nil
}
func (s *MasterService) ValidateMasterKey(masterKey string) (*model.Master, error) {
// This is inefficient. We should query by a hash or an indexed field.
// For now, we iterate. In a real system, this needs optimization.
var masters []model.Master
if err := s.db.Find(&masters).Error; err != nil {
return nil, err
}
for _, master := range masters {
if bcrypt.CompareHashAndPassword([]byte(master.MasterKey), []byte(masterKey)) == nil {
return &master, nil
}
}
return nil, errors.New("invalid master key")
}
func (s *MasterService) IssueChildKey(masterID uint, group string, scopes string) (*model.Key, string, error) {
var master model.Master
if err := s.db.First(&master, masterID).Error; err != nil {
return nil, "", fmt.Errorf("master not found: %w", err)
}
var count int64
s.db.Model(&model.Key{}).Where("master_id = ?", masterID).Count(&count)
if count >= int64(master.MaxChildKeys) {
return nil, "", fmt.Errorf("child key limit reached for master %d", masterID)
}
rawChildKey, err := generateRandomKey(32)
if err != nil {
return nil, "", fmt.Errorf("failed to generate child key: %w", err)
}
key := &model.Key{
MasterID: masterID,
KeySecret: rawChildKey, // In a real system, this should also be hashed
Group: group,
Scopes: scopes,
IssuedAtEpoch: master.Epoch,
Status: "active",
}
if err := s.db.Create(key).Error; err != nil {
return nil, "", err
}
return key, rawChildKey, nil
}
func generateRandomKey(length int) (string, error) {
bytes := make([]byte, length)
if _, err := rand.Read(bytes); err != nil {
return "", err
}
return hex.EncodeToString(bytes), nil
}

View File

@@ -2,13 +2,12 @@ package service
import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"strings"
"github.com/ez-api/ez-api/internal/model"
"github.com/ez-api/ez-api/internal/util"
"github.com/redis/go-redis/v9"
"gorm.io/gorm"
)
@@ -24,26 +23,16 @@ func NewSyncService(rdb *redis.Client) *SyncService {
// SyncKey writes a single key into Redis without rebuilding the entire snapshot.
func (s *SyncService) SyncKey(key *model.Key) error {
ctx := context.Background()
snap := keySnapshot{
ID: key.ID,
TokenHash: hashToken(key.KeySecret),
Group: normalizeGroup(key.Group),
Status: key.Status,
Weight: key.Weight,
Balance: key.Balance,
}
if err := s.hsetJSON(ctx, "config:keys", snap.TokenHash, snap); err != nil {
return err
}
tokenHash := util.HashToken(key.KeySecret)
fields := map[string]interface{}{
"status": snap.Status,
"group": snap.Group,
"weight": snap.Weight,
"balance": snap.Balance,
"master_id": key.MasterID,
"issued_at_epoch": key.IssuedAtEpoch,
"status": key.Status,
"group": key.Group,
"scopes": key.Scopes,
}
if err := s.rdb.HSet(ctx, fmt.Sprintf("auth:token:%s", snap.TokenHash), fields).Err(); err != nil {
if err := s.rdb.HSet(ctx, fmt.Sprintf("auth:token:%s", tokenHash), fields).Err(); err != nil {
return fmt.Errorf("write auth token: %w", err)
}
return nil
@@ -111,14 +100,7 @@ type providerSnapshot struct {
Models []string `json:"models"`
}
type keySnapshot struct {
ID uint `json:"id"`
TokenHash string `json:"token_hash"`
Group string `json:"group"`
Status string `json:"status"`
Weight int `json:"weight"`
Balance float64 `json:"balance"`
}
// keySnapshot is no longer needed as we write directly to auth:token:*
type modelSnapshot struct {
Name string `json:"name"`
@@ -145,6 +127,11 @@ func (s *SyncService) SyncAll(db *gorm.DB) error {
return fmt.Errorf("load keys: %w", err)
}
var masters []model.Master
if err := db.Find(&masters).Error; err != nil {
return fmt.Errorf("load masters: %w", err)
}
var models []model.Model
if err := db.Find(&models).Error; err != nil {
return fmt.Errorf("load models: %w", err)
@@ -152,6 +139,18 @@ func (s *SyncService) SyncAll(db *gorm.DB) error {
pipe := s.rdb.TxPipeline()
pipe.Del(ctx, "config:providers", "config:keys", "meta:models")
// Also clear master keys
var masterKeys []string
iter := s.rdb.Scan(ctx, 0, "auth:master:*", 0).Iterator()
for iter.Next(ctx) {
masterKeys = append(masterKeys, iter.Val())
}
if err := iter.Err(); err != nil {
return fmt.Errorf("scan master keys: %w", err)
}
if len(masterKeys) > 0 {
pipe.Del(ctx, masterKeys...)
}
// Clear old routing tables (pattern scan would be better in prod, but keys are predictable if we knew them)
// For MVP, we rely on the fact that we are rebuilding.
@@ -188,24 +187,21 @@ func (s *SyncService) SyncAll(db *gorm.DB) error {
}
for _, k := range keys {
snap := keySnapshot{
ID: k.ID,
TokenHash: hashToken(k.KeySecret),
Group: normalizeGroup(k.Group),
Status: k.Status,
Weight: k.Weight,
Balance: k.Balance,
}
payload, err := json.Marshal(snap)
if err != nil {
return fmt.Errorf("marshal key %d: %w", k.ID, err)
}
pipe.HSet(ctx, "config:keys", snap.TokenHash, payload)
pipe.HSet(ctx, fmt.Sprintf("auth:token:%s", snap.TokenHash), map[string]interface{}{
"status": snap.Status,
"group": snap.Group,
"weight": snap.Weight,
"balance": snap.Balance,
tokenHash := util.HashToken(k.KeySecret)
pipe.HSet(ctx, fmt.Sprintf("auth:token:%s", tokenHash), map[string]interface{}{
"master_id": k.MasterID,
"issued_at_epoch": k.IssuedAtEpoch,
"status": k.Status,
"group": k.Group,
"scopes": k.Scopes,
})
}
for _, m := range masters {
pipe.HSet(ctx, fmt.Sprintf("auth:master:%d", m.ID), map[string]interface{}{
"epoch": m.Epoch,
"status": m.Status,
"global_qps": m.GlobalQPS,
})
}
@@ -245,12 +241,6 @@ func (s *SyncService) hsetJSON(ctx context.Context, key, field string, val inter
return nil
}
func hashToken(token string) string {
hasher := sha256.New()
hasher.Write([]byte(token))
return hex.EncodeToString(hasher.Sum(nil))
}
func normalizeGroup(group string) string {
if strings.TrimSpace(group) == "" {
return "default"

69
internal/service/token.go Normal file
View File

@@ -0,0 +1,69 @@
package service
import (
"context"
"errors"
"fmt"
"strconv"
"github.com/ez-api/ez-api/internal/util"
"github.com/redis/go-redis/v9"
)
type TokenService struct {
rdb *redis.Client
}
func NewTokenService(rdb *redis.Client) *TokenService {
return &TokenService{rdb: rdb}
}
type TokenInfo struct {
MasterID uint
IssuedAtEpoch int64
Status string
Group string
}
// ValidateToken checks a child key against Redis for validity.
// This is designed to be called by the data plane (balancer).
func (s *TokenService) ValidateToken(ctx context.Context, token string) (*TokenInfo, error) {
tokenHash := util.HashToken(token)
tokenKey := fmt.Sprintf("auth:token:%s", tokenHash)
// 1. Get token metadata from Redis
tokenData, err := s.rdb.HGetAll(ctx, tokenKey).Result()
if err != nil {
return nil, fmt.Errorf("failed to get token data: %w", err)
}
if len(tokenData) == 0 {
return nil, errors.New("token not found")
}
if tokenData["status"] != "active" {
return nil, errors.New("token is not active")
}
masterID, _ := strconv.ParseUint(tokenData["master_id"], 10, 64)
issuedAtEpoch, _ := strconv.ParseInt(tokenData["issued_at_epoch"], 10, 64)
// 2. Get master metadata from Redis
masterKey := fmt.Sprintf("auth:master:%d", masterID)
masterEpochStr, err := s.rdb.HGet(ctx, masterKey, "epoch").Result()
if err != nil {
return nil, fmt.Errorf("failed to get master epoch: %w", err)
}
masterEpoch, _ := strconv.ParseInt(masterEpochStr, 10, 64)
// 3. Core Epoch Validation
if issuedAtEpoch < masterEpoch {
return nil, errors.New("token revoked due to master key rotation")
}
return &TokenInfo{
MasterID: uint(masterID),
IssuedAtEpoch: issuedAtEpoch,
Status: tokenData["status"],
Group: tokenData["group"],
}, nil
}

12
internal/util/hash.go Normal file
View File

@@ -0,0 +1,12 @@
package util
import (
"crypto/sha256"
"encoding/hex"
)
func HashToken(token string) string {
hasher := sha256.New()
hasher.Write([]byte(token))
return hex.EncodeToString(hasher.Sum(nil))
}