plugin.go

package token

import (
	"bufio"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"slices"
	"strings"

	"github.com/avast/retry-go/v4"
	caddy "github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
	"github.com/coreos/go-oidc/v3/oidc"
	"github.com/fsnotify/fsnotify"
	"github.com/golang-jwt/jwt/v5"
	"github.com/loafoe/caddy-token/keys"
	"go.uber.org/zap"
)

const (
	scopeIDHeader    = "X-Scope-OrgID"
	apiKeyHeader     = "X-Api-Key"
	tokenKeyHeader   = "X-Id-Token"
	grafanaOrgHeader = "X-Grafana-Org-Id"
	authHeader       = "Authorization"
)

type Middleware struct {
	logger            *zap.Logger
	TokenFile         string
	tokens            map[string]keys.Key
	Issuer            string
	InjectOrgHeader   bool
	AllowUpstreamAuth bool
	Verify            bool
	verifier          *oidc.IDTokenVerifier
	watcher           *fsnotify.Watcher
	TenantOrgClaim    string
	SigningKey        string
	Groups            []string
	Scopes            []string
	ClientCA          bool
	Debug             bool
	DefaultOrg        string
	Spiffe            *SpiffeConfig
	spiffeValidator   *SpiffeValidator
}

func (m *Middleware) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "http.handlers.token",
		New: func() caddy.Module { return new(Middleware) },
	}
}

func (m *Middleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
	err := m.CheckTokenAndInjectHeaders(r)
	if err != nil {
		return err
	}
	return next.ServeHTTP(w, r)
}

func (m *Middleware) Validate() error {
	return nil
}

func (m *Middleware) Provision(ctx caddy.Context) error {
	var err error
	m.logger = ctx.Logger() // g.logger is a *zap.Logger
	// Create new watcher.
	m.watcher, err = fsnotify.NewWatcher()
	if err != nil {
		return fmt.Errorf("error creating watcher: %w", err)
	}
	//defer watcher.Close()

	if m.Issuer != "" {
		provider, err := oidc.NewProvider(ctx, m.Issuer)
		if err != nil {
			m.logger.Info("error provisioning issuer", zap.String("issuer", m.Issuer), zap.Error(err))
			return fmt.Errorf("erorr provisioning issuer '%s': %w", m.Issuer, err)
		}
		m.verifier = provider.Verifier(&oidc.Config{
			SkipClientIDCheck: true,
		})
		if !m.Verify {
			m.logger.Warn("jwt 'verify false' no longer disables signature verification; tokens with invalid signatures are always rejected",
				zap.String("issuer", m.Issuer))
		}
		m.logger.Info("verifier setup", zap.String("issuer", m.Issuer))
	}
	if m.TokenFile != "" {
		tokens, err := m.readTokenFile(m.TokenFile)
		if err != nil {
			return err
		}
		err = m.watcher.Add(m.TokenFile)
		if err != nil {
			return fmt.Errorf("error watching token file: %w", err)
		}
		m.tokens = tokens
	}
	if m.Spiffe != nil && len(m.Spiffe.TrustDomains) > 0 {
		socketPath := m.Spiffe.GetWorkloadSocket()
		// Count trust domains with/without JWKS URLs
		jwksCount := 0
		workloadCount := 0
		for _, td := range m.Spiffe.TrustDomains {
			if td.JWKSURL != "" {
				jwksCount++
			} else {
				workloadCount++
			}
		}

		if socketPath != "" || jwksCount > 0 {
			// Use hybrid source: JWKS for domains with URLs, Workload API for others
			var err error
			m.spiffeValidator, err = NewSpiffeValidatorWithWorkloadAPI(ctx, m.Spiffe, m.logger)
			if err != nil {
				return fmt.Errorf("creating SPIFFE validator: %w", err)
			}
			m.logger.Info("SPIFFE validator configured",
				zap.String("socket", socketPath),
				zap.Int("trustDomains", len(m.Spiffe.TrustDomains)),
				zap.Int("jwksDomains", jwksCount),
				zap.Int("workloadDomains", workloadCount),
				zap.Int("allowedIDs", len(m.Spiffe.AllowedIDs)))
		} else {
			// No socket and no JWKS URLs - this is a configuration error
			return fmt.Errorf("SPIFFE configured but no workload socket or JWKS URLs provided")
		}
	}
	if m.verifier == nil && len(m.tokens) == 0 && m.SigningKey == "" && !m.ClientCA && m.spiffeValidator == nil {
		return fmt.Errorf("no tokens, issuer, client CA, or SPIFFE config provided")
	}
	m.logger.Info("provisioned caddy-token middleware",
		zap.String("issuer", m.Issuer),
		zap.String("tokenFile", m.TokenFile),
		zap.Int64("apiKeyCount", int64(len(m.tokens))),
		zap.String("TenantOrgClaim", m.TenantOrgClaim),
		zap.Bool("HasSigningKey", m.SigningKey != ""),
		zap.Bool("AllowUpstreamAuth", m.AllowUpstreamAuth),
		zap.Bool("HasSpiffe", m.spiffeValidator != nil))
	// start watching tokenFile
	if m.TokenFile != "" {
		m.logger.Info("starting watcher for token file", zap.String("tokenFile", m.TokenFile))
		// Start listening for events
		go func() {
			for {
				select {
				case event, ok := <-m.watcher.Events:
					if !ok {
						return
					}
					if event.Has(fsnotify.Write) || event.Has(fsnotify.Create) || event.Has(fsnotify.Rename) || event.Has(fsnotify.Remove) {
						tokens := make(map[string]keys.Key)
						err = retry.Do(func() error {
							tokens, err = m.readTokenFile(m.TokenFile)
							return err
						}, retry.Attempts(5), retry.Delay(1))
						if err != nil {
							m.logger.Error("error reloading token file", zap.Error(err))
						} else {
							m.tokens = tokens
							m.logger.Info("reloaded token file", zap.Int("apiKeyCount", len(m.tokens)))
						}
					}
					if event.Has(fsnotify.Rename) || event.Has(fsnotify.Remove) { // Re-add
						_ = m.watcher.Remove(m.TokenFile)
						err = retry.Do(func() error {
							return m.watcher.Add(m.TokenFile)
						}, retry.Attempts(5), retry.Delay(1))
						if err != nil {
							m.logger.Error("error re-adding watcher", zap.Error(err))
						}
					}
				case err, ok := <-m.watcher.Errors:
					if !ok {
						return
					}
					m.logger.Error("watcher error", zap.Error(err))
				}
			}
		}()
	} else {
		m.logger.Info("no token file to watch")
	}
	return nil
}

