AIRouter/internal/auth/auth.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/edgeai/gateway/internal/handler"
	"github.com/edgeai/gateway/internal/middleware"
	"github.com/edgeai/gateway/internal/observability"
	_ "github.com/mattn/go-sqlite3"
)

// AppIdentity represents the authenticated application identity.
type AppIdentity struct {
	AppID             string
	TenantID          string
	Name              string
	AllowedModels     []string
	AllowedPriorities []int
	IsAdmin           bool
	ExpiresAt         *time.Time // 过期时间，nil 表示永不过期
}

type contextKey string

const (
	AppIdentityKey contextKey = "app_identity"
)

// Authenticator manages API Key authentication.
type Authenticator struct {
	mu     sync.RWMutex
	keys   map[string]*AppIdentity // hashed_key -> identity
	db     *sql.DB
	logger *observability.Logger
	stopCh chan struct{} // 停止清理 goroutine
}

// NewAuthenticator creates a new Authenticator with SQLite storage.
func NewAuthenticator(dbPath string, logger *observability.Logger) (*Authenticator, error) {
	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, fmt.Errorf("open auth db: %w", err)
	}

	if err := initAuthDB(db); err != nil {
		return nil, fmt.Errorf("init auth db: %w", err)
	}

	a := &Authenticator{
		keys:   make(map[string]*AppIdentity),
		db:     db,
		logger: logger,
		stopCh: make(chan struct{}),
	}

	if err := a.loadKeys(); err != nil {
		return nil, fmt.Errorf("load api keys: %w", err)
	}

	go a.cleanupExpiredKeys()

	return a, nil
}

func initAuthDB(db *sql.DB) error {
	schema := `
	CREATE TABLE IF NOT EXISTS api_keys (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		app_id TEXT NOT NULL,
		tenant_id TEXT NOT NULL,
		name TEXT NOT NULL,
		key_hash TEXT NOT NULL UNIQUE,
		allowed_models TEXT, -- JSON array, empty = all
		allowed_priorities TEXT, -- JSON array, empty = all
		is_admin INTEGER DEFAULT 0,
		enabled INTEGER DEFAULT 1,
		created_at TEXT NOT NULL DEFAULT (datetime('now')),
		expires_at TEXT
	);`
	_, err := db.Exec(schema)
	return err
}

func (a *Authenticator) loadKeys() error {
	rows, err := a.db.Query(`SELECT key_hash, app_id, tenant_id, name, allowed_models, allowed_priorities, is_admin, expires_at FROM api_keys WHERE enabled = 1`)
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var hash, appID, tenantID, name, allowedModelsJSON, allowedPrioritiesJSON string
		var isAdmin int
		var expiresAtStr sql.NullString
		if err := rows.Scan(&hash, &appID, &tenantID, &name, &allowedModelsJSON, &allowedPrioritiesJSON, &isAdmin, &expiresAtStr); err != nil {
			return err
		}

		identity := &AppIdentity{
			AppID:    appID,
			TenantID: tenantID,
			Name:     name,
			IsAdmin:  isAdmin == 1,
		}
		if expiresAtStr.Valid && expiresAtStr.String != "" {
			if t, err := time.Parse(time.RFC3339, expiresAtStr.String); err == nil {
				identity.ExpiresAt = &t
			}
		}
		if allowedModelsJSON != "" && allowedModelsJSON != "null" {
			json.Unmarshal([]byte(allowedModelsJSON), &identity.AllowedModels)
		}
		if allowedPrioritiesJSON != "" && allowedPrioritiesJSON != "null" {
			json.Unmarshal([]byte(allowedPrioritiesJSON), &identity.AllowedPriorities)
		}

		a.keys[hash] = identity
	}
	return rows.Err()
}

// GenerateAPIKey 生成一个随机的 API Key（前缀 edgeai- + 32 字节随机十六进制）。
func GenerateAPIKey() string {
	b := make([]byte, 32)
	rand.Read(b)
	return "edgeai-" + hex.EncodeToString(b)
}

