fix: set exhausted cookie with env json (#5868)

* fix: set exhausted cookie with env json

* lint

internal/api/http/middleware/access_interceptor.go

@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"context"
 	"net"
 	"net/http"
 	"net/url"
@@ -32,15 +33,54 @@ type AccessConfig struct {
 // 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 {
+func NewAccessInterceptor(svc *logstore.Service, cookieHandler *http_utils.CookieHandler, cookieConfig *AccessConfig) *AccessInterceptor {
 	return &AccessInterceptor{
 		svc:           svc,
 		cookieHandler: cookieHandler,
 		limitConfig:   cookieConfig,
-		storeOnly:     storeOnly,
 	}
 }
 
+func (a *AccessInterceptor) WithoutLimiting() *AccessInterceptor {
+	return &AccessInterceptor{
+		svc:           a.svc,
+		cookieHandler: a.cookieHandler,
+		limitConfig:   a.limitConfig,
+		storeOnly:     true,
+	}
+}
+
+func (a *AccessInterceptor) AccessService() *logstore.Service {
+	return a.svc
+}
+
+func (a *AccessInterceptor) Limit(ctx context.Context) bool {
+	if !a.svc.Enabled() || a.storeOnly {
+		return false
+	}
+	instance := authz.GetInstance(ctx)
+	remaining := a.svc.Limit(ctx, instance.InstanceID())
+	return remaining != nil && *remaining <= 0
+}
+
+func (a *AccessInterceptor) SetExhaustedCookie(writer http.ResponseWriter, 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")
+		}
+	}
+	a.cookieHandler.SetCookie(writer, a.limitConfig.ExhaustedCookieKey, domain, cookieValue)
+}
+
+func (a *AccessInterceptor) DeleteExhaustedCookie(writer http.ResponseWriter, request *http.Request) {
+	a.cookieHandler.DeleteCookie(writer, request, a.limitConfig.ExhaustedCookieKey)
+}
+
 func (a *AccessInterceptor) Handle(next http.Handler) http.Handler {
 	if !a.svc.Enabled() {
 		return next
@@ -49,23 +89,16 @@ func (a *AccessInterceptor) Handle(next http.Handler) http.Handler {
 		ctx := request.Context()
 		tracingCtx, checkSpan := tracing.NewNamedSpan(ctx, "checkAccess")
 		wrappedWriter := &statusRecorder{ResponseWriter: writer, status: 0}
-		instance := authz.GetInstance(ctx)
-		limit := false
-		if !a.storeOnly {
-			remaining := a.svc.Limit(tracingCtx, instance.InstanceID())
-			limit = remaining != nil && *remaining == 0
-		}
+		limited := a.Limit(tracingCtx)
 		checkSpan.End()
-		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)
+		if limited {
+			a.SetExhaustedCookie(wrappedWriter, 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
+		}
+		if !limited && !a.storeOnly {
+			a.DeleteExhaustedCookie(wrappedWriter, request)
+		}
+		if !limited {
 			next.ServeHTTP(wrappedWriter, request)
 		}
 		tracingCtx, writeSpan := tracing.NewNamedSpan(tracingCtx, "writeAccess")
@@ -75,6 +108,7 @@ func (a *AccessInterceptor) Handle(next http.Handler) http.Handler {
 		if err != nil {
 			logging.WithError(err).WithField("url", requestURL).Warning("failed to unescape request url")
 		}
+		instance := authz.GetInstance(tracingCtx)
 		a.svc.Handle(tracingCtx, &access.Record{
 			LogDate:         time.Now(),
 			Protocol:        access.HTTP,
@@ -90,24 +124,6 @@ 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