Primora/apps/backend/internal/middleware/context.go

package middleware

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/redis/go-redis/v9"

	"github.com/tdvorak/primora/apps/backend/internal/auth"
	db "github.com/tdvorak/primora/apps/backend/internal/database/db"
	"github.com/tdvorak/primora/apps/backend/internal/models"
	"github.com/tdvorak/primora/apps/backend/internal/repositories"
	apperrors "github.com/tdvorak/primora/apps/backend/internal/response"
)

const (
	requestIDKey = "request_id"
	actorKey     = "actor"
)

type AuthMiddleware struct {
	Queries    *repositories.CoreRepository
	Logger     *slog.Logger
	Redis      *redis.Client
	Verifier   *auth.Verifier
	RateLimits RateLimitConfig
	// RateLimitResolver overrides RateLimits with live instance settings when
	// set — it resolves "user" and "api_key" limits per request (cached).
	RateLimitResolver func(ctx context.Context, scope string) int
}

type RateLimitConfig struct {
	APIKeyPerMinute int
	UserPerMinute   int
}

func (m AuthMiddleware) rateLimit(ctx context.Context, scope string, fallback int) int {
	if m.RateLimitResolver == nil {
		return fallback
	}
	if v := m.RateLimitResolver(ctx, scope); v >= 0 {
		return v
	}
	return fallback
}

func RequestID() gin.HandlerFunc {
	return func(c *gin.Context) {
		requestID := c.Request.Header.Get("X-Request-ID")
		if requestID == "" {
			requestID = uuid.NewString()
		}
		c.Set(requestIDKey, requestID)
		c.Writer.Header().Set("X-Request-ID", requestID)
		c.Next()
	}
}

func Logger(logger *slog.Logger) gin.HandlerFunc {
	return func(c *gin.Context) {
		start := time.Now()
		c.Next()
		logger.Info("request_complete",
			"method", c.Request.Method,
			"path", c.Request.URL.Path,
			"status", c.Writer.Status(),
			"duration_ms", time.Since(start).Milliseconds(),
			"request_id", RequestIDFromContext(c),
			"client_ip", c.ClientIP(),
		)
	}
}

func (m AuthMiddleware) ResolveActor() gin.HandlerFunc {
	return func(c *gin.Context) {
		apiKey := strings.TrimSpace(c.GetHeader("X-API-Key"))
		if apiKey == "" {
			apiKey = strings.TrimSpace(c.GetHeader("X-Primora-Key"))
		}
		authz := strings.TrimSpace(c.GetHeader("Authorization"))
		// EventSource cannot set headers — the SSE streams accept credentials
		// as query params instead.
		sseStream := strings.HasSuffix(c.Request.URL.Path, "/telemetry/stream") ||
			strings.HasSuffix(c.Request.URL.Path, "/realtime/stream")
		if apiKey == "" && authz == "" && sseStream {
			apiKey = strings.TrimSpace(c.Query("api_key"))
			if apiKey == "" {
				if t := strings.TrimSpace(c.Query("token")); t != "" {
					authz = "Bearer " + t
				}
			}
		}
		if apiKey == "" && strings.HasPrefix(strings.ToLower(authz), "bearer ") {
			token := strings.TrimSpace(strings.TrimPrefix(authz, "Bearer"))
			if strings.HasPrefix(authz, "Bearer ") {
				token = strings.TrimSpace(strings.TrimPrefix(authz, "Bearer "))
			}
			if strings.HasPrefix(strings.ToLower(authz), "bearer ") {
				actor, err := m.resolveJWTActor(c.Request.Context(), token)
				if err != nil {
					apperrors.Abort(c, http.StatusUnauthorized, "invalid_token", err.Error())
					return
				}
				userIdentity := actor.AuthSubject
				if actor.UserID != nil {
					userIdentity = actor.UserID.String()
				}
				if !m.enforceRateLimit(c, "user", userIdentity, m.rateLimit(c.Request.Context(), "user", m.RateLimits.UserPerMinute), "User rate limit exceeded") {
					return
				}
				c.Set(actorKey, actor)
				c.Next()
				return
			}
		}

		if apiKey == "" {
			apiKey = strings.TrimSpace(strings.TrimPrefix(authz, "ApiKey "))
		}
		if apiKey != "" {
			actor, err := m.resolveAPIKeyActor(c.Request.Context(), apiKey)
			if err != nil {
				apperrors.Abort(c, http.StatusUnauthorized, "invalid_api_key", err.Error())
				return
			}
			if !m.enforceRateLimit(c, "apikey", actor.APIKeyPrefix, m.rateLimit(c.Request.Context(), "api_key", m.RateLimits.APIKeyPerMinute), "API key rate limit exceeded") {
				return
			}
			if need := requiredAPIKeyScope(c.Request.Method, c.Request.URL.Path); !actor.HasScope(need) {
				apperrors.Abort(c, http.StatusForbidden, "insufficient_scope", "api key lacks the required scope: "+need)
				return
			}
			c.Set(actorKey, actor)
		}
		c.Next()
	}
}

