[proxy, management] Add header auth, access restrictions, and session idle timeout (#5587)

proxy/internal/auth/header.go

@@ -0,0 +1,69 @@
+package auth
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+
+	"github.com/netbirdio/netbird/proxy/auth"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// ErrHeaderAuthFailed indicates that the header was present but the
+// credential did not validate. Callers should return 401 instead of
+// falling through to other auth schemes.
+var ErrHeaderAuthFailed = errors.New("header authentication failed")
+
+// Header implements header-based authentication. The proxy checks for the
+// configured header in each request and validates its value via gRPC.
+type Header struct {
+	id         types.ServiceID
+	accountId  types.AccountID
+	headerName string
+	client     authenticator
+}
+
+// NewHeader creates a Header authentication scheme for the given header name.
+func NewHeader(client authenticator, id types.ServiceID, accountId types.AccountID, headerName string) Header {
+	return Header{
+		id:         id,
+		accountId:  accountId,
+		headerName: headerName,
+		client:     client,
+	}
+}
+
+// Type returns auth.MethodHeader.
+func (Header) Type() auth.Method {
+	return auth.MethodHeader
+}
+
+// Authenticate checks for the configured header in the request. If absent,
+// returns empty (unauthenticated). If present, validates via gRPC.
+func (h Header) Authenticate(r *http.Request) (string, string, error) {
+	value := r.Header.Get(h.headerName)
+	if value == "" {
+		return "", "", nil
+	}
+
+	res, err := h.client.Authenticate(r.Context(), &proto.AuthenticateRequest{
+		Id:        string(h.id),
+		AccountId: string(h.accountId),
+		Request: &proto.AuthenticateRequest_HeaderAuth{
+			HeaderAuth: &proto.HeaderAuthRequest{
+				HeaderValue: value,
+				HeaderName:  h.headerName,
+			},
+		},
+	})
+	if err != nil {
+		return "", "", fmt.Errorf("authenticate header: %w", err)
+	}
+
+	if res.GetSuccess() {
+		return res.GetSessionToken(), "", nil
+	}
+
+	return "", "", ErrHeaderAuthFailed
+}

proxy/internal/auth/middleware.go

@@ -4,9 +4,12 @@ import (
 	"context"
 	"crypto/ed25519"
 	"encoding/base64"
+	"errors"
 	"fmt"
+	"html"
 	"net"
 	"net/http"
+	"net/netip"
 	"net/url"
 	"sync"
 	"time"
@@ -16,11 +19,16 @@ import (
 
 	"github.com/netbirdio/netbird/proxy/auth"
 	"github.com/netbirdio/netbird/proxy/internal/proxy"
+	"github.com/netbirdio/netbird/proxy/internal/restrict"
 	"github.com/netbirdio/netbird/proxy/internal/types"
 	"github.com/netbirdio/netbird/proxy/web"
 	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
+// errValidationUnavailable indicates that session validation failed due to
+// an infrastructure error (e.g. gRPC unavailable), not an invalid token.
+var errValidationUnavailable = errors.New("session validation unavailable")
+
 type authenticator interface {
 	Authenticate(ctx context.Context, in *proto.AuthenticateRequest, opts ...grpc.CallOption) (*proto.AuthenticateResponse, error)
 }
@@ -40,12 +48,14 @@ type Scheme interface {
 	Authenticate(*http.Request) (token string, promptData string, err error)
 }
 
+// DomainConfig holds the authentication and restriction settings for a protected domain.
 type DomainConfig struct {
 	Schemes           []Scheme
 	SessionPublicKey  ed25519.PublicKey
 	SessionExpiration time.Duration
 	AccountID         types.AccountID
 	ServiceID         types.ServiceID
+	IPRestrictions    *restrict.Filter
 }
 
 type validationResult struct {
@@ -54,17 +64,18 @@ type validationResult struct {
 	DeniedReason string
 }
 
+// Middleware applies per-domain authentication and IP restriction checks.
 type Middleware struct {
 	domainsMux       sync.RWMutex
 	domains          map[string]DomainConfig
 	logger           *log.Logger
 	sessionValidator SessionValidator
+	geo              restrict.GeoResolver
 }
 
-// NewMiddleware creates a new authentication middleware.
-// The sessionValidator is optional; if nil, OIDC session tokens will be validated
-// locally without group access checks.
-func NewMiddleware(logger *log.Logger, sessionValidator SessionValidator) *Middleware {
+// NewMiddleware creates a new authentication middleware. The sessionValidator is
+// optional; if nil, OIDC session tokens are validated locally without group access checks.
+func NewMiddleware(logger *log.Logger, sessionValidator SessionValidator, geo restrict.GeoResolver) *Middleware {
 	if logger == nil {
 		logger = log.StandardLogger()
 	}
@@ -72,18 +83,12 @@ func NewMiddleware(logger *log.Logger, sessionValidator SessionValidator) *Middl
 		domains:          make(map[string]DomainConfig),
 		logger:           logger,
 		sessionValidator: sessionValidator,
+		geo:              geo,
 	}
 }
 
-// Protect applies authentication middleware to the passed handler.
-// For each incoming request it will be checked against the middleware's
-// internal list of protected domains.
-// If the Host domain in the inbound request is not present, then it will
-// simply be passed through.
-// However, if the Host domain is present, then the specified authentication
-// schemes for that domain will be applied to the request.
-// In the event that no authentication schemes are defined for the domain,
-// then the request will also be simply passed through.
+// Protect wraps next with per-domain authentication and IP restriction checks.
+// Requests whose Host is not registered pass through unchanged.
 func (mw *Middleware) Protect(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		host, _, err := net.SplitHostPort(r.Host)
@@ -94,8 +99,7 @@ func (mw *Middleware) Protect(next http.Handler) http.Handler {
 		config, exists := mw.getDomainConfig(host)
 		mw.logger.Debugf("checking authentication for host: %s, exists: %t", host, exists)
 
-		// Domains that are not configured here or have no authentication schemes applied should simply pass through.
-		if !exists || len(config.Schemes) == 0 {
+		if !exists {
 			next.ServeHTTP(w, r)
 			return
 		}
@@ -103,6 +107,16 @@ func (mw *Middleware) Protect(next http.Handler) http.Handler {
 		// Set account and service IDs in captured data for access logging.
 		setCapturedIDs(r, config)
 
+		if !mw.checkIPRestrictions(w, r, config) {
+			return
+		}
+
+		// Domains with no authentication schemes pass through after IP checks.
+		if len(config.Schemes) == 0 {
+			next.ServeHTTP(w, r)
+			return
+		}
+
 		if mw.handleOAuthCallbackError(w, r) {
 			return
 		}
@@ -111,6 +125,10 @@ func (mw *Middleware) Protect(next http.Handler) http.Handler {
 			return
 		}
 
+		if mw.forwardWithHeaderAuth(w, r, host, config, next) {
+			return
+		}
+
 		mw.authenticateWithSchemes(w, r, host, config)
 	})
 }
@@ -124,11 +142,65 @@ func (mw *Middleware) getDomainConfig(host string) (DomainConfig, bool) {
 
 func setCapturedIDs(r *http.Request, config DomainConfig) {
 	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
-		cd.SetAccountId(config.AccountID)
-		cd.SetServiceId(config.ServiceID)
+		cd.SetAccountID(config.AccountID)
+		cd.SetServiceID(config.ServiceID)
 	}
 }
 
+// checkIPRestrictions validates the client IP against the domain's IP restrictions.
+// Uses the resolved client IP from CapturedData (which accounts for trusted proxies)
+// rather than r.RemoteAddr directly.
+func (mw *Middleware) checkIPRestrictions(w http.ResponseWriter, r *http.Request, config DomainConfig) bool {
+	if config.IPRestrictions == nil {
+		return true
+	}
+
+	clientIP := mw.resolveClientIP(r)
+	if !clientIP.IsValid() {
+		mw.logger.Debugf("IP restriction: cannot resolve client address for %q, denying", r.RemoteAddr)
+		http.Error(w, "Forbidden", http.StatusForbidden)
+		return false
+	}
+
+	verdict := config.IPRestrictions.Check(clientIP, mw.geo)
+	if verdict == restrict.Allow {
+		return true
+	}
+
+	reason := verdict.String()
+	mw.blockIPRestriction(r, reason)
+	http.Error(w, "Forbidden", http.StatusForbidden)
+	return false
+}
+
+// resolveClientIP extracts the real client IP from CapturedData, falling back to r.RemoteAddr.
+func (mw *Middleware) resolveClientIP(r *http.Request) netip.Addr {
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		if ip := cd.GetClientIP(); ip.IsValid() {
+			return ip
+		}
+	}
+
+	clientIPStr, _, _ := net.SplitHostPort(r.RemoteAddr)
+	if clientIPStr == "" {
+		clientIPStr = r.RemoteAddr
+	}
+	addr, err := netip.ParseAddr(clientIPStr)
+	if err != nil {
+		return netip.Addr{}
+	}
+	return addr.Unmap()
+}
+
+// blockIPRestriction sets captured data fields for an IP-restriction block event.
+func (mw *Middleware) blockIPRestriction(r *http.Request, reason string) {
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetOrigin(proxy.OriginAuth)
+		cd.SetAuthMethod(reason)
+	}
+	mw.logger.Debugf("IP restriction: %s for %s", reason, r.RemoteAddr)
+}
+
 // handleOAuthCallbackError checks for error query parameters from an OAuth
 // callback and renders the access denied page if present.
 func (mw *Middleware) handleOAuthCallbackError(w http.ResponseWriter, r *http.Request) bool {
@@ -146,6 +218,8 @@ func (mw *Middleware) handleOAuthCallbackError(w http.ResponseWriter, r *http.Re
 	errDesc := r.URL.Query().Get("error_description")
 	if errDesc == "" {
 		errDesc = "An error occurred during authentication"
+	} else {
+		errDesc = html.EscapeString(errDesc)
 	}
 	web.ServeAccessDeniedPage(w, r, http.StatusForbidden, "Access Denied", errDesc, requestID)
 	return true
@@ -170,6 +244,85 @@ func (mw *Middleware) forwardWithSessionCookie(w http.ResponseWriter, r *http.Re
 	return true
 }
 
+// forwardWithHeaderAuth checks for a Header auth scheme. If the header validates,
+// the request is forwarded directly (no redirect), which is important for API clients.
+func (mw *Middleware) forwardWithHeaderAuth(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, next http.Handler) bool {
+	for _, scheme := range config.Schemes {
+		hdr, ok := scheme.(Header)
+		if !ok {
+			continue
+		}
+
+		handled := mw.tryHeaderScheme(w, r, host, config, hdr, next)
+		if handled {
+			return true
+		}
+	}
+	return false
+}
+
+func (mw *Middleware) tryHeaderScheme(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, hdr Header, next http.Handler) bool {
+	token, _, err := hdr.Authenticate(r)
+	if err != nil {
+		return mw.handleHeaderAuthError(w, r, err)
+	}
+	if token == "" {
+		return false
+	}
+
+	result, err := mw.validateSessionToken(r.Context(), host, token, config.SessionPublicKey, auth.MethodHeader)
+	if err != nil {
+		setHeaderCapturedData(r.Context(), "")
+		status := http.StatusBadRequest
+		msg := "invalid session token"
+		if errors.Is(err, errValidationUnavailable) {
+			status = http.StatusBadGateway
+			msg = "authentication service unavailable"
+		}
+		http.Error(w, msg, status)
+		return true
+	}
+
+	if !result.Valid {
+		setHeaderCapturedData(r.Context(), result.UserID)
+		http.Error(w, "Unauthorized", http.StatusUnauthorized)
+		return true
+	}
+
+	setSessionCookie(w, token, config.SessionExpiration)
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetUserID(result.UserID)
+		cd.SetAuthMethod(auth.MethodHeader.String())
+	}
+
+	next.ServeHTTP(w, r)
+	return true
+}
+
+func (mw *Middleware) handleHeaderAuthError(w http.ResponseWriter, r *http.Request, err error) bool {
+	if errors.Is(err, ErrHeaderAuthFailed) {
+		setHeaderCapturedData(r.Context(), "")
+		http.Error(w, "Unauthorized", http.StatusUnauthorized)
+		return true
+	}
+	mw.logger.WithField("scheme", "header").Warnf("header auth infrastructure error: %v", err)
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetOrigin(proxy.OriginAuth)
+	}
+	http.Error(w, "authentication service unavailable", http.StatusBadGateway)
+	return true
+}
+
+func setHeaderCapturedData(ctx context.Context, userID string) {
+	cd := proxy.CapturedDataFromContext(ctx)
+	if cd == nil {
+		return
+	}
+	cd.SetOrigin(proxy.OriginAuth)
+	cd.SetAuthMethod(auth.MethodHeader.String())
+	cd.SetUserID(userID)
+}
+
 // authenticateWithSchemes tries each configured auth scheme in order.
 // On success it sets a session cookie and redirects; on failure it renders the login page.
 func (mw *Middleware) authenticateWithSchemes(w http.ResponseWriter, r *http.Request, host string, config DomainConfig) {
@@ -217,7 +370,13 @@ func (mw *Middleware) handleAuthenticatedToken(w http.ResponseWriter, r *http.Re
 			cd.SetOrigin(proxy.OriginAuth)
 			cd.SetAuthMethod(scheme.Type().String())
 		}
-		http.Error(w, err.Error(), http.StatusBadRequest)
+		status := http.StatusBadRequest
+		msg := "invalid session token"
+		if errors.Is(err, errValidationUnavailable) {
+			status = http.StatusBadGateway
+			msg = "authentication service unavailable"
+		}
+		http.Error(w, msg, status)
 		return
 	}
 
@@ -233,7 +392,21 @@ func (mw *Middleware) handleAuthenticatedToken(w http.ResponseWriter, r *http.Re
 		return
 	}
 
-	expiration := config.SessionExpiration
+	setSessionCookie(w, token, config.SessionExpiration)
+
+	// Redirect instead of forwarding the auth POST to the backend.
+	// The browser will follow with a GET carrying the new session cookie.
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetOrigin(proxy.OriginAuth)
+		cd.SetUserID(result.UserID)
+		cd.SetAuthMethod(scheme.Type().String())
+	}
+	redirectURL := stripSessionTokenParam(r.URL)
+	http.Redirect(w, r, redirectURL, http.StatusSeeOther)
+}
+
+// setSessionCookie writes a session cookie with secure defaults.
+func setSessionCookie(w http.ResponseWriter, token string, expiration time.Duration) {
 	if expiration == 0 {
 		expiration = auth.DefaultSessionExpiry
 	}
@@ -245,16 +418,6 @@ func (mw *Middleware) handleAuthenticatedToken(w http.ResponseWriter, r *http.Re
 		SameSite: http.SameSiteLaxMode,
 		MaxAge:   int(expiration.Seconds()),
 	})
-
-	// Redirect instead of forwarding the auth POST to the backend.
-	// The browser will follow with a GET carrying the new session cookie.
-	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
-		cd.SetOrigin(proxy.OriginAuth)
-		cd.SetUserID(result.UserID)
-		cd.SetAuthMethod(scheme.Type().String())
-	}
-	redirectURL := stripSessionTokenParam(r.URL)
-	http.Redirect(w, r, redirectURL, http.StatusSeeOther)
 }
 
 // wasCredentialSubmitted checks if credentials were submitted for the given auth method.
