初始提交：边缘AI算力机统一AI通讯层

2026-08-03 07:44:05 +08:00
commit 93a469061d
51 changed files with 11565 additions and 0 deletions
@@ -0,0 +1,260 @@
+package auth
+
+import (
+	"context"
+	"crypto/sha256"
+	"database/sql"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"strings"
+	"sync"
+
+	"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
+}
+
+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
+}
+
+// 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,
+	}
+
+	if err := a.loadKeys(); err != nil {
+		return nil, fmt.Errorf("load api keys: %w", err)
+	}
+
+	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 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
+		if err := rows.Scan(&hash, &appID, &tenantID, &name, &allowedModelsJSON, &allowedPrioritiesJSON, &isAdmin); err != nil {
+			return err
+		}
+
+		identity := &AppIdentity{
+			AppID:    appID,
+			TenantID: tenantID,
+			Name:     name,
+			IsAdmin:  isAdmin == 1,
+		}
+		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()
+}
+
+// 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.
+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
+	}
+	return identity, true
+}
+
+// AddKey adds a new API key (for management API).
+func (a *Authenticator) AddKey(apiKey string, identity *AppIdentity) error {
+	hash := hashKey(apiKey)
+	allowedModelsJSON, _ := json.Marshal(identity.AllowedModels)
+	allowedPrioritiesJSON, _ := json.Marshal(identity.AllowedPriorities)
+
+	_, err := a.db.Exec(
+		`INSERT INTO api_keys (app_id, tenant_id, name, key_hash, allowed_models, allowed_priorities, is_admin, enabled)
+		 VALUES (?, ?, ?, ?, ?, ?, ?, 1)`,
+		identity.AppID, identity.TenantID, identity.Name, hash, string(allowedModelsJSON), string(allowedPrioritiesJSON), isAdminInt(identity.IsAdmin),
+	)
+	if err != nil {
+		return err
+	}
+
+	a.mu.Lock()
+	a.keys[hash] = identity
+	a.mu.Unlock()
+	return 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 {
+			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.
+func (a *Authenticator) Close() error {
+	return a.db.Close()
+}
+
+// Ensure middleware import is used.
+var _ = middleware.GetRequestID