func (m AuthMiddleware) enforceRateLimit(c *gin.Context, scope, identity string, limit int, exceededMessage string) bool {
	if m.Redis == nil || limit <= 0 || identity == "" {
		return true
	}
	key := "ratelimit:" + scope + ":" + identity + ":" + time.Now().UTC().Format("200601021504")
	count, err := m.Redis.Incr(c.Request.Context(), key).Result()
	if err != nil {
		if m.Logger != nil {
			m.Logger.Warn("rate limit increment failed", "scope", scope, "error", err)
		}
		return true
	}
	if count == 1 {
		if err := m.Redis.Expire(c.Request.Context(), key, time.Minute).Err(); err != nil && m.Logger != nil {
			m.Logger.Warn("rate limit expiry update failed", "scope", scope, "error", err)
		}
	}
	ttl, err := m.Redis.TTL(c.Request.Context(), key).Result()
	if err != nil {
		if m.Logger != nil {
			m.Logger.Warn("rate limit ttl lookup failed", "scope", scope, "error", err)
		}
		ttl = time.Minute
	}
	if ttl <= 0 {
		ttl = time.Minute
	}
	resetSeconds := int((ttl + time.Second - 1) / time.Second)
	if resetSeconds < 1 {
		resetSeconds = 1
	}
	remaining := limit - int(count)
	if remaining < 0 {
		remaining = 0
	}
	c.Header("X-RateLimit-Limit", strconv.Itoa(limit))
	c.Header("X-RateLimit-Remaining", strconv.Itoa(remaining))
	c.Header("X-RateLimit-Reset", strconv.Itoa(resetSeconds))
	if count > int64(limit) {
		c.Header("Retry-After", strconv.Itoa(resetSeconds))
		apperrors.Abort(c, http.StatusTooManyRequests, "rate_limited", exceededMessage)
		return false
	}
	return true
}

func RequireActor(c *gin.Context) (*models.Actor, bool) {
	actor, ok := ActorFromContext(c)
	if !ok {
		apperrors.Abort(c, http.StatusUnauthorized, "authentication_required", "authentication is required")
		return nil, false
	}
	return actor, true
}

func ActorFromContext(c *gin.Context) (*models.Actor, bool) {
	value, ok := c.Get(actorKey)
	if !ok {
		return nil, false
	}
	actor, ok := value.(*models.Actor)
	return actor, ok
}

func RequestIDFromContext(c *gin.Context) string {
	value, ok := c.Get(requestIDKey)
	if !ok {
		return ""
	}
	requestID, _ := value.(string)
	return requestID
}

func (m AuthMiddleware) resolveJWTActor(ctx context.Context, token string) (*models.Actor, error) {
	if m.Verifier == nil {
		return nil, errors.New("jwt verifier not configured")
	}
	claims, err := m.Verifier.ParseToken(token)
	if err != nil {
		return nil, err
	}
	user, err := m.Queries.UpsertUser(ctx, db.UpsertUserParams{
		AuthSubject:   claims.Subject,
		Email:         strings.ToLower(claims.Email),
		Name:          claims.Name,
		EmailVerified: claims.EmailVerified,
	})
	if err != nil {
		return nil, fmt.Errorf("upsert user: %w", err)
	}
	return &models.Actor{
		Type:          models.ActorTypeUser,
		UserID:        &user.ID,
		AuthSubject:   claims.Subject,
		Email:         strings.ToLower(claims.Email),
		EmailVerified: claims.EmailVerified,
		Name:          claims.Name,
		SessionID:     claims.SessionID,
	}, nil
}

func (m AuthMiddleware) resolveAPIKeyActor(ctx context.Context, rawKey string) (*models.Actor, error) {
	parts := strings.Split(rawKey, "_")
	if len(parts) < 3 {
		return nil, errors.New("malformed api key")
	}
	prefix := strings.Join(parts[:2], "_")
	row, err := m.Queries.GetAPIKeyByPrefix(ctx, prefix)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, errors.New("unknown api key")
		}
		return nil, err
	}
	if row.RevokedAt.Valid {
		return nil, errors.New("api key revoked")
	}
	sum := sha256.Sum256([]byte(rawKey))
	if subtle.ConstantTimeCompare(row.SecretHash, sum[:]) != 1 {
		return nil, errors.New("invalid api key secret")
	}
	if err := m.Queries.TouchAPIKey(ctx, row.ID); err != nil {
		m.Logger.Warn("failed to touch api key", "error", err)
	}
	projectID := row.ProjectID
	orgID := row.OrganizationID
	apiKeyID := row.ID
	return &models.Actor{
		Type:           models.ActorTypeAPIKey,
		ProjectID:      &projectID,
		OrganizationID: &orgID,
		APIKeyID:       &apiKeyID,
		APIKeyPrefix:   prefix,
		Scopes:         row.Scopes,
	}, nil
}

// requiredAPIKeyScope maps a request to the scope an API key must hold.
// admin-only surfaces: org management, invitations, bootstrap, project
// settings/members, and API key minting (a write key must never mint keys).
func requiredAPIKeyScope(method, path string) string {
	switch method {
	case http.MethodGet, http.MethodHead, http.MethodOptions:
		return "read"
	}
	if path == "/api/v1/ingest" {
		return "ingest"
	}
	rest := strings.TrimPrefix(path, "/api/v1/")
	segs := strings.SplitN(rest, "/", 3)
	switch segs[0] {
	case "organizations", "invitations", "bootstrap":
		return "admin"
	case "projects":
		if len(segs) < 3 ||
			strings.HasPrefix(segs[2], "api-keys") ||
			strings.HasPrefix(segs[2], "members") {
			return "admin"
		}
	}
	return "write"
}

func PgUUID(id uuid.UUID) pgtype.UUID {
	return pgtype.UUID{Bytes: id, Valid: true}
}

func HexDigest(input string) string {
	sum := sha256.Sum256([]byte(input))
	return hex.EncodeToString(sum[:])
}