deep review: fixes

Author: xcolwell

http.go

@@ -169,12 +169,22 @@ func (self *HttpProxy) handleHttp(w http.ResponseWriter, r *http.Request) {
 	// r.Header.Del("Proxy-Authorization")
 
 	b := bytes.NewBuffer(nil)
-	_, err := copyBufferWithTimeout(b, io.LimitReader(r.Body, self.MaxHttpBodyBytes), nil, self.ProxyReadTimeout, self.ProxyWriteTimeout)
-	r.Body.Close()
+	bodyReader := io.Reader(r.Body)
+	if 0 < self.MaxHttpBodyBytes {
+		bodyReader = io.LimitReader(r.Body, self.MaxHttpBodyBytes+1)
+	}
+	_, err := copyBufferWithTimeout(b, bodyReader, nil, self.ProxyReadTimeout, self.ProxyWriteTimeout)
+	if closeErr := r.Body.Close(); err == nil {
+		err = closeErr
+	}
 	if err != nil {
 		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
 		return
 	}
+	if 0 < self.MaxHttpBodyBytes && self.MaxHttpBodyBytes < int64(b.Len()) {
+		http.Error(w, http.StatusText(http.StatusRequestEntityTooLarge), http.StatusRequestEntityTooLarge)
+		return
+	}
 	bodyBytes := b.Bytes()
 
 	handleCtx, handleCancel := context.WithCancel(r.Context())
@@ -195,16 +205,7 @@ func (self *HttpProxy) handleHttp(w http.ResponseWriter, r *http.Request) {
 
 	var response *http.Response
 	for {
-		r2, err := http.NewRequestWithContext(
-			r.Context(),
-			r.Method,
-			r.URL.String(),
-			bytes.NewReader(bodyBytes),
-		)
-		if err != nil {
-			http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
-			return
-		}
+		r2 := cloneProxyRequest(handleCtx, r, bodyBytes)
 		reconnect := connect.NewReconnect(self.ProxyConnectTimeout)
 		response, err = tr.RoundTrip(r2)
 		if err == nil {
@@ -265,8 +266,9 @@ func (self *HttpProxy) handleHttp(w http.ResponseWriter, r *http.Request) {
 			chunked = true
 		}
 		if chunked {
-			f := w.(http.Flusher)
-			flush = f.Flush
+			if f, ok := w.(http.Flusher); ok {
+				flush = f.Flush
+			}
 		}
 		_, err := copyBufferWithTimeoutAndFlush(w, response.Body, nil, self.ProxyReadTimeout, self.ProxyWriteTimeout, flush)
 		if err != nil {
@@ -275,6 +277,33 @@ func (self *HttpProxy) handleHttp(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func cloneProxyRequest(ctx context.Context, r *http.Request, bodyBytes []byte) *http.Request {
+	r2 := r.Clone(ctx)
+	r2.RequestURI = ""
+	r2.Header = r.Header.Clone()
+	removeProxyRequestHeaders(r2.Header)
+	if len(bodyBytes) == 0 {
+		r2.Body = http.NoBody
+		r2.GetBody = func() (io.ReadCloser, error) {
+			return http.NoBody, nil
+		}
+		r2.ContentLength = 0
+	} else {
+		r2.Body = io.NopCloser(bytes.NewReader(bodyBytes))
+		r2.GetBody = func() (io.ReadCloser, error) {
+			return io.NopCloser(bytes.NewReader(bodyBytes)), nil
+		}
+		r2.ContentLength = int64(len(bodyBytes))
+	}
+	return r2
+}
+
+func removeProxyRequestHeaders(h http.Header) {
+	h.Del("Proxy-Authenticate")
+	h.Del("Proxy-Authorization")
+	h.Del("Proxy-Connection")
+}
+
 // for a hijacked connection
 func httpError(w io.Writer, statusCode int, err error) error {
 	errorMessage := err.Error()

http_test.go

@@ -0,0 +1,183 @@
+package proxy
+
+import (
+	"bufio"
+	"context"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestHttpProxyForwardsRequestMetadataAndBody(t *testing.T) {
+	type receivedRequest struct {
+		host               string
+		authorization      string
+		contentType        string
+		proxyAuthorization string
+		body               string
+	}
+
+	received := make(chan receivedRequest, 1)
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			t.Errorf("backend read body: %v", err)
+			return
+		}
+		received <- receivedRequest{
+			host:               r.Host,
+			authorization:      r.Header.Get("Authorization"),
+			contentType:        r.Header.Get("Content-Type"),
+			proxyAuthorization: r.Header.Get("Proxy-Authorization"),
+			body:               string(body),
+		}
+		w.Header().Set("X-Backend", "ok")
+		w.WriteHeader(http.StatusCreated)
+		_, _ = w.Write([]byte("created"))
+	}))
+	defer backend.Close()
+
+	proxy := NewHttpProxy()
+	proxy.ConnectDialWithRequest = func(r *http.Request, network string, addr string) (net.Conn, error) {
+		var d net.Dialer
+		return d.DialContext(r.Context(), network, addr)
+	}
+
+	req := httptest.NewRequest(http.MethodPost, backend.URL+"/resource", strings.NewReader("request-body"))
+	req.Host = "forwarded.example"
+	req.Header.Set("Authorization", "Bearer token")
+	req.Header.Set("Content-Type", "text/plain")
+	req.Header.Set("Proxy-Authorization", "proxy-secret")
+
+	rr := httptest.NewRecorder()
+	proxy.ServeHTTP(rr, req)
+
+	if rr.Code != http.StatusCreated {
+		t.Fatalf("status = %d, want %d; body=%q", rr.Code, http.StatusCreated, rr.Body.String())
+	}
+	if rr.Header().Get("X-Backend") != "ok" {
+		t.Fatalf("missing backend response header")
+	}
+
+	got := <-received
+	if got.host != "forwarded.example" {
+		t.Fatalf("host = %q, want forwarded.example", got.host)
+	}
+	if got.authorization != "Bearer token" {
+		t.Fatalf("authorization = %q", got.authorization)
+	}
+	if got.contentType != "text/plain" {
+		t.Fatalf("content-type = %q", got.contentType)
+	}
+	if got.proxyAuthorization != "" {
+		t.Fatalf("proxy authorization was forwarded: %q", got.proxyAuthorization)
+	}
+	if got.body != "request-body" {
+		t.Fatalf("body = %q", got.body)
+	}
+}
+
+func TestHttpProxyRejectsOversizedRequestBody(t *testing.T) {
+	proxy := NewHttpProxy()
+	proxy.MaxHttpBodyBytes = 3
+	proxy.ConnectDialWithRequest = func(r *http.Request, network string, addr string) (net.Conn, error) {
+		t.Fatalf("dial should not be called for oversized request")
+		return nil, context.Canceled
+	}
+
+	req := httptest.NewRequest(http.MethodPost, "http://example.test/upload", strings.NewReader("four"))
+	rr := httptest.NewRecorder()
+
+	proxy.ServeHTTP(rr, req)
+
+	if rr.Code != http.StatusRequestEntityTooLarge {
+		t.Fatalf("status = %d, want %d", rr.Code, http.StatusRequestEntityTooLarge)
+	}
+}
+
+func TestHttpProxyConnectTunnel(t *testing.T) {
+	backendAddr := startTCPBackend(t, func(conn net.Conn) {
+		defer conn.Close()
+		buf := make([]byte, 4)
+		if _, err := io.ReadFull(conn, buf); err != nil {
+			t.Errorf("backend read: %v", err)
+			return
+		}
+		if string(buf) != "ping" {
+			t.Errorf("backend got %q, want ping", string(buf))
+			return
+		}
+		_, _ = conn.Write([]byte("pong"))
+	})
+
+	proxy := NewHttpProxy()
+	proxy.ConnectDialWithRequest = func(r *http.Request, network string, addr string) (net.Conn, error) {
+		var d net.Dialer
+		return d.DialContext(r.Context(), "tcp", backendAddr)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	proxyAddr := freeTCPAddr(t)
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- proxy.ListenAndServe(ctx, "tcp", proxyAddr)
+	}()
+	waitForTCP(t, proxyAddr)
+	defer func() {
+		cancel()
+		select {
+		case err := <-errCh:
+			if err != nil && err != http.ErrServerClosed {
+				t.Fatalf("http proxy returned error: %v", err)
+			}
+		case <-time.After(2 * time.Second):
+			t.Fatalf("http proxy did not stop")
+		}
+	}()
+
+	conn, err := net.DialTimeout("tcp", proxyAddr, 2*time.Second)
+	if err != nil {
+		t.Fatalf("dial proxy: %v", err)
+	}
+	defer conn.Close()
+	if err := conn.SetDeadline(time.Now().Add(3 * time.Second)); err != nil {
+		t.Fatalf("set deadline: %v", err)
+	}
+
+	if _, err := io.WriteString(conn, "CONNECT example.test:443 HTTP/1.1\r\nHost: example.test:443\r\n\r\n"); err != nil {
+		t.Fatalf("write connect request: %v", err)
+	}
+	reader := bufio.NewReader(conn)
+	status, err := reader.ReadString('\n')
+	if err != nil {
+		t.Fatalf("read connect status: %v", err)
+	}
+	if !strings.Contains(status, "200") {
+		t.Fatalf("connect status = %q", status)
+	}
+	for {
+		line, err := reader.ReadString('\n')
+		if err != nil {
+			t.Fatalf("read connect headers: %v", err)
+		}
+		if line == "\r\n" {
+			break
+		}
+	}
+
+	if _, err := conn.Write([]byte("ping")); err != nil {
+		t.Fatalf("write tunnel: %v", err)
+	}
+	buf := make([]byte, 4)
+	if _, err := io.ReadFull(reader, buf); err != nil {
+		t.Fatalf("read tunnel: %v", err)
+	}
+	if string(buf) != "pong" {
+		t.Fatalf("tunnel got %q, want pong", string(buf))
+	}
+}

profile/cpu

@@ -93,6 +93,7 @@ func (self *SocksProxy) ListenAndServe(ctx context.Context, network string, addr
 }
 
 func (self *SocksProxy) connectHandle(ctx context.Context, writer io.Writer, r SocksRequest) error {
+	clientConn, _ := writer.(net.Conn)
 	proxyConn, err := self.ConnectDialWithRequest(ctx, r, "tcp", r.DestAddr.String())
 	if err != nil {
 		resp := mapDialErrorToSocksReply(err)
@@ -103,6 +104,9 @@ func (self *SocksProxy) connectHandle(ctx context.Context, writer io.Writer, r S
 	defer handleCancel()
 	go connect.HandleError(func() {
 		defer proxyConn.Close()
+		if clientConn != nil {
+			defer clientConn.Close()
+		}
 		select {
 		case <-handleCtx.Done():
 		}
@@ -131,8 +135,9 @@ func (self *SocksProxy) Valid(username string, password string, userAddr string)
 
 // socks.NameResolver
 func (self *SocksProxy) Resolve(ctx context.Context, name string) (context.Context, net.IP, error) {
-	// names are not resolved locally
-	return ctx, net.ParseIP("0.0.0.0").To4(), nil
+	// Names are resolved by the proxied dialer. Returning nil preserves the FQDN
+	// in the request address instead of replacing it with a local resolver result.
+	return ctx, nil, nil
 }
 
 func mapDialErrorToSocksReply(err error) uint8 {