archivmail/internal/auth/apikey_middleware.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"
)

// APIKeySession holds the resolved API-key context for downstream handlers.
type APIKeySession struct {
	KeyID    int64
	KeyName  string
	TenantID int64
	Role     string
}

type apiKeyContextKey string

const apiKeySessionKey apiKeyContextKey = "apikey_session"

// APIKeySessionFromCtx extracts the API-key session from the request context.
func APIKeySessionFromCtx(ctx context.Context) *APIKeySession {
	v, _ := ctx.Value(apiKeySessionKey).(*APIKeySession)
	return v
}

// APIKeyLookup is the interface the middleware uses to resolve a hashed token.
type APIKeyLookup interface {
	// LookupAPIKey returns the key metadata for the given SHA-256 hex hash.
	// Returns nil if not found or inactive.
	LookupAPIKey(ctx context.Context, tokenHash string) (*APIKeyRow, error)
	// TouchAPIKeyLastUsed updates last_used_at for the given key ID.
	TouchAPIKeyLastUsed(ctx context.Context, keyID int64) error
}

// APIKeyRow holds a single API key record from the database.
type APIKeyRow struct {
	ID        int64
	TenantID  int64
	Name      string
	Role      string
	Active    bool
	RateLimit int
}

// tokenBucket implements a simple per-key token-bucket rate limiter.
type tokenBucket struct {
	tokens    float64
	limit     float64
	lastCheck time.Time
}

// APIKeyMiddleware returns an http middleware that authenticates requests
// via "Authorization: Bearer am_<token>" and enforces per-key rate limits.
type APIKeyMiddleware struct {
	lookup  APIKeyLookup
	buckets sync.Map // map[int64]*tokenBucket
}

// NewAPIKeyMiddleware creates a new API-key authentication middleware.
func NewAPIKeyMiddleware(lookup APIKeyLookup) *APIKeyMiddleware {
	return &APIKeyMiddleware{lookup: lookup}
}

// Wrap returns an http.HandlerFunc that performs API-key auth before calling next.
func (m *APIKeyMiddleware) Wrap(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		token := extractAPIKeyToken(r)
		if token == "" {
			writeAPIError(w, http.StatusUnauthorized, "missing or invalid authorization")
			return
		}

		// Compute SHA-256 hash of the raw token.
		hash := sha256.Sum256([]byte(token))
		tokenHash := hex.EncodeToString(hash[:])

		row, err := m.lookup.LookupAPIKey(r.Context(), tokenHash)
		if err != nil {
			writeAPIError(w, http.StatusInternalServerError, "internal error")
			return
		}
		if row == nil || !row.Active {
			writeAPIError(w, http.StatusUnauthorized, "invalid or inactive API key")
			return
		}

		// Rate limiting (token bucket).
		if !m.allow(row.ID, row.RateLimit) {
			w.Header().Set("Retry-After", "30")
			writeAPIError(w, http.StatusTooManyRequests, "rate limit exceeded")
			return
		}

		// Update last_used_at (best-effort, do not block request).
		go func() {
			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
			defer cancel()
			_ = m.lookup.TouchAPIKeyLastUsed(ctx, row.ID)
		}()

		// Inject API-key session into context.
		sess := &APIKeySession{
			KeyID:    row.ID,
			KeyName:  row.Name,
			TenantID: row.TenantID,
			Role:     row.Role,
		}
		ctx := context.WithValue(r.Context(), apiKeySessionKey, sess)
		next(w, r.WithContext(ctx))
	}
}

// allow checks and updates the token-bucket for the given key.
func (m *APIKeyMiddleware) allow(keyID int64, limitPerMin int) bool {
	now := time.Now()
	limit := float64(limitPerMin)

	val, _ := m.buckets.LoadOrStore(keyID, &tokenBucket{
		tokens:    limit,
		limit:     limit,
		lastCheck: now,
	})
	bucket := val.(*tokenBucket)

	elapsed := now.Sub(bucket.lastCheck).Seconds()
	bucket.lastCheck = now

	// Refill tokens based on elapsed time.
	bucket.tokens += elapsed * (limit / 60.0)
	if bucket.tokens > limit {
		bucket.tokens = limit
	}

	if bucket.tokens < 1.0 {
		return false
	}
	bucket.tokens -= 1.0
	return true
}

// extractAPIKeyToken extracts the token from "Authorization: Bearer am_<token>".
func extractAPIKeyToken(r *http.Request) string {
	h := r.Header.Get("Authorization")
	if !strings.HasPrefix(h, "Bearer ") {
		return ""
	}
	token := strings.TrimPrefix(h, "Bearer ")
	if !strings.HasPrefix(token, "am_") {
		return ""
	}
	return token
}

// GenerateAPIKey creates a new random API key with the "am_" prefix.
// Returns the raw token (to show once) and its SHA-256 hex hash (to store).
func GenerateAPIKey() (rawToken string, tokenHash string, err error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", "", fmt.Errorf("auth: generate api key: %w", err)
	}
	rawToken = "am_" + base64.RawURLEncoding.EncodeToString(b)
	hash := sha256.Sum256([]byte(rawToken))
	tokenHash = hex.EncodeToString(hash[:])
	return rawToken, tokenHash, nil
}

// writeAPIError writes a JSON error response. Duplicated here to avoid
// importing the api package (which would create a circular dependency).
func writeAPIError(w http.ResponseWriter, code int, msg string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	fmt.Fprintf(w, `{"error":%q}`, msg)
}