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/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)
+}