AIRouter/internal/auth/auth.go

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