@@ -275,13 +438,14 @@ func wasCredentialSubmitted(r *http.Request, method auth.Method) bool {
 // session JWTs. Returns an error if the key is missing or invalid.
 // Callers must not serve the domain if this returns an error, to avoid
 // exposing an unauthenticated service.
-func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 string, expiration time.Duration, accountID types.AccountID, serviceID types.ServiceID) error {
+func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 string, expiration time.Duration, accountID types.AccountID, serviceID types.ServiceID, ipRestrictions *restrict.Filter) error {
 	if len(schemes) == 0 {
 		mw.domainsMux.Lock()
 		defer mw.domainsMux.Unlock()
 		mw.domains[domain] = DomainConfig{
-			AccountID: accountID,
-			ServiceID: serviceID,
+			AccountID:      accountID,
+			ServiceID:      serviceID,
+			IPRestrictions: ipRestrictions,
 		}
 		return nil
 	}
@@ -302,30 +466,28 @@ func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 st
 		SessionExpiration: expiration,
 		AccountID:         accountID,
 		ServiceID:         serviceID,
+		IPRestrictions:    ipRestrictions,
 	}
 	return nil
 }
 
+// RemoveDomain unregisters authentication for the given domain.
 func (mw *Middleware) RemoveDomain(domain string) {
 	mw.domainsMux.Lock()
 	defer mw.domainsMux.Unlock()
 	delete(mw.domains, domain)
 }
 
