httpcaddyfile: Fix memory leak in HTTP/2 push handler

modules/caddyhttp/push/handler.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"net/http"
 	"strings"
+	"sync"
 
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
@@ -84,57 +85,75 @@ func (h Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhtt
 		return next.ServeHTTP(w, r)
 	}
 
-	repl := r.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)
-	server := r.Context().Value(caddyhttp.ServerCtxKey).(*caddyhttp.Server)
-	shouldLogCredentials := server.Logs != nil && server.Logs.ShouldLogCredentials
-
-	// create header for push requests
-	hdr := h.initializePushHeaders(r, repl)
-
-	// push first!
-	for _, resource := range h.Resources {
-		if c := h.logger.Check(zapcore.DebugLevel, "pushing resource"); c != nil {
-			c.Write(
-				zap.String("uri", r.RequestURI),
-				zap.String("push_method", resource.Method),
-				zap.String("push_target", resource.Target),
-				zap.Object("push_headers", caddyhttp.LoggableHTTPHeader{
-					Header:               hdr,
-					ShouldLogCredentials: shouldLogCredentials,
-				}),
-			)
-		}
-		err := pusher.Push(repl.ReplaceAll(resource.Target, "."), &http.PushOptions{
-			Method: resource.Method,
-			Header: hdr,
-		})
-		if err != nil {
-			// usually this means either that push is not
-			// supported or concurrent streams are full
-			break
+	var repl *caddy.Replacer
+	var hdr http.Header
+
+	if len(h.Resources) > 0 {
+		repl = r.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)
+		server := r.Context().Value(caddyhttp.ServerCtxKey).(*caddyhttp.Server)
+		shouldLogCredentials := server.Logs != nil && server.Logs.ShouldLogCredentials
+
+		for _, resource := range h.Resources {
+			if hdr == nil {
+				hdr = h.initializePushHeaders(r, repl)
+			}
+
+			if c := h.logger.Check(zapcore.DebugLevel, "pushing resource"); c != nil {
+				c.Write(
+					zap.String("uri", r.RequestURI),
+					zap.String("push_method", resource.Method),
+					zap.String("push_target", resource.Target),
+					zap.Object("push_headers", caddyhttp.LoggableHTTPHeader{
+						Header:               hdr,
+						ShouldLogCredentials: shouldLogCredentials,
+					}),
+				)
+			}
+
+			target := resource.Target
+			if strings.Contains(target, "{") {
+				target = repl.ReplaceAll(target, ".")
+			}
+
+			err := pusher.Push(target, &http.PushOptions{
+				Method: resource.Method,
+				Header: hdr,
+			})
+			if err != nil {
+				break
+			}
 		}
 	}
 
 	// wrap the response writer so that we can initiate push of any resources
 	// described in Link header fields before the response is written
-	lp := linkPusher{
-		ResponseWriterWrapper: &caddyhttp.ResponseWriterWrapper{ResponseWriter: w},
-		handler:               h,
-		pusher:                pusher,
-		header:                hdr,
-		request:               r,
-	}
+	lp := linkPusherPool.Get().(*linkPusher)
+	lp.ResponseWriterWrapper.ResponseWriter = w
+	lp.handler = &h
+	lp.pusher = pusher
+	lp.header = hdr // reuse header if already initialized
+	lp.request = r
+	lp.repl = repl
+
+	// clear references and return to pool after serving
+	defer func() {
+		caddyhttp.SetVar(r.Context(), pushedLink, nil)
+		lp.ResponseWriterWrapper.ResponseWriter = nil
+		lp.handler = nil
+		lp.pusher = nil
+		lp.header = nil
+		lp.request = nil
+		lp.repl = nil
+		linkPusherPool.Put(lp)
+	}()
 
 	// serve only after pushing!
-	if err := next.ServeHTTP(lp, r); err != nil {
-		return err
-	}
-
-	return nil
+	return next.ServeHTTP(lp, r)
 }
 
 func (h Handler) initializePushHeaders(r *http.Request, repl *caddy.Replacer) http.Header {
-	hdr := make(http.Header)
+	// Pre-allocate capacity for safeHeaders + pushHeader
+	hdr := make(http.Header, len(safeHeaders)+1)
 
 	// prevent recursive pushes
 	hdr.Set(pushHeader, "1")
@@ -153,6 +172,9 @@ func (h Handler) initializePushHeaders(r *http.Request, repl *caddy.Replacer) ht
 
 	// user can customize the push request headers
 	if h.Headers != nil {
+		if repl == nil {
+			repl = r.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)
+		}
 		h.Headers.ApplyTo(hdr, repl)
 	}
 
@@ -163,10 +185,10 @@ func (h Handler) initializePushHeaders(r *http.Request, repl *caddy.Replacer) ht
 // resources described by them. If a resource has the "nopush"
 // attribute or describes an external entity (meaning, the resource
 // URI includes a scheme), it will not be pushed.
