fix: improve exhausted SetCookie header (#5789)

* fix: remove access interceptor for console

* feat: template quota cookie value

* fix: send exhausted cookie from grpc-gateway

* refactor: remove ineffectual err assignments

* Update internal/api/grpc/server/gateway.go

Co-authored-by: Livio Spring <livio.a@gmail.com>

* use dynamic host header to find instance

* add instance mgmt url to environment.json

* support hosts with default ports

* fix linting

* docs: update lb example

* print access logs to stdout

* fix grpc gateway exhausted cookies

* cleanup

---------

Co-authored-by: Livio Spring <livio.a@gmail.com>

internal/api/api.go

@@ -34,6 +34,9 @@ type API struct {
 	http1HostName string
 	grpcGateway   *server.Gateway
 	healthServer  *health.Server
+	cookieHandler *http_util.CookieHandler
+	cookieConfig  *http_mw.AccessConfig
+	queries       *query.Queries
 }
 
 type healthCheck interface {
@@ -49,6 +52,8 @@ func New(
 	authZ internal_authz.Config,
 	tlsConfig *tls.Config, http2HostName, http1HostName string,
 	accessSvc *logstore.Service,
+	cookieHandler *http_util.CookieHandler,
+	cookieConfig *http_mw.AccessConfig,
 ) (_ *API, err error) {
 	api := &API{
 		port:          port,
@@ -56,10 +61,13 @@ func New(
 		health:        queries,
 		router:        router,
 		http1HostName: http1HostName,
+		cookieConfig:  cookieConfig,
+		cookieHandler: cookieHandler,
+		queries:       queries,
 	}
 
 	api.grpcServer = server.CreateServer(api.verifier, authZ, queries, http2HostName, tlsConfig, accessSvc)
-	api.grpcGateway, err = server.CreateGateway(ctx, port, http1HostName)
+	api.grpcGateway, err = server.CreateGateway(ctx, port, http1HostName, cookieHandler, cookieConfig)
 	if err != nil {
 		return nil, err
 	}
@@ -77,7 +85,15 @@ func New(
 // used for v1 api (system, admin, mgmt, auth)
 func (a *API) RegisterServer(ctx context.Context, grpcServer server.WithGatewayPrefix) error {
 	grpcServer.RegisterServer(a.grpcServer)
-	handler, prefix, err := server.CreateGatewayWithPrefix(ctx, grpcServer, a.port, a.http1HostName)
+	handler, prefix, err := server.CreateGatewayWithPrefix(
+		ctx,
+		grpcServer,
+		a.port,
+		a.http1HostName,
+		a.cookieHandler,
+		a.cookieConfig,
+		a.queries,
+	)
 	if err != nil {
 		return err
 	}

internal/api/grpc/server/gateway.go

@@ -16,7 +16,9 @@ import (
 
 	client_middleware "github.com/zitadel/zitadel/internal/api/grpc/client/middleware"
 	"github.com/zitadel/zitadel/internal/api/grpc/server/middleware"
+	http_utils "github.com/zitadel/zitadel/internal/api/http"
 	http_mw "github.com/zitadel/zitadel/internal/api/http/middleware"
+	"github.com/zitadel/zitadel/internal/query"
 )
 
 const (
@@ -67,10 +69,13 @@ type Gateway struct {
 	mux           *runtime.ServeMux
 	http1HostName string
 	connection    *grpc.ClientConn
+	cookieHandler *http_utils.CookieHandler
+	cookieConfig  *http_mw.AccessConfig
+	queries       *query.Queries
 }
 
 func (g *Gateway) Handler() http.Handler {
-	return addInterceptors(g.mux, g.http1HostName)
+	return addInterceptors(g.mux, g.http1HostName, g.cookieHandler, g.cookieConfig, g.queries)
 }
 
 type CustomHTTPResponse interface {
@@ -79,7 +84,15 @@ type CustomHTTPResponse interface {
 
 type RegisterGatewayFunc func(ctx context.Context, mux *runtime.ServeMux, conn *grpc.ClientConn) error
 
-func CreateGatewayWithPrefix(ctx context.Context, g WithGatewayPrefix, port uint16, http1HostName string) (http.Handler, string, error) {
+func CreateGatewayWithPrefix(
+	ctx context.Context,
+	g WithGatewayPrefix,
+	port uint16,
+	http1HostName string,
+	cookieHandler *http_utils.CookieHandler,
+	cookieConfig *http_mw.AccessConfig,
+	queries *query.Queries,
+) (http.Handler, string, error) {
 	runtimeMux := runtime.NewServeMux(serveMuxOptions...)
 	opts := []grpc.DialOption{
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
@@ -93,10 +106,10 @@ func CreateGatewayWithPrefix(ctx context.Context, g WithGatewayPrefix, port uint
 	if err != nil {
 		return nil, "", fmt.Errorf("failed to register grpc gateway: %w", err)
 	}
-	return addInterceptors(runtimeMux, http1HostName), g.GatewayPathPrefix(), nil
+	return addInterceptors(runtimeMux, http1HostName, cookieHandler, cookieConfig, queries), g.GatewayPathPrefix(), nil
 }
 
-func CreateGateway(ctx context.Context, port uint16, http1HostName string) (*Gateway, error) {
+func CreateGateway(ctx context.Context, port uint16, http1HostName string, cookieHandler *http_utils.CookieHandler, cookieConfig *http_mw.AccessConfig) (*Gateway, error) {
 	connection, err := dial(ctx,
 		port,
 		[]grpc.DialOption{
@@ -111,6 +124,8 @@ func CreateGateway(ctx context.Context, port uint16, http1HostName string) (*Gat
 		mux:           runtimeMux,
 		http1HostName: http1HostName,
 		connection:    connection,
+		cookieHandler: cookieHandler,
+		cookieConfig:  cookieConfig,
 	}, nil
 }
 
@@ -145,13 +160,23 @@ func dial(ctx context.Context, port uint16, opts []grpc.DialOption) (*grpc.Clien
 	return conn, nil
 }
 
-func addInterceptors(handler http.Handler, http1HostName string) http.Handler {
+func addInterceptors(
+	handler http.Handler,
+	http1HostName string,
+	cookieHandler *http_utils.CookieHandler,
+	cookieConfig *http_mw.AccessConfig,
+	queries *query.Queries,
+) http.Handler {
 	handler = http_mw.CallDurationHandler(handler)
 	handler = http1Host(handler, http1HostName)
 	handler = http_mw.CORSInterceptor(handler)
 	handler = http_mw.RobotsTagHandler(handler)
 	handler = http_mw.DefaultTelemetryHandler(handler)
-	return http_mw.DefaultMetricsHandler(handler)
+	// For some non-obvious reason, the exhaustedCookieInterceptor sends the SetCookie header
+	// only if it follows the http_mw.DefaultTelemetryHandler
+	handler = exhaustedCookieInterceptor(handler, cookieHandler, cookieConfig, queries)
+	handler = http_mw.DefaultMetricsHandler(handler)
+	return handler
 }
 
 func http1Host(next http.Handler, http1HostName string) http.Handler {
@@ -165,3 +190,38 @@ func http1Host(next http.Handler, http1HostName string) http.Handler {
 		next.ServeHTTP(w, r)
 	})
 }
+
+func exhaustedCookieInterceptor(
+	next http.Handler,
+	cookieHandler *http_utils.CookieHandler,
+	cookieConfig *http_mw.AccessConfig,
+	queries *query.Queries,
+) http.Handler {
+	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		next.ServeHTTP(&cookieResponseWriter{
+			ResponseWriter: writer,
+			cookieHandler:  cookieHandler,
+			cookieConfig:   cookieConfig,
+			request:        request,
+			queries:        queries,
+		}, request)
+	})
+}
+
+type cookieResponseWriter struct {
+	http.ResponseWriter
+	cookieHandler *http_utils.CookieHandler
+	cookieConfig  *http_mw.AccessConfig
+	request       *http.Request
+	queries       *query.Queries
+}
+
+func (r *cookieResponseWriter) WriteHeader(status int) {
+	if status >= 200 && status < 300 {
+		http_mw.DeleteExhaustedCookie(r.cookieHandler, r.ResponseWriter, r.request, r.cookieConfig)
+	}
+	if status == http.StatusTooManyRequests {
+		http_mw.SetExhaustedCookie(r.cookieHandler, r.ResponseWriter, r.cookieConfig, r.request)
+	}
+	r.ResponseWriter.WriteHeader(status)
+}

internal/api/http/middleware/access_interceptor.go

@@ -1,15 +1,16 @@
 package middleware
 
 import (
-	"math"
+	"net"
 	"net/http"
 	"net/url"
-	"strconv"
+	"strings"
 	"time"
 
 	"github.com/zitadel/logging"
 
 	"github.com/zitadel/zitadel/internal/api/authz"
+	"github.com/zitadel/zitadel/internal/api/grpc/server/middleware"
 	http_utils "github.com/zitadel/zitadel/internal/api/http"
 	"github.com/zitadel/zitadel/internal/logstore"
 	"github.com/zitadel/zitadel/internal/logstore/emitters/access"
@@ -20,6 +21,7 @@ type AccessInterceptor struct {
 	svc           *logstore.Service
 	cookieHandler *http_utils.CookieHandler
 	limitConfig   *AccessConfig
+	storeOnly     bool
 }
 
 type AccessConfig struct {
@@ -27,14 +29,15 @@ type AccessConfig struct {
 	ExhaustedCookieMaxAge time.Duration
 }
 
-func NewAccessInterceptor(svc *logstore.Service, cookieConfig *AccessConfig) *AccessInterceptor {
+// NewAccessInterceptor intercepts all requests and stores them to the logstore.
+// If storeOnly is false, it also checks if requests are exhausted.
+// If requests are exhausted, it also returns http.StatusTooManyRequests and sets a cookie
+func NewAccessInterceptor(svc *logstore.Service, cookieHandler *http_utils.CookieHandler, cookieConfig *AccessConfig, storeOnly bool) *AccessInterceptor {
 	return &AccessInterceptor{
-		svc: svc,
-		cookieHandler: http_utils.NewCookieHandler(
-			http_utils.WithUnsecure(),
-			http_utils.WithMaxAge(int(math.Floor(cookieConfig.ExhaustedCookieMaxAge.Seconds()))),
-		),
-		limitConfig: cookieConfig,
+		svc:           svc,
+		cookieHandler: cookieHandler,
+		limitConfig:   cookieConfig,
+		storeOnly:     storeOnly,
 	}
 }
 
@@ -44,36 +47,33 @@ func (a *AccessInterceptor) Handle(next http.Handler) http.Handler {
 	}
 	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
 		ctx := request.Context()
-		var err error
-
 		tracingCtx, checkSpan := tracing.NewNamedSpan(ctx, "checkAccess")
-
 		wrappedWriter := &statusRecorder{ResponseWriter: writer, status: 0}
-
 		instance := authz.GetInstance(ctx)
-		remaining := a.svc.Limit(tracingCtx, instance.InstanceID())
-		limit := remaining != nil && *remaining == 0
-
-		a.cookieHandler.SetCookie(wrappedWriter, a.limitConfig.ExhaustedCookieKey, request.Host, strconv.FormatBool(limit))
-
-		if limit {
-			wrappedWriter.WriteHeader(http.StatusTooManyRequests)
-			wrappedWriter.ignoreWrites = true
+		limit := false
+		if !a.storeOnly {
+			remaining := a.svc.Limit(tracingCtx, instance.InstanceID())
+			limit = remaining != nil && *remaining == 0
 		}
-
 		checkSpan.End()
-
-		next.ServeHTTP(wrappedWriter, request)
-
+		if limit {
+			// Limit can only be true when storeOnly is false, so set the cookie and the response code
+			SetExhaustedCookie(a.cookieHandler, wrappedWriter, a.limitConfig, request)
+			http.Error(wrappedWriter, "quota for authenticated requests is exhausted", http.StatusTooManyRequests)
+		} else {
+			if !a.storeOnly {
+				// If not limited and not storeOnly, ensure the cookie is deleted
+				DeleteExhaustedCookie(a.cookieHandler, wrappedWriter, request, a.limitConfig)
+			}
+			// Always serve if not limited
+			next.ServeHTTP(wrappedWriter, request)
+		}
 		tracingCtx, writeSpan := tracing.NewNamedSpan(tracingCtx, "writeAccess")
 		defer writeSpan.End()
-
 		requestURL := request.RequestURI
 		unescapedURL, err := url.QueryUnescape(requestURL)
 		if err != nil {
 			logging.WithError(err).WithField("url", requestURL).Warning("failed to unescape request url")
-			// err = nil is effective because of deferred tracing span end
-			err = nil
 		}
 		a.svc.Handle(tracingCtx, &access.Record{
 			LogDate:         time.Now(),
@@ -90,6 +90,24 @@ func (a *AccessInterceptor) Handle(next http.Handler) http.Handler {
 	})
 }
 
+func SetExhaustedCookie(cookieHandler *http_utils.CookieHandler, writer http.ResponseWriter, cookieConfig *AccessConfig, request *http.Request) {
+	cookieValue := "true"
+	host := request.Header.Get(middleware.HTTP1Host)
+	domain := host
+	if strings.ContainsAny(host, ":") {
+		var err error
+		domain, _, err = net.SplitHostPort(host)
+		if err != nil {
+			logging.WithError(err).WithField("host", host).Warning("failed to extract cookie domain from request host")
+		}
+	}
+	cookieHandler.SetCookie(writer, cookieConfig.ExhaustedCookieKey, domain, cookieValue)
+}
+
+func DeleteExhaustedCookie(cookieHandler *http_utils.CookieHandler, writer http.ResponseWriter, request *http.Request, cookieConfig *AccessConfig) {
+	cookieHandler.DeleteCookie(writer, request, cookieConfig.ExhaustedCookieKey)
+}
+
 type statusRecorder struct {
 	http.ResponseWriter
 	status       int

internal/api/ui/console/console.go

@@ -1,9 +1,11 @@
 package console
 
 import (
+	"bytes"
 	"embed"
 	"encoding/json"
 	"fmt"
+	"html/template"
 	"io/fs"
 	"net/http"
 	"os"
@@ -22,8 +24,9 @@ import (
 )
 
 type Config struct {
-	ShortCache middleware.CacheConfig
-	LongCache  middleware.CacheConfig
+	ShortCache            middleware.CacheConfig
+	LongCache             middleware.CacheConfig
+	InstanceManagementURL string
 }
 
 type spaHandler struct {
@@ -106,7 +109,13 @@ func Start(config Config, externalSecure bool, issuer op.IssuerFromRequest, call
 	handler.Use(callDurationInterceptor, instanceHandler, security, accessInterceptor)
 	handler.Handle(envRequestPath, middleware.TelemetryHandler()(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		url := http_util.BuildOrigin(r.Host, externalSecure)
-		environmentJSON, err := createEnvironmentJSON(url, issuer(r), authz.GetInstance(r.Context()).ConsoleClientID(), customerPortal)
+		instance := authz.GetInstance(r.Context())
+		instanceMgmtURL, err := templateInstanceManagementURL(config.InstanceManagementURL, instance)
+		if err != nil {
+			http.Error(w, fmt.Sprintf("unable to template instance management url for console: %v", err), http.StatusInternalServerError)
+			return
+		}
+		environmentJSON, err := createEnvironmentJSON(url, issuer(r), instance.ConsoleClientID(), customerPortal, instanceMgmtURL)
 		if err != nil {
 			http.Error(w, fmt.Sprintf("unable to marshal env for console: %v", err), http.StatusInternalServerError)
 			return
@@ -118,6 +127,18 @@ func Start(config Config, externalSecure bool, issuer op.IssuerFromRequest, call
 	return handler, nil
 }
 
+func templateInstanceManagementURL(templateableCookieValue string, instance authz.Instance) (string, error) {
+	cookieValueTemplate, err := template.New("cookievalue").Parse(templateableCookieValue)
+	if err != nil {
+		return templateableCookieValue, err
+	}
+	cookieValue := new(bytes.Buffer)
+	if err = cookieValueTemplate.Execute(cookieValue, instance); err != nil {
+		return templateableCookieValue, err
+	}
+	return cookieValue.String(), nil
+}
+
 func csp() *middleware.CSP {
 	csp := middleware.DefaultSCP
 	csp.StyleSrc = csp.StyleSrc.AddInline()
@@ -127,17 +148,19 @@ func csp() *middleware.CSP {
 	return &csp
 }
 
-func createEnvironmentJSON(api, issuer, clientID, customerPortal string) ([]byte, error) {
+func createEnvironmentJSON(api, issuer, clientID, customerPortal, instanceMgmtUrl string) ([]byte, error) {
 	environment := struct {
-		API            string `json:"api,omitempty"`
-		Issuer         string `json:"issuer,omitempty"`
-		ClientID       string `json:"clientid,omitempty"`
-		CustomerPortal string `json:"customer_portal,omitempty"`
+		API                   string `json:"api,omitempty"`
+		Issuer                string `json:"issuer,omitempty"`
+		ClientID              string `json:"clientid,omitempty"`
+		CustomerPortal        string `json:"customer_portal,omitempty"`
+		InstanceManagementURL string `json:"instance_management_url,omitempty"`
 	}{
-		API:            api,
-		Issuer:         issuer,
-		ClientID:       clientID,
-		CustomerPortal: customerPortal,
+		API:                   api,
+		Issuer:                issuer,
+		ClientID:              clientID,
+		CustomerPortal:        customerPortal,
+		InstanceManagementURL: instanceMgmtUrl,
 	}
 	return json.Marshal(environment)
 }