Repository: caddyserver/caddy
File: modules/caddyhttp/reverseproxy/reverseproxy.go (Go)

Code before:
	if (r.ProtoMajor == 2 && r.Method == http.MethodConnect && r.Header.Get(":protocol") == "websocket") ||
		(r.ProtoMajor == 3 && r.Method == http.MethodConnect && r.Proto == "websocket") {
		clonedReq.Header.Del(":protocol")
		// keep the body for later use. http1.1 upgrade uses http.NoBody
		caddyhttp.SetVar(clonedReq.Context(), "extended_connect_websocket_body", clonedReq.Body)
		clonedReq.Body = http.NoBody
		clonedReq.Method = http.MethodGet
		clonedReq.Header.Set("Upgrade", "websocket")
		clonedReq.Header.Set("Connection", "Upgrade")
		key := make([]byte, 16)
		_, randErr := rand.Read(key)
		if randErr != nil {
			return randErr
		}
		clonedReq.Header["Sec-WebSocket-Key"] = []string{base64.StdEncoding.EncodeToString(key)}
	}

	// we will need the original headers and Host value if
	// header operations are configured; this is so that each
	// retry can apply the modifications, because placeholders
	// may be used which depend on the selected upstream for
	// their values
	reqHost := clonedReq.Host
	reqHeader := clonedReq.Header

	var bufferedReqBody []byte
	if reqBodyBuf, ok := clonedReq.Body.(bodyReadCloser); ok && reqBodyBuf.body == nil {
		bufferedReqBody = append([]byte(nil), reqBodyBuf.buf.Bytes()...)
	}

	start := time.Now()
	defer func() {
		// total proxying duration, including time spent on LB and retries
		repl.Set("http.reverse_proxy.duration", time.Since(start))
		repl.Set("http.reverse_proxy.duration_ms", time.Since(start).Seconds()*1e3) // multiply seconds to preserve decimal (see #4666)
	}()

	// in the proxy loop, each iteration is an attempt to proxy the request,
	// and because we may retry some number of times, carry over the error
	// from previous tries because of the nuances of load balancing & retries
	var proxyErr error
	var retries int
	for {
		// if the request body was buffered (and only the entire body, hence no body
		// set to read from after the buffer), make reading from the body idempotent
		// and reusable, so if a backend partially or fully reads the body but then
		// produces an error, the request can be repeated to the next backend with
		// the full body (retries should only happen for idempotent requests) (see #6259)
		if bufferedReqBody != nil {
			clonedReq.Body = io.NopCloser(bytes.NewReader(bufferedReqBody))
		}

Review comment (suggestion):
We can modify the `Close` behavior of `bodyReadCloser`. Instead of putting the buffer to the pool, just close the body if not nil. Using `defer` below to put the buffer to the pool, so there is no need to copy the bytes.

Code after:
	if (r.ProtoMajor == 2 && r.Method == http.MethodConnect && r.Header.Get(":protocol") == "websocket") ||
		(r.ProtoMajor == 3 && r.Method == http.MethodConnect && r.Proto == "websocket") {
		clonedReq.Header.Del(":protocol")
		// keep the body for later use. http1.1 upgrade uses http.NoBody
		caddyhttp.SetVar(clonedReq.Context(), "extended_connect_websocket_body", clonedReq.Body)
		clonedReq.Body = http.NoBody
		clonedReq.Method = http.MethodGet
		clonedReq.Header.Set("Upgrade", "websocket")
		clonedReq.Header.Set("Connection", "Upgrade")
		key := make([]byte, 16)
		_, randErr := rand.Read(key)
		if randErr != nil {
			return randErr
		}
		clonedReq.Header["Sec-WebSocket-Key"] = []string{base64.StdEncoding.EncodeToString(key)}
	}

	// we will need the original headers and Host value if
	// header operations are configured; this is so that each
	// retry can apply the modifications, because placeholders
	// may be used which depend on the selected upstream for
	// their values
	reqHost := clonedReq.Host
	reqHeader := clonedReq.Header

	// If the cloned request body was fully buffered, keep a reference to its
	// buffer so we can reuse it across retries and return it to the pool
	// once we’re done.
	var bufferedReqBody *bytes.Buffer
	if reqBodyBuf, ok := clonedReq.Body.(bodyReadCloser); ok && reqBodyBuf.body == nil && reqBodyBuf.buf != nil {
		bufferedReqBody = reqBodyBuf.buf
		reqBodyBuf.buf = nil

		defer func() {
			bufferedReqBody.Reset()
			bufPool.Put(bufferedReqBody)
		}()
	}

	start := time.Now()
	defer func() {
		// total proxying duration, including time spent on LB and retries
		repl.Set("http.reverse_proxy.duration", time.Since(start))
		repl.Set("http.reverse_proxy.duration_ms", time.Since(start).Seconds()*1e3) // multiply seconds to preserve decimal (see #4666)
	}()

	// in the proxy loop, each iteration is an attempt to proxy the request,
	// and because we may retry some number of times, carry over the error
	// from previous tries because of the nuances of load balancing & retries
	var proxyErr error
	var retries int