// hashKey hashes an API key with SHA-256.
func hashKey(key string) string {
	h := sha256.Sum256([]byte(key))
	return hex.EncodeToString(h[:])
}

// Authenticate validates an API key and returns the AppIdentity.
// 检查 Key 是否存在且未过期。
func (a *Authenticator) Authenticate(apiKey string) (*AppIdentity, bool) {
	hash := hashKey(apiKey)
	a.mu.RLock()
	defer a.mu.RUnlock()
	identity, ok := a.keys[hash]
	if !ok {
		return nil, false
	}
	// 检查过期
	if identity.ExpiresAt != nil && time.Now().After(*identity.ExpiresAt) {
		return nil, false
	}
	return identity, true
}

// AddKey adds a new API key (for management API).
// 如果 identity.ExpiresAt 不为 nil，则设置过期时间。
func (a *Authenticator) AddKey(apiKey string, identity *AppIdentity) error {
	hash := hashKey(apiKey)
	allowedModelsJSON, _ := json.Marshal(identity.AllowedModels)
	allowedPrioritiesJSON, _ := json.Marshal(identity.AllowedPriorities)

	var expiresAt interface{}
	if identity.ExpiresAt != nil {
		expiresAt = identity.ExpiresAt.Format(time.RFC3339)
	}

	_, err := a.db.Exec(
		`INSERT INTO api_keys (app_id, tenant_id, name, key_hash, allowed_models, allowed_priorities, is_admin, enabled, expires_at)
		 VALUES (?, ?, ?, ?, ?, ?, ?, 1, ?)`,
		identity.AppID, identity.TenantID, identity.Name, hash, string(allowedModelsJSON), string(allowedPrioritiesJSON), isAdminInt(identity.IsAdmin), expiresAt,
	)
	if err != nil {
		return err
	}

	a.mu.Lock()
	a.keys[hash] = identity
	a.mu.Unlock()
	return nil
}

// KeyInfo 是 API Key 的元信息（不含哈希值）。
// 用于管理 API 列出所有 Key。
type KeyInfo struct {
	ID                int64    `json:"id"`
	AppID             string   `json:"app_id"`
	TenantID          string   `json:"tenant_id"`
	Name              string   `json:"name"`
	AllowedModels     []string `json:"allowed_models"`
	AllowedPriorities []int    `json:"allowed_priorities"`
	IsAdmin           bool     `json:"is_admin"`
	Enabled           bool     `json:"enabled"`
	CreatedAt         string   `json:"created_at"`
	ExpiresAt         string   `json:"expires_at,omitempty"`
}

// ListKeys 列出所有 API Key 的元信息（不含哈希值）。
func (a *Authenticator) ListKeys() ([]KeyInfo, error) {
	rows, err := a.db.Query(
		`SELECT id, app_id, tenant_id, name, allowed_models, allowed_priorities, is_admin, enabled, created_at, expires_at FROM api_keys ORDER BY id DESC`,
	)
	if err != nil {
		return nil, fmt.Errorf("query api keys: %w", err)
	}
	defer rows.Close()

	var keys []KeyInfo
	for rows.Next() {
		var ki KeyInfo
		var allowedModelsJSON, allowedPrioritiesJSON string
		var isAdmin, enabled int
		var expiresAt sql.NullString
		if err := rows.Scan(&ki.ID, &ki.AppID, &ki.TenantID, &ki.Name, &allowedModelsJSON, &allowedPrioritiesJSON, &isAdmin, &enabled, &ki.CreatedAt, &expiresAt); err != nil {
			return nil, err
		}
		ki.IsAdmin = isAdmin == 1
		ki.Enabled = enabled == 1
		if expiresAt.Valid && expiresAt.String != "" {
			ki.ExpiresAt = expiresAt.String
		}
		if allowedModelsJSON != "" && allowedModelsJSON != "null" {
			json.Unmarshal([]byte(allowedModelsJSON), &ki.AllowedModels)
		}
		if allowedPrioritiesJSON != "" && allowedPrioritiesJSON != "null" {
			json.Unmarshal([]byte(allowedPrioritiesJSON), &ki.AllowedPriorities)
		}
		keys = append(keys, ki)
	}
	return keys, rows.Err()
}