func (m *Middleware) CheckTokenAndInjectHeaders(r *http.Request) error {
	grafanaOrgId := r.Header.Get(grafanaOrgHeader)
	idToken := r.Header.Get(tokenKeyHeader)
	apiKey := r.Header.Get(apiKeyHeader)
	// Check for upstream auth
	upstreamAuth := r.Header.Get(scopeIDHeader)
	if grafanaOrgId != "" && m.Debug {
		m.logger.Info("Grafana Org context detected", zap.String("value", grafanaOrgId))
	}
	if upstreamAuth != "" {
		if m.Debug {
			m.logger.Info("upstream X-Scope-OrgID detected", zap.String("value", upstreamAuth))
		}
		if m.AllowUpstreamAuth {
			return nil
		}
		if m.Debug {
			m.logger.Info("ignoring upstream X-Scope-OrgID", zap.Bool("AllowUpstreamAuth", m.AllowUpstreamAuth))
		}
	}

	// Strip any client-supplied tenant headers before authenticating. Unless
	// AllowUpstreamAuth is enabled (handled above), these headers must never be
	// trusted: leaving them in place would let a caller spoof a tenant on any
	// success path that does not explicitly overwrite X-Scope-OrgID (e.g.
	// InjectOrgHeader=false, or a JWT/SPIFFE org that resolves to empty).
	r.Header.Del(scopeIDHeader)
	r.Header.Del(grafanaOrgHeader)

	// Check for client certificate authentication first. We require a verified
	// chain (r.TLS.VerifiedChains), not merely a presented certificate: a bare
	// PeerCertificates entry can be any self-signed cert. Verified chains are
	// only populated when the TLS listener is configured with
	// require_and_verify against a trusted client CA.
	if m.ClientCA && r.TLS != nil && len(r.TLS.PeerCertificates) > 0 && len(r.TLS.VerifiedChains) > 0 {
		if m.Debug {
			m.logger.Info("client certificate detected",
				zap.Int("certCount", len(r.TLS.PeerCertificates)),
				zap.String("subject", r.TLS.PeerCertificates[0].Subject.String()))
		}
		// Set the default organization header
		r.Header.Set(scopeIDHeader, m.DefaultOrg)
		if m.Debug {
			m.logger.Info("client certificate authenticated", zap.String("defaultOrg", m.DefaultOrg))
		}
		return nil
	}

	// Check for SPIFFE JWT SVID in Authorization header
	if m.spiffeValidator != nil {
		authHeader := r.Header.Get(authHeader)
		if authHeader != "" {
			parts := strings.Split(authHeader, " ")
			if len(parts) == 2 && strings.ToLower(parts[0]) == "bearer" {
				result, err := m.spiffeValidator.ValidateJWT(r.Context(), parts[1])
				if err == nil {
					// Valid SPIFFE JWT
					if m.InjectOrgHeader && result.Org != "" {
						r.Header.Set(scopeIDHeader, result.Org)
					}
					if m.Debug {
						m.logger.Debug("SPIFFE JWT authenticated",
							zap.String("spiffeID", result.SpiffeID.String()),
							zap.String("org", result.Org))
					}
					return nil
				}
				// If SPIFFE validation failed but we have other auth methods, continue
				// Otherwise, if SPIFFE is the only method, return the error
				if m.verifier == nil && len(m.tokens) == 0 && m.SigningKey == "" {
					if m.Debug {
						m.logger.Error("SPIFFE JWT validation failed", zap.Error(err))
					}
					return caddyhttp.Error(http.StatusUnauthorized, err)
				}
				// Log but continue to try other auth methods
				if m.Debug {
					m.logger.Debug("SPIFFE JWT validation failed, trying other auth methods", zap.Error(err))
				}
			}
		}
	}

	// Check if API key is there in header
	// Try to extract token from Basic Auth
	_, password, ok := r.BasicAuth()
	if ok && password != "" {
		apiKey = password
	}
	// Also support bearer token
	if apiKey == "" {
		authHeader := r.Header.Get(authHeader)
		if authHeader != "" {
			parts := strings.Split(authHeader, " ")
			if len(parts) == 2 && strings.ToLower(parts[0]) == "bearer" {
				apiKey = parts[1]
			}
		}
	}
	if apiKey != "" { // API Key flow
		// Check v2 API keys first
		if ok, token, _ := keys.VerifyAPIKey(apiKey, m.SigningKey); ok {
			r.Header.Set(scopeIDHeader, token.Organization)
			return nil
		}
		token, ok := m.tokens[apiKey]
		if !ok {
			if m.Debug {
				m.logger.Error("invalid token detected",
					zap.String("apiKey", "..."+LastNChars(6, apiKey)),
					zap.Int64("count", int64(len(m.tokens))),
					zap.String("remoteAddr", r.RemoteAddr))
			}
			return caddyhttp.Error(http.StatusForbidden, nil)
		}

		// Check scopes
		if len(m.Scopes) > 0 {
			for _, scope := range m.Scopes {
				if !slices.Contains(token.Scopes, scope) {
					m.logger.Error("missing required scope",
						zap.String("scope", scope),
						zap.String("remoteAddr", r.RemoteAddr))
					return caddyhttp.Error(http.StatusForbidden, nil)
				}
			}
		}

		if m.InjectOrgHeader {
			r.Header.Set(scopeIDHeader, token.Organization)
		}
		return nil
	}
	if m.verifier != nil && idToken != "" { // OIDC flow
		_, err := m.verifier.Verify(r.Context(), idToken)
		if err != nil {
			// Fail closed: a token whose signature cannot be verified is
			// rejected regardless of the `verify` setting. Authorizing on
			// claims from an unverified JWT would let anyone forge group and
			// tenant membership. The `verify false` option only suppresses
			// issuer/expiry strictness within the verifier, not the signature
			// check itself.
			m.logger.Error("invalid token detected",
				zap.Error(err),
				zap.String("remoteAddr", r.RemoteAddr))
			return caddyhttp.Error(http.StatusUnauthorized, err)
		}
		type DexClaims struct {
			ObservabilityReadTenants  []string `json:"ort,omitempty"`
			ObservabilityWriteTenants []string `json:"owt,omitempty"`
			Groups                    []string `json:"groups,omitempty"`
			jwt.RegisteredClaims
		}
		token, err := jwt.ParseWithClaims(idToken, &DexClaims{}, func(token *jwt.Token) (any, error) {
			return []byte(""), jwt.ErrTokenUnverifiable // We already verified
		})
		if !errors.Is(err, jwt.ErrTokenUnverifiable) {
			return caddyhttp.Error(http.StatusUnauthorized, err)
		}
		// Verified
		claims, ok := token.Claims.(*DexClaims)
		if !ok {
			err := fmt.Errorf("invalid claims detected: %w", err)
			m.logger.Error("invalid claims detected", zap.Error(err))
			return caddyhttp.Error(http.StatusUnauthorized, err)
		}

		// Check group claims
		for _, group := range m.Groups {
			if !slices.Contains(claims.Groups, group) {
				m.logger.Error("missing group claim", zap.String("group", group))
				return caddyhttp.Error(http.StatusUnauthorized, nil)
			}
		}

		// Inject X-Scope-OrgID header
		if m.InjectOrgHeader {
			switch m.TenantOrgClaim {
			case "ort":
				if len(claims.ObservabilityReadTenants) > 0 {
					if m.Debug {
						m.logger.Info("ort X-Scope-OrgID", zap.String("value", strings.Join(claims.ObservabilityReadTenants, "|")))
					}
					r.Header.Set(scopeIDHeader, strings.Join(claims.ObservabilityReadTenants, "|"))
				}
			case "owt":
				if len(claims.ObservabilityWriteTenants) > 0 {
					if m.Debug {
						m.logger.Info("owt X-Scope-OrgID", zap.String("value", strings.Join(claims.ObservabilityWriteTenants, "|")))
					}
					r.Header.Set(scopeIDHeader, strings.Join(claims.ObservabilityWriteTenants, "|"))
				}
			default:
				if m.Debug {
					m.logger.Info("not injecting X-Scope-OrgID header")
				}
			}
		} else {
			if m.Debug {
				m.logger.Info("not injecting X-Scope-OrgID header")
			}
		}
		return nil
	}
	// No valid token found
	if m.Debug {
		m.logger.Error("no valid token found")
	}
	return caddyhttp.Error(http.StatusUnauthorized, nil)
}

