From 5ff9cd1257463961b5b5d040712beef51d588dc1 Mon Sep 17 00:00:00 2001
From: Weidi Deng <weidi_deng@icloud.com>
Date: Thu, 26 Oct 2023 21:52:20 +0800
Subject: [PATCH] only set read limit for http redirect listener when request
 is http

---
 modules/caddyhttp/httpredirectlistener.go | 48 +++++++++++++----------
 1 file changed, 28 insertions(+), 20 deletions(-)

diff --git a/modules/caddyhttp/httpredirectlistener.go b/modules/caddyhttp/httpredirectlistener.go
index 082dc7ce8bb..03860a2d9fb 100644
--- a/modules/caddyhttp/httpredirectlistener.go
+++ b/modules/caddyhttp/httpredirectlistener.go
@@ -16,11 +16,11 @@ package caddyhttp
 
 import (
 	"bufio"
+	"bytes"
 	"fmt"
 	"io"
 	"net"
 	"net/http"
-	"sync"
 
 	"github.com/caddyserver/caddy/v2"
 	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
@@ -86,15 +86,17 @@ func (l *httpRedirectListener) Accept() (net.Conn, error) {
 	}
 
 	return &httpRedirectConn{
-		Conn: c,
-		r:    bufio.NewReader(io.LimitReader(c, maxHeaderBytes)),
+		Conn:  c,
+		limit: maxHeaderBytes,
+		r:     bufio.NewReader(c),
 	}, nil
 }
 
 type httpRedirectConn struct {
 	net.Conn
-	once sync.Once
-	r    *bufio.Reader
+	once  bool
+	limit int64
+	r     *bufio.Reader
 }
 
 // Read tries to peek at the first few bytes of the request, and if we get
@@ -102,23 +104,36 @@ type httpRedirectConn struct {
 // like an HTTP request, then we perform a HTTP->HTTPS redirect on the same
 // port as the original connection.
 func (c *httpRedirectConn) Read(p []byte) (int, error) {
-	var errReturn error
-	c.once.Do(func() {
+	// no need to use sync.Once - net.Conn is not read from concurrently.
+	if !c.once {
+		c.once = true
+
 		firstBytes, err := c.r.Peek(5)
 		if err != nil {
-			return
+			return 0, err
 		}
 
 		// If the request doesn't look like HTTP, then it's probably
-		// TLS bytes and we don't need to do anything.
+		// TLS bytes, and we don't need to do anything.
+		// 0, nil are valid return values - caller most likely will try again.
 		if !firstBytesLookLikeHTTP(firstBytes) {
-			return
+			return 0, nil
 		}
 
+		// From now on, we can be almost certain the request is HTTP.
+		// The returned error will be non nil and caller are expected to
+		// close the connection.
+
+		// Set the read limit, io.MultiReader is needed because
+		// when resetting, *bufio.Reader discards buffered data.
+		buffered, _ := c.r.Peek(c.r.Buffered())
+		mr := io.MultiReader(bytes.NewReader(buffered), c.Conn)
+		c.r.Reset(io.LimitReader(mr, c.limit))
+
 		// Parse the HTTP request, so we can get the Host and URL to redirect to.
 		req, err := http.ReadRequest(c.r)
 		if err != nil {
-			return
+			return 0, fmt.Errorf("couldn't read HTTP request")
 		}
 
 		// Build the redirect response, using the same Host and URL,
@@ -136,18 +151,11 @@ func (c *httpRedirectConn) Read(p []byte) (int, error) {
 
 		err = resp.Write(c.Conn)
 		if err != nil {
-			errReturn = fmt.Errorf("couldn't write HTTP->HTTPS redirect")
-			return
+			return 0, fmt.Errorf("couldn't write HTTP->HTTPS redirect")
 		}
 
-		errReturn = fmt.Errorf("redirected HTTP request on HTTPS port")
-		c.Conn.Close()
-	})
-
-	if errReturn != nil {
-		return 0, errReturn
+		return 0, fmt.Errorf("redirected HTTP request on HTTPS port")
 	}
-
 	return c.r.Read(p)
 }