fix(proxy): preserve streamed request bodies

pkg/logger/logger.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/asaskevich/govalidator"
@@ -120,7 +121,7 @@ func (l *Logger) LogRequest(req *http.Request, userdata types.UserData) error {
 	// send to writer channel
 	l.asyncqueue <- types.HTTPTransaction{
 		Userdata: userdata,
-		Request:  req,
+		Request:  cloneRequestForLogging(req),
 	}
 
 	if l.harLogger != nil {
@@ -132,6 +133,56 @@ func (l *Logger) LogRequest(req *http.Request, userdata types.UserData) error {
 	return nil
 }
 
+// cloneRequestForLogging gives the asynchronous logger its own request body
+// while the original body continues to stream to the upstream server.
+func cloneRequestForLogging(req *http.Request) *http.Request {
+	clone := req.Clone(req.Context())
+	if req.Body == nil || req.Body == http.NoBody {
+		return clone
+	}
+
+	reader, writer := io.Pipe()
+	clone.Body = reader
+	clone.GetBody = nil
+	req.Body = &loggingBody{
+		body:   req.Body,
+		writer: writer,
+	}
+	return clone
+}
+
+type loggingBody struct {
+	body      io.ReadCloser
+	writer    *io.PipeWriter
+	closeOnce sync.Once
+}
+
+func (b *loggingBody) Read(data []byte) (int, error) {
+	n, readErr := b.body.Read(data)
+	if n > 0 {
+		// Logging is best effort. If its reader stops, forwarding must continue.
+		if _, writeErr := b.writer.Write(data[:n]); writeErr != nil {
+			b.closeLogging(nil)
+		}
+	}
+	if readErr != nil {
+		b.closeLogging(readErr)
+	}
+	return n, readErr
+}
+
+func (b *loggingBody) Close() error {
+	err := b.body.Close()
+	b.closeLogging(err)
+	return err
+}
+
+func (b *loggingBody) closeLogging(err error) {
+	b.closeOnce.Do(func() {
+		_ = b.writer.CloseWithError(err)
+	})
+}
+
 // LogResponse and user data
 func (l *Logger) LogResponse(resp *http.Response, userdata types.UserData) error {
 	if resp == nil {

pkg/logger/request_test.go

@@ -0,0 +1,102 @@
+package logger
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+func TestCloneRequestForLoggingStreamsIndependentBody(t *testing.T) {
+	const body = "username=admin&password=secret"
+	req := httptest.NewRequest(http.MethodPost, "http://example.com/login", strings.NewReader(body))
+	originalContentLength := req.ContentLength
+
+	logRequest := cloneRequestForLogging(req)
+	loggedBody := make(chan struct {
+		body []byte
+		err  error
+	}, 1)
+	go func() {
+		data, err := io.ReadAll(logRequest.Body)
+		loggedBody <- struct {
+			body []byte
+			err  error
+		}{body: data, err: err}
+	}()
+
+	forwardedBody, err := io.ReadAll(req.Body)
+	if err != nil {
+		t.Fatalf("read forwarding body: %v", err)
+	}
+	if err := req.Body.Close(); err != nil {
+		t.Fatalf("close forwarding body: %v", err)
+	}
+
+	logged := <-loggedBody
+	if logged.err != nil {
+		t.Fatalf("read logging body: %v", logged.err)
+	}
+	if string(forwardedBody) != body {
+		t.Fatalf("forwarded body = %q, want %q", forwardedBody, body)
+	}
+	if string(logged.body) != body {
+		t.Fatalf("logged body = %q, want %q", logged.body, body)
+	}
+	if req.ContentLength != originalContentLength {
+		t.Fatalf("forwarding content length = %d, want %d", req.ContentLength, originalContentLength)
+	}
+	if logRequest.ContentLength != originalContentLength {
+		t.Fatalf("logging content length = %d, want %d", logRequest.ContentLength, originalContentLength)
+	}
+}
+
+func TestCloneRequestForLoggingPreservesUnknownLength(t *testing.T) {
+	const body = "streamed-body"
+	req := httptest.NewRequest(http.MethodPost, "http://example.com/upload", strings.NewReader(body))
+	req.ContentLength = -1
+
+	logRequest := cloneRequestForLogging(req)
+	loggedBody := make(chan []byte, 1)
+	go func() {
+		data, _ := io.ReadAll(logRequest.Body)
+		loggedBody <- data
+	}()
+
+	forwardedBody, err := io.ReadAll(req.Body)
+	if err != nil {
+		t.Fatalf("read forwarding body: %v", err)
+	}
+	if err := req.Body.Close(); err != nil {
+		t.Fatalf("close forwarding body: %v", err)
+	}
+
+	if string(forwardedBody) != body {
+		t.Fatalf("forwarded body = %q, want %q", forwardedBody, body)
+	}
+	if string(<-loggedBody) != body {
+		t.Fatalf("logged body does not match %q", body)
+	}
+	if req.ContentLength != -1 || logRequest.ContentLength != -1 {
+		t.Fatalf("unknown content length changed: forwarding=%d logging=%d", req.ContentLength, logRequest.ContentLength)
+	}
+}
+
+func TestCloneRequestForLoggingDoesNotFailForwardingWhenLoggerStops(t *testing.T) {
+	const body = "forward-even-if-logging-stops"
+	req := httptest.NewRequest(http.MethodPost, "http://example.com/upload", strings.NewReader(body))
+
+	logRequest := cloneRequestForLogging(req)
+	if err := logRequest.Body.Close(); err != nil {
+		t.Fatalf("close logging body: %v", err)
+	}
+
+	forwardedBody, err := io.ReadAll(req.Body)
+	if err != nil {
+		t.Fatalf("read forwarding body after logger close: %v", err)
+	}
+	if string(forwardedBody) != body {
+		t.Fatalf("forwarded body = %q, want %q", forwardedBody, body)
+	}
+}

proxy.go

@@ -219,9 +219,6 @@ func NewProxy(options *Options) (*Proxy, error) {
 
 // ModifyRequest
 func (p *Proxy) ModifyRequest(req *http.Request) error {
-	// // Set Content-Length to zero to allow automatic calculation
-	req.ContentLength = -1
-
 	ctx := martian.NewContext(req)
 	// disable upgrading http connections to https by default
 	ctx.Session().MarkInsecure()

proxy_test.go

@@ -0,0 +1,84 @@
+package proxify
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/projectdiscovery/martian/v3"
+	proxifylogger "github.com/projectdiscovery/proxify/pkg/logger"
+	"github.com/projectdiscovery/proxify/pkg/logger/elastic"
+	"github.com/projectdiscovery/proxify/pkg/logger/kafka"
+)
+
+func TestModifyRequestPreservesContentLengthAndBody(t *testing.T) {
+	const body = "username=admin&password=secret"
+	type receivedRequest struct {
+		body             []byte
+		contentLength    int64
+		transferEncoding []string
+		err              error
+	}
+	received := make(chan receivedRequest, 1)
+	upstream := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
+		data, err := io.ReadAll(request.Body)
+		received <- receivedRequest{
+			body:             data,
+			contentLength:    request.ContentLength,
+			transferEncoding: request.TransferEncoding,
+			err:              err,
+		}
+		response.WriteHeader(http.StatusNoContent)
+	}))
+	t.Cleanup(upstream.Close)
+
+	request, err := http.NewRequest(http.MethodPost, upstream.URL+"/login", strings.NewReader(body))
+	if err != nil {
+		t.Fatalf("create request: %v", err)
+	}
+	originalContentLength := request.ContentLength
+	_, removeContext, err := martian.TestContext(request, nil, nil)
+	if err != nil {
+		t.Fatalf("create proxy request context: %v", err)
+	}
+	t.Cleanup(removeContext)
+	requestLogger := proxifylogger.NewLogger(&proxifylogger.OptionsLogger{
+		Elastic: &elastic.Options{},
+		Kafka:   &kafka.Options{},
+	})
+	t.Cleanup(requestLogger.Close)
+	proxy := &Proxy{
+		options: &Options{},
+		logger:  requestLogger,
+	}
+
+	if err := proxy.ModifyRequest(request); err != nil {
+		t.Fatalf("modify request: %v", err)
+	}
+	response, err := http.DefaultTransport.RoundTrip(request)
+	if err != nil {
+		t.Fatalf("forward request: %v", err)
+	}
+	if err := response.Body.Close(); err != nil {
+		t.Fatalf("close response body: %v", err)
+	}
+
+	forwarded := <-received
+	if forwarded.err != nil {
+		t.Fatalf("upstream read body: %v", forwarded.err)
+	}
+	if string(forwarded.body) != body {
+		t.Fatalf("forwarded body = %q, want %q", forwarded.body, body)
+	}
+	if request.ContentLength != originalContentLength {
+		t.Fatalf("content length = %d, want %d", request.ContentLength, originalContentLength)
+	}
+	if forwarded.contentLength != originalContentLength {
+		t.Fatalf("upstream content length = %d, want %d", forwarded.contentLength, originalContentLength)
+	}
+	if len(forwarded.transferEncoding) != 0 {
+		t.Fatalf("upstream transfer encoding = %v, want none", forwarded.transferEncoding)
+	}
+}