// readTokenFile reads a static token file and returns a map of tokens
func (m *Middleware) readTokenFile(filename string) (map[string]keys.Key, error) {
	tokens := make(map[string]keys.Key)

	file, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("opening file %s: %w", filename, err)
	}
	defer func(file *os.File) {
		_ = file.Close()
	}(file)

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		var decoded keys.Key
		trimmedLine := strings.TrimSpace(scanner.Text())
		if len(trimmedLine) == 0 { // Skip empty lines
			continue
		}
		prefixRemoved := strings.TrimPrefix(trimmedLine, keys.Prefix)
		decodedString, err := base64.StdEncoding.DecodeString(prefixRemoved)
		if err != nil {
			return nil, fmt.Errorf("decode token: %w", err)
		}
		err = json.Unmarshal([]byte(decodedString), &decoded)
		if err != nil {
			return nil, fmt.Errorf("unmarshal token: %w '%s'", err, decodedString)
		}
		tokens[trimmedLine] = decoded
	}

	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("scanner: %w", err)
	}
	m.logger.Info("loaded tokens", zap.Int("apiKeyCount", len(tokens)))
	return tokens, nil
}

// LastNChars returns the last n characters of a string.
func LastNChars(n int, s string) string {
	if len(s) > n {
		return s[len(s)-n:]
	}
	return s
}

// Interface guards
var (
	_ caddy.Provisioner           = (*Middleware)(nil)
	_ caddy.Validator             = (*Middleware)(nil)
	_ caddyhttp.MiddlewareHandler = (*Middleware)(nil)
	_ caddyfile.Unmarshaler       = (*Middleware)(nil)
	_ caddy.Module                = (*Middleware)(nil)
)