// DisableKey 禁用一个 API Key（通过数据库 ID）。
func (a *Authenticator) DisableKey(id int64) error {
	// 先查出 hash 以便从内存中移除
	var hash string
	err := a.db.QueryRow(`SELECT key_hash FROM api_keys WHERE id = ?`, id).Scan(&hash)
	if err == sql.ErrNoRows {
		return fmt.Errorf("api key not found: %d", id)
	}
	if err != nil {
		return fmt.Errorf("query api key: %w", err)
	}

	_, err = a.db.Exec(`UPDATE api_keys SET enabled = 0 WHERE id = ?`, id)
	if err != nil {
		return fmt.Errorf("disable api key: %w", err)
	}

	a.mu.Lock()
	delete(a.keys, hash)
	a.mu.Unlock()
	return nil
}

// DeleteKey 彻底删除一个 API Key（通过数据库 ID）。
func (a *Authenticator) DeleteKey(id int64) error {
	var hash string
	err := a.db.QueryRow(`SELECT key_hash FROM api_keys WHERE id = ?`, id).Scan(&hash)
	if err == sql.ErrNoRows {
		return fmt.Errorf("api key not found: %d", id)
	}
	if err != nil {
		return fmt.Errorf("query api key: %w", err)
	}

	_, err = a.db.Exec(`DELETE FROM api_keys WHERE id = ?`, id)
	if err != nil {
		return fmt.Errorf("delete api key: %w", err)
	}

	a.mu.Lock()
	delete(a.keys, hash)
	a.mu.Unlock()
	return nil
}

// RotateKey 轮换 API Key：生成新 Key，禁用旧 Key，返回新 Key 值。
func (a *Authenticator) RotateKey(id int64) (string, error) {
	var oldHash string
	var appID, tenantID, name, allowedModelsJSON, allowedPrioritiesJSON string
	var isAdmin int
	var expiresAt sql.NullString
	err := a.db.QueryRow(
		`SELECT key_hash, app_id, tenant_id, name, allowed_models, allowed_priorities, is_admin, expires_at FROM api_keys WHERE id = ?`,
		id,
	).Scan(&oldHash, &appID, &tenantID, &name, &allowedModelsJSON, &allowedPrioritiesJSON, &isAdmin, &expiresAt)
	if err == sql.ErrNoRows {
		return "", fmt.Errorf("api key not found: %d", id)
	}
	if err != nil {
		return "", fmt.Errorf("query api key: %w", err)
	}

	newKey := GenerateAPIKey()
	newHash := hashKey(newKey)

	// 插入新 Key，继承旧 Key 的所有属性
	_, err = a.db.Exec(
		`INSERT INTO api_keys (app_id, tenant_id, name, key_hash, allowed_models, allowed_priorities, is_admin, enabled, expires_at)
		 VALUES (?, ?, ?, ?, ?, ?, ?, 1, ?)`,
		appID, tenantID, name+" (rotated)", newHash, allowedModelsJSON, allowedPrioritiesJSON, isAdmin, expiresAt,
	)
	if err != nil {
		return "", fmt.Errorf("insert rotated key: %w", err)
	}

	// 禁用旧 Key
	_, err = a.db.Exec(`UPDATE api_keys SET enabled = 0 WHERE id = ?`, id)
	if err != nil {
		return "", fmt.Errorf("disable old key: %w", err)
	}

	// 更新内存：移除旧 Key，添加新 Key
	identity := &AppIdentity{
		AppID:    appID,
		TenantID: tenantID,
		Name:     name + " (rotated)",
		IsAdmin:  isAdmin == 1,
	}
	if allowedModelsJSON != "" && allowedModelsJSON != "null" {
		json.Unmarshal([]byte(allowedModelsJSON), &identity.AllowedModels)
	}
	if allowedPrioritiesJSON != "" && allowedPrioritiesJSON != "null" {
		json.Unmarshal([]byte(allowedPrioritiesJSON), &identity.AllowedPriorities)
	}
	if expiresAt.Valid && expiresAt.String != "" {
		if t, err := time.Parse(time.RFC3339, expiresAt.String); err == nil {
			identity.ExpiresAt = &t
		}
	}

	a.mu.Lock()
	delete(a.keys, oldHash)
	a.keys[newHash] = identity
	a.mu.Unlock()

	return newKey, nil
}