-// validateSessionToken validates a session token, optionally checking group access via gRPC.
-// For OIDC tokens with a configured validator, it calls ValidateSession to check group access.
-// For other auth methods (PIN, password), it validates the JWT locally.
-// Returns a validationResult with user ID and validity status, or error for invalid tokens.
+// validateSessionToken validates a session token. OIDC tokens with a configured
+// validator go through gRPC for group access checks; other methods validate locally.
 func (mw *Middleware) validateSessionToken(ctx context.Context, host, token string, publicKey ed25519.PublicKey, method auth.Method) (*validationResult, error) {
-	// For OIDC with a session validator, call the gRPC service to check group access
 	if method == auth.MethodOIDC && mw.sessionValidator != nil {
 		resp, err := mw.sessionValidator.ValidateSession(ctx, &proto.ValidateSessionRequest{
 			Domain:       host,
 			SessionToken: token,
 		})
 		if err != nil {
-			mw.logger.WithError(err).Error("ValidateSession gRPC call failed")
-			return nil, fmt.Errorf("session validation failed")
+			return nil, fmt.Errorf("%w: %w", errValidationUnavailable, err)
 		}
 		if !resp.Valid {
 			mw.logger.WithFields(log.Fields{
@@ -342,7 +504,6 @@ func (mw *Middleware) validateSessionToken(ctx context.Context, host, token stri
 		return &validationResult{UserID: resp.UserId, Valid: true}, nil
 	}
 
-	// For non-OIDC methods or when no validator is configured, validate JWT locally
 	userID, _, err := auth.ValidateSessionJWT(token, host, publicKey)
 	if err != nil {
 		return nil, err

proxy/internal/auth/middleware_test.go

@@ -1,11 +1,14 @@
 package auth
 
 import (
+	"context"
 	"crypto/ed25519"
 	"crypto/rand"
 	"encoding/base64"
+	"errors"
 	"net/http"
 	"net/http/httptest"
+	"net/netip"
 	"net/url"
 	"strings"
 	"testing"
@@ -14,10 +17,13 @@ import (
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
 
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
 	"github.com/netbirdio/netbird/proxy/auth"
 	"github.com/netbirdio/netbird/proxy/internal/proxy"
+	"github.com/netbirdio/netbird/proxy/internal/restrict"
+	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
 func generateTestKeyPair(t *testing.T) *sessionkey.KeyPair {
@@ -52,11 +58,11 @@ func newPassthroughHandler() http.Handler {
 }
 
 func TestAddDomain_ValidKey(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	err := mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "")
+	err := mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil)
 	require.NoError(t, err)
 
 	mw.domainsMux.RLock()
@@ -70,10 +76,10 @@ func TestAddDomain_ValidKey(t *testing.T) {
 }
 
 func TestAddDomain_EmptyKey(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	err := mw.AddDomain("example.com", []Scheme{scheme}, "", time.Hour, "", "")
+	err := mw.AddDomain("example.com", []Scheme{scheme}, "", time.Hour, "", "", nil)
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "invalid session public key size")
 
@@ -84,10 +90,10 @@ func TestAddDomain_EmptyKey(t *testing.T) {
 }
 
 func TestAddDomain_InvalidBase64(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	err := mw.AddDomain("example.com", []Scheme{scheme}, "not-valid-base64!!!", time.Hour, "", "")
+	err := mw.AddDomain("example.com", []Scheme{scheme}, "not-valid-base64!!!", time.Hour, "", "", nil)
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "decode session public key")
 
@@ -98,11 +104,11 @@ func TestAddDomain_InvalidBase64(t *testing.T) {
 }
 
 func TestAddDomain_WrongKeySize(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 
 	shortKey := base64.StdEncoding.EncodeToString([]byte("tooshort"))
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	err := mw.AddDomain("example.com", []Scheme{scheme}, shortKey, time.Hour, "", "")
+	err := mw.AddDomain("example.com", []Scheme{scheme}, shortKey, time.Hour, "", "", nil)
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "invalid session public key size")
 
@@ -113,9 +119,9 @@ func TestAddDomain_WrongKeySize(t *testing.T) {
 }
 
 func TestAddDomain_NoSchemes_NoKeyRequired(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 
-	err := mw.AddDomain("example.com", nil, "", time.Hour, "", "")
+	err := mw.AddDomain("example.com", nil, "", time.Hour, "", "", nil)
 	require.NoError(t, err, "domains with no auth schemes should not require a key")
 
 	mw.domainsMux.RLock()
@@ -125,14 +131,14 @@ func TestAddDomain_NoSchemes_NoKeyRequired(t *testing.T) {
 }
 
 func TestAddDomain_OverwritesPreviousConfig(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp1 := generateTestKeyPair(t)
 	kp2 := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
 
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp1.PublicKey, time.Hour, "", ""))
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp2.PublicKey, 2*time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp1.PublicKey, time.Hour, "", "", nil))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp2.PublicKey, 2*time.Hour, "", "", nil))
 
 	mw.domainsMux.RLock()
 	config := mw.domains["example.com"]
@@ -144,11 +150,11 @@ func TestAddDomain_OverwritesPreviousConfig(t *testing.T) {
 }
 
 func TestRemoveDomain(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	mw.RemoveDomain("example.com")
 
@@ -159,7 +165,7 @@ func TestRemoveDomain(t *testing.T) {
 }
 
 func TestProtect_UnknownDomainPassesThrough(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	handler := mw.Protect(newPassthroughHandler())
 
 	req := httptest.NewRequest(http.MethodGet, "http://unknown.com/", nil)
@@ -171,8 +177,8 @@ func TestProtect_UnknownDomainPassesThrough(t *testing.T) {
 }
 
 func TestProtect_DomainWithNoSchemesPassesThrough(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
-	require.NoError(t, mw.AddDomain("example.com", nil, "", time.Hour, "", ""))
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	require.NoError(t, mw.AddDomain("example.com", nil, "", time.Hour, "", "", nil))
 
 	handler := mw.Protect(newPassthroughHandler())
 
@@ -185,11 +191,11 @@ func TestProtect_DomainWithNoSchemesPassesThrough(t *testing.T) {
 }
 
 func TestProtect_UnauthenticatedRequestIsBlocked(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	var backendCalled bool
 	backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -206,11 +212,11 @@ func TestProtect_UnauthenticatedRequestIsBlocked(t *testing.T) {
 }
 
 func TestProtect_HostWithPortIsMatched(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	var backendCalled bool
 	backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -227,16 +233,16 @@ func TestProtect_HostWithPortIsMatched(t *testing.T) {
 }
 
 func TestProtect_ValidSessionCookiePassesThrough(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	token, err := sessionkey.SignToken(kp.PrivateKey, "test-user", "example.com", auth.MethodPIN, time.Hour)
 	require.NoError(t, err)
 
-	capturedData := &proxy.CapturedData{}
+	capturedData := proxy.NewCapturedData("")
 	handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		cd := proxy.CapturedDataFromContext(r.Context())
 		require.NotNil(t, cd)
@@ -257,11 +263,11 @@ func TestProtect_ValidSessionCookiePassesThrough(t *testing.T) {
 }
 
 func TestProtect_ExpiredSessionCookieIsRejected(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	// Sign a token that expired 1 second ago.
 	token, err := sessionkey.SignToken(kp.PrivateKey, "test-user", "example.com", auth.MethodPIN, -time.Second)
@@ -283,11 +289,11 @@ func TestProtect_ExpiredSessionCookieIsRejected(t *testing.T) {
 }
 
 func TestProtect_WrongDomainCookieIsRejected(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	// Token signed for a different domain audience.
 	token, err := sessionkey.SignToken(kp.PrivateKey, "test-user", "other.com", auth.MethodPIN, time.Hour)
@@ -309,12 +315,12 @@ func TestProtect_WrongDomainCookieIsRejected(t *testing.T) {
 }
 
 func TestProtect_WrongKeyCookieIsRejected(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp1 := generateTestKeyPair(t)
 	kp2 := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp1.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp1.PublicKey, time.Hour, "", "", nil))
 
 	// Token signed with a different private key.
 	token, err := sessionkey.SignToken(kp2.PrivateKey, "test-user", "example.com", auth.MethodPIN, time.Hour)
@@ -336,7 +342,7 @@ func TestProtect_WrongKeyCookieIsRejected(t *testing.T) {
 }
 
 func TestProtect_SchemeAuthRedirectsWithCookie(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	token, err := sessionkey.SignToken(kp.PrivateKey, "pin-user", "example.com", auth.MethodPIN, time.Hour)
@@ -351,7 +357,7 @@ func TestProtect_SchemeAuthRedirectsWithCookie(t *testing.T) {
 			return "", "pin", nil
 		},
 	}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	var backendCalled bool
 	backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -386,7 +392,7 @@ func TestProtect_SchemeAuthRedirectsWithCookie(t *testing.T) {
 }
 
 func TestProtect_FailedAuthDoesNotSetCookie(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{
@@ -395,7 +401,7 @@ func TestProtect_FailedAuthDoesNotSetCookie(t *testing.T) {
 			return "", "pin", nil
 		},
 	}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	handler := mw.Protect(newPassthroughHandler())
 
@@ -409,7 +415,7 @@ func TestProtect_FailedAuthDoesNotSetCookie(t *testing.T) {
 }
 
 func TestProtect_MultipleSchemes(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	token, err := sessionkey.SignToken(kp.PrivateKey, "password-user", "example.com", auth.MethodPassword, time.Hour)
@@ -431,7 +437,7 @@ func TestProtect_MultipleSchemes(t *testing.T) {
 			return "", "password", nil
 		},
 	}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{pinScheme, passwordScheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{pinScheme, passwordScheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	var backendCalled bool
 	backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -451,7 +457,7 @@ func TestProtect_MultipleSchemes(t *testing.T) {
 }
 
 func TestProtect_InvalidTokenFromSchemeReturns400(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	// Return a garbage token that won't validate.
@@ -461,7 +467,7 @@ func TestProtect_InvalidTokenFromSchemeReturns400(t *testing.T) {
 			return "invalid-jwt-token", "", nil
 		},
 	}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	handler := mw.Protect(newPassthroughHandler())
 
@@ -473,7 +479,7 @@ func TestProtect_InvalidTokenFromSchemeReturns400(t *testing.T) {
 }
 
 func TestAddDomain_RandomBytes32NotEd25519(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 
 	// 32 random bytes that happen to be valid base64 and correct size
 	// but are actually a valid ed25519 public key length-wise.
@@ -485,19 +491,19 @@ func TestAddDomain_RandomBytes32NotEd25519(t *testing.T) {
 	key := base64.StdEncoding.EncodeToString(randomBytes)
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
 
-	err = mw.AddDomain("example.com", []Scheme{scheme}, key, time.Hour, "", "")
+	err = mw.AddDomain("example.com", []Scheme{scheme}, key, time.Hour, "", "", nil)
 	require.NoError(t, err, "any 32-byte key should be accepted at registration time")
 }
 
 func TestAddDomain_InvalidKeyDoesNotCorruptExistingConfig(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
 	// Attempt to overwrite with an invalid key.
-	err := mw.AddDomain("example.com", []Scheme{scheme}, "bad", time.Hour, "", "")
+	err := mw.AddDomain("example.com", []Scheme{scheme}, "bad", time.Hour, "", "", nil)
 	require.Error(t, err)
 
 	// The original valid config should still be intact.
@@ -511,7 +517,7 @@ func TestAddDomain_InvalidKeyDoesNotCorruptExistingConfig(t *testing.T) {
 }
 
 func TestProtect_FailedPinAuthCapturesAuthMethod(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	// Scheme that always fails authentication (returns empty token)
@@ -521,9 +527,9 @@ func TestProtect_FailedPinAuthCapturesAuthMethod(t *testing.T) {
 			return "", "pin", nil
 		},
 	}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
-	capturedData := &proxy.CapturedData{}
+	capturedData := proxy.NewCapturedData("")
 	handler := mw.Protect(newPassthroughHandler())
 
 	// Submit wrong PIN - should capture auth method
@@ -539,7 +545,7 @@ func TestProtect_FailedPinAuthCapturesAuthMethod(t *testing.T) {
 }
 
 func TestProtect_FailedPasswordAuthCapturesAuthMethod(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{
@@ -548,9 +554,9 @@ func TestProtect_FailedPasswordAuthCapturesAuthMethod(t *testing.T) {
 			return "", "password", nil
 		},
 	}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
-	capturedData := &proxy.CapturedData{}
+	capturedData := proxy.NewCapturedData("")
 	handler := mw.Protect(newPassthroughHandler())
 
 	// Submit wrong password - should capture auth method
@@ -566,7 +572,7 @@ func TestProtect_FailedPasswordAuthCapturesAuthMethod(t *testing.T) {
 }
 
 func TestProtect_NoCredentialsDoesNotCaptureAuthMethod(t *testing.T) {
-	mw := NewMiddleware(log.StandardLogger(), nil)
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
 	kp := generateTestKeyPair(t)
 
 	scheme := &stubScheme{
@@ -575,9 +581,9 @@ func TestProtect_NoCredentialsDoesNotCaptureAuthMethod(t *testing.T) {
 			return "", "pin", nil
 		},
 	}
-	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", ""))
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil))
 
-	capturedData := &proxy.CapturedData{}
+	capturedData := proxy.NewCapturedData("")
 	handler := mw.Protect(newPassthroughHandler())
 
 	// No credentials submitted - should not capture auth method
@@ -658,3 +664,271 @@ func TestWasCredentialSubmitted(t *testing.T) {
 		})
 	}
 }
+
+func TestCheckIPRestrictions_UnparseableAddress(t *testing.T) {
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+
+	err := mw.AddDomain("example.com", nil, "", 0, "acc1", "svc1",
+		restrict.ParseFilter([]string{"10.0.0.0/8"}, nil, nil, nil))
+	require.NoError(t, err)
+
+	handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	tests := []struct {
+		name       string
+		remoteAddr string
+		wantCode   int
+	}{
+		{"unparsable address denies", "not-an-ip:1234", http.StatusForbidden},
+		{"empty address denies", "", http.StatusForbidden},
+		{"allowed address passes", "10.1.2.3:5678", http.StatusOK},
+		{"denied address blocked", "192.168.1.1:5678", http.StatusForbidden},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+			req.RemoteAddr = tt.remoteAddr
+			req.Host = "example.com"
+			rr := httptest.NewRecorder()
+			handler.ServeHTTP(rr, req)
+			assert.Equal(t, tt.wantCode, rr.Code)
+		})
+	}
+}
+
+func TestCheckIPRestrictions_UsesCapturedDataClientIP(t *testing.T) {
+	// When CapturedData is set (by the access log middleware, which resolves
+	// trusted proxies), checkIPRestrictions should use that IP, not RemoteAddr.
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+
+	err := mw.AddDomain("example.com", nil, "", 0, "acc1", "svc1",
+		restrict.ParseFilter([]string{"203.0.113.0/24"}, nil, nil, nil))
+	require.NoError(t, err)
+
+	handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	// RemoteAddr is a trusted proxy, but CapturedData has the real client IP.
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	req.RemoteAddr = "10.0.0.1:5000"
+	req.Host = "example.com"
+
+	cd := proxy.NewCapturedData("")
+	cd.SetClientIP(netip.MustParseAddr("203.0.113.50"))
+	ctx := proxy.WithCapturedData(req.Context(), cd)
+	req = req.WithContext(ctx)
+
+	rr := httptest.NewRecorder()
+	handler.ServeHTTP(rr, req)
+	assert.Equal(t, http.StatusOK, rr.Code, "should use CapturedData IP (203.0.113.50), not RemoteAddr (10.0.0.1)")
+
+	// Same request but CapturedData has a blocked IP.
+	req2 := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	req2.RemoteAddr = "203.0.113.50:5000"
+	req2.Host = "example.com"
+
+	cd2 := proxy.NewCapturedData("")
+	cd2.SetClientIP(netip.MustParseAddr("10.0.0.1"))
+	ctx2 := proxy.WithCapturedData(req2.Context(), cd2)
+	req2 = req2.WithContext(ctx2)
+
+	rr2 := httptest.NewRecorder()
+	handler.ServeHTTP(rr2, req2)
+	assert.Equal(t, http.StatusForbidden, rr2.Code, "should use CapturedData IP (10.0.0.1), not RemoteAddr (203.0.113.50)")
+}
+
+func TestCheckIPRestrictions_NilGeoWithCountryRules(t *testing.T) {
+	// Geo is nil, country restrictions are configured: must deny (fail-close).
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+
+	err := mw.AddDomain("example.com", nil, "", 0, "acc1", "svc1",
+		restrict.ParseFilter(nil, nil, []string{"US"}, nil))
+	require.NoError(t, err)
+
+	handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	req.RemoteAddr = "1.2.3.4:5678"
+	req.Host = "example.com"
+	rr := httptest.NewRecorder()
+	handler.ServeHTTP(rr, req)
+	assert.Equal(t, http.StatusForbidden, rr.Code, "country restrictions with nil geo must deny")
+}
+
+// mockAuthenticator is a minimal mock for the authenticator gRPC interface
+// used by the Header scheme.
+type mockAuthenticator struct {
+	fn func(ctx context.Context, req *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error)
+}
+
+func (m *mockAuthenticator) Authenticate(ctx context.Context, in *proto.AuthenticateRequest, _ ...grpc.CallOption) (*proto.AuthenticateResponse, error) {
+	return m.fn(ctx, in)
+}
+
+// newHeaderSchemeWithToken creates a Header scheme backed by a mock that
+// returns a signed session token when the expected header value is provided.
+func newHeaderSchemeWithToken(t *testing.T, kp *sessionkey.KeyPair, headerName, expectedValue string) Header {
+	t.Helper()
+	token, err := sessionkey.SignToken(kp.PrivateKey, "header-user", "example.com", auth.MethodHeader, time.Hour)
+	require.NoError(t, err)
+
+	mock := &mockAuthenticator{fn: func(_ context.Context, req *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
+		ha := req.GetHeaderAuth()
+		if ha != nil && ha.GetHeaderValue() == expectedValue {
+			return &proto.AuthenticateResponse{Success: true, SessionToken: token}, nil
+		}
+		return &proto.AuthenticateResponse{Success: false}, nil
+	}}
+	return NewHeader(mock, "svc1", "acc1", headerName)
+}
+
+func TestProtect_HeaderAuth_ForwardsOnSuccess(t *testing.T) {
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+
+	hdr := newHeaderSchemeWithToken(t, kp, "X-API-Key", "secret-key")
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil))
+
+	var backendCalled bool
+	capturedData := proxy.NewCapturedData("")
+	handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		backendCalled = true
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("ok"))
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/path", nil)
+	req.Header.Set("X-API-Key", "secret-key")
+	req = req.WithContext(proxy.WithCapturedData(req.Context(), capturedData))
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.True(t, backendCalled, "backend should be called directly for header auth (no redirect)")
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "ok", rec.Body.String())
+
+	// Session cookie should be set.
+	var sessionCookie *http.Cookie
+	for _, c := range rec.Result().Cookies() {
+		if c.Name == auth.SessionCookieName {
+			sessionCookie = c
+			break
+		}
+	}
+	require.NotNil(t, sessionCookie, "session cookie should be set after successful header auth")
+	assert.True(t, sessionCookie.HttpOnly)
+	assert.True(t, sessionCookie.Secure)
+
+	assert.Equal(t, "header-user", capturedData.GetUserID())
+	assert.Equal(t, "header", capturedData.GetAuthMethod())
+}
+
+func TestProtect_HeaderAuth_MissingHeaderFallsThrough(t *testing.T) {
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+
+	hdr := newHeaderSchemeWithToken(t, kp, "X-API-Key", "secret-key")
+	// Also add a PIN scheme so we can verify fallthrough behavior.
+	pinScheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr, pinScheme}, kp.PublicKey, time.Hour, "acc1", "svc1", nil))
+
+	handler := mw.Protect(newPassthroughHandler())
+
+	// No X-API-Key header: should fall through to PIN login page (401).
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusUnauthorized, rec.Code, "missing header should fall through to login page")
+}
+
+func TestProtect_HeaderAuth_WrongValueReturns401(t *testing.T) {
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+
+	mock := &mockAuthenticator{fn: func(_ context.Context, _ *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
+		return &proto.AuthenticateResponse{Success: false}, nil
+	}}
+	hdr := NewHeader(mock, "svc1", "acc1", "X-API-Key")
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil))
+
+	capturedData := proxy.NewCapturedData("")
+	handler := mw.Protect(newPassthroughHandler())
+
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	req.Header.Set("X-API-Key", "wrong-key")
+	req = req.WithContext(proxy.WithCapturedData(req.Context(), capturedData))
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusUnauthorized, rec.Code)
+	assert.Equal(t, "header", capturedData.GetAuthMethod())
+}
+
+func TestProtect_HeaderAuth_InfraErrorReturns502(t *testing.T) {
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+
+	mock := &mockAuthenticator{fn: func(_ context.Context, _ *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
+		return nil, errors.New("gRPC unavailable")
+	}}
+	hdr := NewHeader(mock, "svc1", "acc1", "X-API-Key")
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil))
+
+	handler := mw.Protect(newPassthroughHandler())
+
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	req.Header.Set("X-API-Key", "some-key")
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusBadGateway, rec.Code)
+}
+
+func TestProtect_HeaderAuth_SubsequentRequestUsesSessionCookie(t *testing.T) {
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+
+	hdr := newHeaderSchemeWithToken(t, kp, "X-API-Key", "secret-key")
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil))
+
+	handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	// First request with header auth.
+	req1 := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	req1.Header.Set("X-API-Key", "secret-key")
+	req1 = req1.WithContext(proxy.WithCapturedData(req1.Context(), proxy.NewCapturedData("")))
+	rec1 := httptest.NewRecorder()
+	handler.ServeHTTP(rec1, req1)
+	require.Equal(t, http.StatusOK, rec1.Code)
+
+	// Extract session cookie.
+	var sessionCookie *http.Cookie
+	for _, c := range rec1.Result().Cookies() {
+		if c.Name == auth.SessionCookieName {
+			sessionCookie = c
+			break
+		}
+	}
+	require.NotNil(t, sessionCookie)
+
+	// Second request with only the session cookie (no header).
+	capturedData2 := proxy.NewCapturedData("")
+	req2 := httptest.NewRequest(http.MethodGet, "http://example.com/other", nil)
+	req2.AddCookie(sessionCookie)
+	req2 = req2.WithContext(proxy.WithCapturedData(req2.Context(), capturedData2))
+	rec2 := httptest.NewRecorder()
+	handler.ServeHTTP(rec2, req2)
+
+	assert.Equal(t, http.StatusOK, rec2.Code)
+	assert.Equal(t, "header-user", capturedData2.GetUserID())
+	assert.Equal(t, "header", capturedData2.GetAuthMethod())
+}