-func (h Handler) servePreloadLinks(pusher http.Pusher, hdr http.Header, resources []string) {
-	for _, resource := range resources {
-		for _, resource := range parseLinkHeader(resource) {
-			if _, ok := resource.params["nopush"]; ok {
+func (h Handler) servePreloadLinks(pusher http.Pusher, hdr http.Header, links []string) {
+	for _, linkHdr := range links {
+		for _, resource := range parseLinkHeader(linkHdr) {
+			if resource.noPush {
 				continue
 			}
 			if isRemoteResource(resource.uri) {
@@ -202,21 +224,39 @@ type HeaderConfig struct {
 // described by Link response headers get pushed before
 // the response is allowed to be written.
 type linkPusher struct {
-	*caddyhttp.ResponseWriterWrapper
-	handler Handler
+	caddyhttp.ResponseWriterWrapper
+	handler *Handler
 	pusher  http.Pusher
 	header  http.Header
 	request *http.Request
+	repl    *caddy.Replacer
 }
 
-func (lp linkPusher) WriteHeader(statusCode int) {
+// Push implements http.Pusher.
+func (lp *linkPusher) Push(target string, opts *http.PushOptions) error {
+	return lp.pusher.Push(target, opts)
+}
+
+// linkPusherPool minimizes allocations for the middleware wrapper.
+var linkPusherPool = sync.Pool{
+	New: func() any {
+		return new(linkPusher)
+	},
+}
+
+func (lp *linkPusher) WriteHeader(statusCode int) {
 	if links, ok := lp.ResponseWriter.Header()["Link"]; ok {
 		// only initiate these pushes if it hasn't been done yet
 		if val := caddyhttp.GetVar(lp.request.Context(), pushedLink); val == nil {
 			if c := lp.handler.logger.Check(zapcore.DebugLevel, "pushing Link resources"); c != nil {
 				c.Write(zap.Strings("linked", links))
 			}
 			caddyhttp.SetVar(lp.request.Context(), pushedLink, true)
+
+			if lp.header == nil {
+				lp.header = lp.handler.initializePushHeaders(lp.request, lp.repl)
+			}
+
 			lp.handler.servePreloadLinks(lp.pusher, lp.header, links)
 		}
 	}

modules/caddyhttp/push/handler_test.go

@@ -0,0 +1,56 @@
+package push
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/caddyserver/caddy/v2"
+	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
+	"go.uber.org/zap"
+)
+
+type mockPusher struct {
+	http.ResponseWriter
+	pushed bool
+}
+
+func (m *mockPusher) Push(target string, opts *http.PushOptions) error {
+	m.pushed = true
+	return nil
+}
+
+func BenchmarkServeHTTP(b *testing.B) {
+	logger := zap.NewNop()
+	h := Handler{
+		Resources: []Resource{
+			{Target: "/style.css", Method: "GET"},
+		},
+		logger: logger,
+	}
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		req, _ := http.NewRequest("GET", "/", nil)
+
+		// Setup Replacer and Server in context as required by ServeHTTP
+		repl := caddy.NewReplacer()
+		ctx := context.WithValue(req.Context(), caddy.ReplacerCtxKey, repl)
+		ctx = context.WithValue(ctx, caddyhttp.ServerCtxKey, &caddyhttp.Server{})
+		req = req.WithContext(ctx)
+
+		rr := httptest.NewRecorder()
+
+		// Wrap with a pusher
+		pusher := &mockPusher{ResponseWriter: rr}
+
+		_ = h.ServeHTTP(pusher, req, caddyhttp.HandlerFunc(func(w http.ResponseWriter, r *http.Request) error {
+			w.Header().Set("Link", "</script.js>; rel=preload")
+			w.WriteHeader(http.StatusOK)
+			return nil
+		}))
+	}
+}

modules/caddyhttp/push/link.go

@@ -21,46 +21,48 @@ import (
 // linkResource contains the results of a parsed Link header.
 type linkResource struct {
 	uri    string
-	params map[string]string
+	noPush bool
 }
 
-// parseLinkHeader is responsible for parsing Link header
-// and returning list of found resources.
-//
-// Accepted formats are:
-//
-//	Link: <resource>; as=script
-//	Link: <resource>; as=script,<resource>; as=style
-//	Link: <resource>;<resource2>
-//
-// where <resource> begins with a forward slash (/).
 func parseLinkHeader(header string) []linkResource {
 	resources := []linkResource{}
 
-	if header == "" {
-		return resources
-	}
-
-	for link := range strings.SplitSeq(header, comma) {
-		l := linkResource{params: make(map[string]string)}
+	for len(header) > 0 {
+		var link string
+		idx := strings.IndexByte(header, ',')
+		if idx >= 0 {
+			link = header[:idx]
+			header = header[idx+1:]
+		} else {
+			link = header
+			header = ""
+		}
 
-		li, ri := strings.Index(link, "<"), strings.Index(link, ">")
+		li, ri := strings.IndexByte(link, '<'), strings.IndexByte(link, '>')
 		if li == -1 || ri == -1 {
 			continue
 		}
 
-		l.uri = strings.TrimSpace(link[li+1 : ri])
+		l := linkResource{
+			uri: strings.TrimSpace(link[li+1 : ri]),
+		}
 
-		for param := range strings.SplitSeq(strings.TrimSpace(link[ri+1:]), semicolon) {
-			before, after, isCut := strings.Cut(strings.TrimSpace(param), equal)
-			key := strings.TrimSpace(before)
-			if key == "" {
-				continue
-			}
-			if isCut {
-				l.params[key] = strings.TrimSpace(after)
+		paramsPart := strings.TrimSpace(link[ri+1:])
+		for len(paramsPart) > 0 {
+			var param string
+			pidx := strings.IndexByte(paramsPart, ';')
+			if pidx >= 0 {
+				param = paramsPart[:pidx]
+				paramsPart = paramsPart[pidx+1:]
 			} else {
-				l.params[key] = key
+				param = paramsPart
+				paramsPart = ""
+			}
+
+			before, _, _ := strings.Cut(strings.TrimSpace(param), "=")
+			if strings.TrimSpace(before) == "nopush" {
+				l.noPush = true
+				break
 			}
 		}
 
@@ -69,9 +71,3 @@ func parseLinkHeader(header string) []linkResource {
 
 	return resources
 }
-
-const (
-	comma     = ","
-	semicolon = ";"
-	equal     = "="
-)