// Middleware returns an HTTP middleware that enforces API Key authentication.
func (a *Authenticator) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Skip auth for health/ready endpoints
		if r.URL.Path == "/health" || r.URL.Path == "/ready" || r.URL.Path == "/metrics" {
			next.ServeHTTP(w, r)
			return
		}

		authHeader := r.Header.Get("Authorization")
		if authHeader == "" {
			handler.WriteError(w, handler.NewGatewayError(handler.ErrAuthFailed, "missing Authorization header"))
			return
		}

		parts := strings.SplitN(authHeader, " ", 2)
		if len(parts) != 2 || parts[0] != "Bearer" {
			handler.WriteError(w, handler.NewGatewayError(handler.ErrAuthFailed, "invalid Authorization format, expected Bearer <api_key>"))
			return
		}

		apiKey := parts[1]
		if apiKey == "" {
			handler.WriteError(w, handler.NewGatewayError(handler.ErrAuthFailed, "empty API key"))
			return
		}

		identity, ok := a.Authenticate(apiKey)
		if !ok {
			handler.WriteError(w, handler.NewGatewayError(handler.ErrAuthFailed, "invalid API key"))
			return
		}

		ctx := context.WithValue(r.Context(), AppIdentityKey, identity)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// GetAppIdentity extracts the AppIdentity from request context.
func GetAppIdentity(ctx context.Context) *AppIdentity {
	if v, ok := ctx.Value(AppIdentityKey).(*AppIdentity); ok {
		return v
	}
	return nil
}

// GetAppIdentityFromRequest is a convenience wrapper.
func GetAppIdentityFromRequest(r *http.Request) *AppIdentity {
	return GetAppIdentity(r.Context())
}

// RequireAdmin checks if the request is from an admin app.
func RequireAdmin(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		identity := GetAppIdentityFromRequest(r)
		if identity == nil || !identity.IsAdmin {
			handler.WriteError(w, handler.NewGatewayError(handler.ErrPermissionDenied, "admin access required"))
			return
		}
		next.ServeHTTP(w, r)
	})
}

// CheckModelPermission verifies the app can access the given model.
// 支持通配符 "*" 匹配所有模型。
func CheckModelPermission(identity *AppIdentity, model string) bool {
	if len(identity.AllowedModels) == 0 {
		return true // empty = all models allowed
	}
	for _, m := range identity.AllowedModels {
		if m == model || m == "*" {
			return true
		}
	}
	return false
}

// CheckPriorityPermission verifies the app can use the given priority.
func CheckPriorityPermission(identity *AppIdentity, priority int) bool {
	if len(identity.AllowedPriorities) == 0 {
		return true
	}
	for _, p := range identity.AllowedPriorities {
		if p == priority {
			return true
		}
	}
	return false
}

func isAdminInt(b bool) int {
	if b {
		return 1
	}
	return 0
}

// Close closes the database connection and stops background cleanup.
func (a *Authenticator) Close() error {
	close(a.stopCh)
	return a.db.Close()
}

// cleanupExpiredKeys 定期清理内存中已过期的 API Key，防止内存泄漏。
func (a *Authenticator) cleanupExpiredKeys() {
	ticker := time.NewTicker(5 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			now := time.Now()
			a.mu.Lock()
			for hash, identity := range a.keys {
				if identity.ExpiresAt != nil && now.After(*identity.ExpiresAt) {
					delete(a.keys, hash)
				}
			}
			a.mu.Unlock()
		case <-a.stopCh:
			return
		}
	}
}

// Ensure middleware import is used.
var _ = middleware.GetRequestID