From 6e0513a38655dd90594ba682eebf5c0ae622983a Mon Sep 17 00:00:00 2001 From: Dwi Siswanto Date: Mon, 14 Sep 2026 04:01:28 +0700 Subject: [PATCH 1/2] fix: body truncation caused by async logging Capture HTTP bodies as forwarding reads them and give the logger independent readers. This prevents logging from consuming bytes needed by the upstream server or downstream client. response processing could therefore read, replace, or close bodies while the transport was using them. Verbose request logging also consumed the body without restoring it. Setting ContentLength to -1 changed framing but did not resolve the shared reader ownership. Add a body capture wrapper that records forwarding reads and publishes immutable data when capture completes. Preserve read errors and forward Close to the original body without holding the capture lock, allowing Close to unblock an active Read. Retain the existing 4096-byte response log limit without limiting the forwarded response, and leave upgraded connections unwrapped. Queue metadata snapshots and capture handles. Retain the request capture in its context so response logs can include the request body after the transport consumes it. Use completed trailer copies instead of cloning live request trailers during an early response. Restore the private request body after verbose inspection for subsequent writers. Keep HAR processing on the original request pointer for Martian context lookup, before publishing the asynchronous transaction. Preserve its existing synchronous behavior and log-and-continue error handling. Preserve incoming ContentLength instead of forcing unknown length. When rebuilding a request after match-and-replace, copy ContentLength, TransferEncoding, and Trailer from the parsed request and clear GetBody so it cannot replay the old body. Fixes #749 Signed-off-by: Dwi Siswanto --- pkg/logger/body.go | 100 ++++++++++ pkg/logger/logger.go | 120 ++++++++++-- pkg/logger/logger_test.go | 379 ++++++++++++++++++++++++++++++++++++++ proxy.go | 8 +- proxy_test.go | 143 ++++++++++++++ 5 files changed, 727 insertions(+), 23 deletions(-) create mode 100644 pkg/logger/body.go create mode 100644 pkg/logger/logger_test.go create mode 100644 proxy_test.go diff --git a/pkg/logger/body.go b/pkg/logger/body.go new file mode 100644 index 00000000..f40bfdad --- /dev/null +++ b/pkg/logger/body.go @@ -0,0 +1,100 @@ +package logger + +import ( + "bytes" + "io" + "net/http" + "sync" +) + +// bodyCapture observes forwarding reads. The logger only accesses data after done +// closes, so it never reads, replaces, or closes the live stream. +type bodyCapture struct { + original io.ReadCloser + trailer http.Header + limit int + mu sync.Mutex + activeReads int + closed bool + complete bool + data []byte + err error + capturedTrailer http.Header + done chan struct{} +} + +func captureBody(body io.ReadCloser, trailer http.Header, limit int) *bodyCapture { + c := &bodyCapture{original: body, trailer: trailer, limit: limit, done: make(chan struct{})} + if body == nil || body == http.NoBody { + c.original = http.NoBody + c.finish(nil) + } + + return c +} + +func (c *bodyCapture) Read(p []byte) (int, error) { + c.mu.Lock() + c.activeReads++ + c.mu.Unlock() + + n, err := c.original.Read(p) + + c.mu.Lock() + c.activeReads-- + + if !c.complete { + size := n + if c.limit > 0 && size > c.limit-len(c.data) { + size = c.limit - len(c.data) + } + + c.data = append(c.data, p[:size]...) + if err != nil || c.closed || (c.limit > 0 && len(c.data) == c.limit) { + c.finish(err) + } + } + + c.mu.Unlock() + + return n, err +} + +func (c *bodyCapture) Close() error { + // Close must unblock an active Read; never hold the capture lock around it. + err := c.original.Close() + + c.mu.Lock() + c.closed = true + if c.activeReads == 0 { + c.finish(nil) + } + c.mu.Unlock() + + return err +} + +// finish runs under mu, or during construction before the wrapper is published. +func (c *bodyCapture) finish(err error) { + if c.complete { + return + } + + c.complete = true + if err != io.EOF { + c.err = err + } + + // Trailer values are final only after the forwarding reader reaches EOF. + if err == io.EOF { + c.capturedTrailer = c.trailer.Clone() + } + + close(c.done) +} + +func (c *bodyCapture) snapshot() (io.ReadCloser, http.Header, error) { + <-c.done + + return io.NopCloser(bytes.NewReader(c.data)), c.capturedTrailer.Clone(), c.err +} diff --git a/pkg/logger/logger.go b/pkg/logger/logger.go index 96131b4b..c2cd9820 100644 --- a/pkg/logger/logger.go +++ b/pkg/logger/logger.go @@ -1,6 +1,8 @@ package logger import ( + "bytes" + "context" "fmt" "io" "net/http" @@ -22,9 +24,10 @@ import ( ) const ( - dataWithNewLine = "%s\n" - dataWithoutNewLine = "%s" - LoggerConfigFilename = "export-config.yaml" + dataWithNewLine = "%s\n" + dataWithoutNewLine = "%s" + LoggerConfigFilename = "export-config.yaml" + maxLoggedResponseBody = 4096 ) type OptionsLogger struct { @@ -46,7 +49,7 @@ type Store interface { type Logger struct { options *OptionsLogger - asyncqueue chan types.HTTPTransaction + asyncqueue chan logTransaction Store []Store sWriter OutputFileWriter // sWriter is the structured writer harLogger *har.Logger @@ -56,7 +59,7 @@ type Logger struct { func NewLogger(options *OptionsLogger) *Logger { logger := &Logger{ options: options, - asyncqueue: make(chan types.HTTPTransaction, 1000), + asyncqueue: make(chan logTransaction, 1000), } if options.Elastic.Addr != "" { store, err := elastic.New(options.Elastic) @@ -111,24 +114,40 @@ func NewLogger(options *OptionsLogger) *Logger { return logger } +// logTransaction contains metadata snapshots and bodies captured by forwarding. +type logTransaction struct { + types.HTTPTransaction + requestBody, responseBody *bodyCapture +} + +type requestBodyKey struct{} + // LogRequest and user data func (l *Logger) LogRequest(req *http.Request, userdata types.UserData) error { if req == nil { return nil } - // send to writer channel - l.asyncqueue <- types.HTTPTransaction{ - Userdata: userdata, - Request: req, - } - + // HAR requires the original request pointer for Martian context lookup. if l.harLogger != nil { if err := l.harLogger.ModifyRequest(req); err != nil { gologger.Error().Msgf("Could not modify HAR request: %s\n", err) } } + body := captureBody(req.Body, req.Trailer, 0) + if req.Body != nil && req.Body != http.NoBody { + req.Body = body + } + + *req = *req.WithContext(context.WithValue(req.Context(), requestBodyKey{}, body)) + loggedReq := req.Clone(req.Context()) + loggedReq.Body = http.NoBody + l.asyncqueue <- logTransaction{ + HTTPTransaction: types.HTTPTransaction{Userdata: userdata, Request: loggedReq}, + requestBody: body, + } + return nil } @@ -138,29 +157,76 @@ func (l *Logger) LogResponse(resp *http.Response, userdata types.UserData) error return nil } - // send to writer channel - l.asyncqueue <- types.HTTPTransaction{ - Userdata: userdata, - Response: resp, - Request: resp.Request, - } - if l.harLogger != nil { if err := l.harLogger.ModifyResponse(resp); err != nil { gologger.Error().Msgf("Could not modify HAR response: %s\n", err) } } + loggedResp := new(http.Response) + *loggedResp = *resp + loggedResp.Header = resp.Header.Clone() + loggedResp.Trailer = resp.Trailer.Clone() + loggedResp.Body = http.NoBody + + var requestBody, responseBody *bodyCapture + + if resp.Request != nil { + // An early response can arrive while the upload is still filling trailers. + // The worker uses the capture's final trailer copy instead. + request := *resp.Request + request.Trailer = nil + request.Body = http.NoBody + loggedResp.Request = request.Clone(request.Context()) + requestBody, _ = resp.Request.Context().Value(requestBodyKey{}).(*bodyCapture) + } + + // An upgraded body is a live connection, not an HTTP message body. + if resp.StatusCode != http.StatusSwitchingProtocols { + responseBody = captureBody(resp.Body, resp.Trailer, maxLoggedResponseBody) + if resp.Body != nil && resp.Body != http.NoBody { + resp.Body = responseBody + } + } + + l.asyncqueue <- logTransaction{ + HTTPTransaction: types.HTTPTransaction{Userdata: userdata, Response: loggedResp, Request: loggedResp.Request}, + requestBody: requestBody, + responseBody: responseBody, + } + return nil } // AsyncWrite data func (l *Logger) AsyncWrite() { - for httpData := range l.asyncqueue { + for transaction := range l.asyncqueue { + httpData := transaction.HTTPTransaction + if transaction.requestBody != nil { + body, trailer, err := transaction.requestBody.snapshot() + if err != nil { + gologger.Warning().Msgf("Error capturing request body: %s", err) + } + + httpData.Request.Body = body + httpData.Request.Trailer = trailer + } + + if transaction.responseBody != nil { + body, trailer, err := transaction.responseBody.snapshot() + if err != nil { + gologger.Warning().Msgf("Error capturing response body: %s", err) + } + + httpData.Response.Body = body + httpData.Response.Trailer = trailer + } + if httpData.Request == nil { // we can't do anything without request continue } + // we have better options to handle this // i.e Buffer reuse and normalizing request/response body (removing encoding etc) reqDump, err := httputil.DumpRequest(httpData.Request, true) @@ -173,13 +239,14 @@ func (l *Logger) AsyncWrite() { var respChain *pdhttpUtils.ResponseChain if httpData.Response != nil { - respChainx := pdhttpUtils.NewResponseChain(httpData.Response, 4096) + respChainx := pdhttpUtils.NewResponseChain(httpData.Response, maxLoggedResponseBody) if err := respChainx.Fill(); err == nil { respChain = respChainx } else { gologger.Warning().Msgf("responseChain: Error while dumping response: %s", err) } } + // debug log response if true if respChain != nil { if err := l.debugLogResponse(respChain); err != nil { @@ -201,6 +268,7 @@ func (l *Logger) AsyncWrite() { Timestamp: time.Now().Format(time.RFC3339), URL: httpData.Request.URL.String(), } + defer func() { if sData.Response != nil { // write to structured writer with whatever data we have @@ -210,17 +278,20 @@ func (l *Logger) AsyncWrite() { } } }() + sRequest, err := types.NewHttpRequestData(httpData.Request) if err != nil { gologger.Warning().Msgf("Error while creating request: %s", err) return } + sData.Request = sRequest if respChain != nil { sResponse, err := types.NewHttpResponseData(respChain) if err != nil { gologger.Warning().Msgf("Error while creating response: %s", err) } + sData.Response = sResponse } }() @@ -253,11 +324,13 @@ func (l *Logger) Close() { gologger.Error().Msgf("Could not close HAR logger: %s\n", err) } } + if l.sWriter != nil { if err := l.sWriter.Close(); err != nil { gologger.Error().Msgf("Could not close OutputFileWriter: %s", err) } } + close(l.asyncqueue) } @@ -266,9 +339,12 @@ func (l *Logger) debugLogRequest(reqdump []byte, req *http.Request) { if l.options.Verbosity >= types.VerbosityVeryVerbose { contentType := req.Header.Get("Content-Type") b, _ := io.ReadAll(req.Body) + req.Body = io.NopCloser(bytes.NewReader(b)) + if isASCIICheckRequired(contentType) && !govalidator.IsPrintableASCII(string(b)) { reqdump, _ = httputil.DumpRequest(req, false) } + gologger.Silent().Msgf("%s", string(reqdump)) } } @@ -283,6 +359,7 @@ func (l *Logger) debugLogResponse(respChain *pdhttpUtils.ResponseChain) error { gologger.Silent().Msgf("%s", respChain.FullResponse().String()) } } + return nil } @@ -300,12 +377,14 @@ func (l *Logger) storeWriter(outputdata types.HTTPTransaction) { } else { return } + outputdata.Name = fmt.Sprintf("%s%s-%s", outputdata.Userdata.Host, outputdata.PartSuffix, outputdata.Userdata.ID) if outputdata.Userdata.HasResponse && (!l.options.DumpRequest && !l.options.DumpResponse) { if outputdata.Userdata.Match != nil && *outputdata.Userdata.Match { outputdata.Name = outputdata.Name + ".match" } } + outputdata.Format = dataWithoutNewLine if !strings.HasSuffix(string(outputdata.Data), "\n") { outputdata.Format = dataWithNewLine @@ -321,6 +400,7 @@ func (l *Logger) storeWriter(outputdata types.HTTPTransaction) { outputdata.DataString = stringsutil.Truncate(outputdata.DataString, l.options.MaxSize) outputdata.RawData = []byte(stringsutil.Truncate(string(outputdata.RawData), l.options.MaxSize)) } + for _, store := range l.Store { err := store.Save(outputdata) if err != nil { diff --git a/pkg/logger/logger_test.go b/pkg/logger/logger_test.go new file mode 100644 index 00000000..21fd403d --- /dev/null +++ b/pkg/logger/logger_test.go @@ -0,0 +1,379 @@ +package logger + +import ( + "bufio" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "github.com/projectdiscovery/martian/v3" + "github.com/projectdiscovery/proxify/pkg/logger/har" + "github.com/projectdiscovery/proxify/pkg/types" +) + +func TestLogRequestPreservesForwardedBody(t *testing.T) { + for _, chunked := range []bool{false, true} { + t.Run(fmt.Sprintf("chunked=%t", chunked), func(t *testing.T) { + const body = "original POST body" + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read upstream body: %v", err) + } + if string(got) != body { + t.Errorf("upstream body = %q, want %q", got, body) + } + w.WriteHeader(http.StatusNoContent) + })) + defer upstream.Close() + framing := fmt.Sprintf("Content-Length: %d\r\n\r\n%s", len(body), body) + if chunked { + framing = fmt.Sprintf("Transfer-Encoding: chunked\r\n\r\n%x\r\n%s\r\n0\r\n\r\n", len(body), body) + } + req, err := http.ReadRequest(bufio.NewReader(strings.NewReader("POST " + upstream.URL + "/ HTTP/1.1\r\nHost: example.test\r\n" + framing))) + if err != nil { + t.Fatal(err) + } + defer req.Body.Close() + l := &Logger{options: &OptionsLogger{Verbosity: types.VerbosityVeryVerbose}, asyncqueue: make(chan logTransaction, 1)} + if err := l.LogRequest(req, types.UserData{}); err != nil { + t.Fatal(err) + } + // Logging runs concurrently with forwarding, using captured bytes only. + close(l.asyncqueue) + done := make(chan struct{}) + go func() { defer close(done); l.AsyncWrite() }() + defer func() { req.Body.Close(); <-done }() + req.RequestURI = "" + client := &http.Client{Timeout: 3 * time.Second} + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + }) + } +} + +func TestLogResponseUsesCapturedRequest(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "http://example.test/", strings.NewReader("request payload")) + l := &Logger{options: &OptionsLogger{}, asyncqueue: make(chan logTransaction, 2)} + if err := l.LogRequest(req, types.UserData{}); err != nil { + t.Fatal(err) + } + requestLog := <-l.asyncqueue + // The transport may consume and close the request before response logging starts. + if _, err := io.Copy(io.Discard, req.Body); err != nil { + t.Fatal(err) + } + req.Body.Close() + resp := &http.Response{Request: req, Body: io.NopCloser(strings.NewReader("response payload")), Header: make(http.Header)} + if err := l.LogResponse(resp, types.UserData{}); err != nil { + t.Fatal(err) + } + responseLog := <-l.asyncqueue + for name, capture := range map[string]*bodyCapture{"request log": requestLog.requestBody, "response log request": responseLog.requestBody} { + body, _, err := capture.snapshot() + if err != nil { + t.Fatal(err) + } + got, err := io.ReadAll(body) + if err != nil || string(got) != "request payload" { + t.Errorf("%s = %q, %v", name, got, err) + } + } + got, err := io.ReadAll(resp.Body) + if err != nil || string(got) != "response payload" { + t.Fatalf("live response = %q, %v", got, err) + } + resp.Body.Close() + loggedBody, _, err := responseLog.responseBody.snapshot() + if err != nil { + t.Fatal(err) + } + got, err = io.ReadAll(loggedBody) + if err != nil || string(got) != "response payload" { + t.Fatalf("logged response = %q, %v", got, err) + } + req.Header.Set("X-Test", "live") + resp.Header.Set("X-Test", "live") + if requestLog.Request.Header.Get("X-Test") != "" || responseLog.Response.Header.Get("X-Test") != "" { + t.Fatal("logger shares live headers") + } +} + +type failingBody struct { + read bool + closed bool +} + +func (b *failingBody) Read(p []byte) (int, error) { + if b.read { + return 0, io.EOF + } + b.read = true + return copy(p, "partial"), io.ErrUnexpectedEOF +} +func (b *failingBody) Close() error { b.closed = true; return nil } + +func TestLogRequestReadErrorPreservesFailure(t *testing.T) { + body := &failingBody{} + req := httptest.NewRequest(http.MethodPost, "http://example.test/", nil) + req.Body = body + l := &Logger{asyncqueue: make(chan logTransaction, 1)} + if err := l.LogRequest(req, types.UserData{}); err != nil { + t.Fatal(err) + } + got, err := io.ReadAll(req.Body) + if string(got) != "partial" || !errors.Is(err, io.ErrUnexpectedEOF) { + t.Fatalf("forwarded body = %q, %v", got, err) + } + req.Body.Close() + if !body.closed { + t.Fatal("original body was not closed") + } + logged := <-l.asyncqueue + _, _, err = logged.requestBody.snapshot() + if !errors.Is(err, io.ErrUnexpectedEOF) { + t.Fatalf("captured error = %v", err) + } +} + +func TestLogResponsePreservesUpgradeBody(t *testing.T) { + body := &failingBody{} + resp := &http.Response{StatusCode: http.StatusSwitchingProtocols, Body: body} + l := &Logger{asyncqueue: make(chan logTransaction, 1)} + if err := l.LogResponse(resp, types.UserData{}); err != nil { + t.Fatal(err) + } + if body.read || body.closed || resp.Body != body { + t.Fatal("logging touched upgraded connection") + } +} + +func TestLogResponsePreservesBodyBeyondLogLimit(t *testing.T) { + payload := strings.Repeat("response", 2048) + resp := &http.Response{Body: io.NopCloser(strings.NewReader(payload))} + l := &Logger{asyncqueue: make(chan logTransaction, 1)} + if err := l.LogResponse(resp, types.UserData{}); err != nil { + t.Fatal(err) + } + logged := <-l.asyncqueue + got, err := io.ReadAll(resp.Body) + resp.Body.Close() + if err != nil || string(got) != payload { + t.Fatalf("live body length = %d, error = %v", len(got), err) + } + body, _, err := logged.responseBody.snapshot() + if err != nil { + t.Fatal(err) + } + got, err = io.ReadAll(body) + if err != nil || len(got) != 4096 { + t.Fatalf("logged body length = %d, error = %v", len(got), err) + } +} + +func TestLoggingDoesNotReadLiveBodies(t *testing.T) { + for _, response := range []bool{false, true} { + t.Run(fmt.Sprintf("response=%t", response), func(t *testing.T) { + body := &failingBody{} + l := &Logger{asyncqueue: make(chan logTransaction, 1)} + var err error + if response { + err = l.LogResponse(&http.Response{Body: body}, types.UserData{}) + } else { + req := httptest.NewRequest(http.MethodPost, "http://example.test/", nil) + req.Body = body + err = l.LogRequest(req, types.UserData{}) + } + if err != nil || body.read { + t.Fatalf("logging read live body: read=%t, error=%v", body.read, err) + } + }) + } +} + +func TestHARLoggingPreservesMartianContext(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "http://example.test/", strings.NewReader("request payload")) + _, remove, err := martian.TestContext(req, nil, nil) + if err != nil { + t.Fatal(err) + } + defer remove() + req.Header.Set("Content-Type", "text/plain") + path := filepath.Join(t.TempDir(), "traffic.har") + harLogger, err := har.NewLogger(path, time.Hour) + if err != nil { + t.Fatal(err) + } + l := &Logger{harLogger: harLogger, asyncqueue: make(chan logTransaction, 2)} + if err := l.LogRequest(req, types.UserData{}); err != nil { + t.Fatal(err) + } + got, err := io.ReadAll(req.Body) + req.Body.Close() + if err != nil || string(got) != "request payload" { + t.Fatalf("forwarded request = %q, %v", got, err) + } + resp := &http.Response{StatusCode: http.StatusOK, Proto: "HTTP/1.1", ProtoMajor: 1, ProtoMinor: 1, Request: req, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("response payload"))} + resp.Header.Set("Content-Type", "text/plain") + if err := l.LogResponse(resp, types.UserData{}); err != nil { + t.Fatal(err) + } + got, err = io.ReadAll(resp.Body) + resp.Body.Close() + if err != nil || string(got) != "response payload" { + t.Fatalf("forwarded response = %q, %v", got, err) + } + if err := harLogger.Close(); err != nil { + t.Fatal(err) + } + data, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if !json.Valid(data) || !strings.Contains(string(data), "request payload") || !strings.Contains(string(data), base64.StdEncoding.EncodeToString([]byte("response payload"))) { + t.Fatalf("HAR does not contain both bodies: %s", data) + } +} + +func TestBodyCaptureCloseUnblocksRead(t *testing.T) { + reader, writer := io.Pipe() + defer writer.Close() + body := captureBody(reader, nil, 0) + done := make(chan struct{}) + go func() { defer close(done); _, _ = io.Copy(io.Discard, body) }() + if err := body.Close(); err != nil { + t.Fatal(err) + } + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("Close did not unblock Read") + } + _, _, _ = body.snapshot() +} + +func TestBodyCaptureStreamsBeforeEOF(t *testing.T) { + reader, writer := io.Pipe() + defer writer.Close() + body := captureBody(reader, nil, 0) + defer body.Close() + wrote := make(chan error, 1) + go func() { _, err := io.WriteString(writer, "prefix"); wrote <- err }() + got := make([]byte, len("prefix")) + if _, err := io.ReadFull(body, got); err != nil { + t.Fatal(err) + } + if string(got) != "prefix" { + t.Fatalf("forwarded prefix = %q", got) + } + if err := <-wrote; err != nil { + t.Fatal(err) + } + select { + case <-body.done: + t.Fatal("capture completed before EOF") + default: + } + writer.Close() + if _, err := io.Copy(io.Discard, body); err != nil { + t.Fatal(err) + } + snapshot, _, err := body.snapshot() + if err != nil { + t.Fatal(err) + } + got, err = io.ReadAll(snapshot) + if err != nil || string(got) != "prefix" { + t.Fatalf("snapshot = %q, %v", got, err) + } +} + +func TestHARLoggingErrorDoesNotStopRequestLogging(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "http://example.test/", strings.NewReader("payload")) + req.Header.Set("Content-Type", "text/plain") + _, remove, err := martian.TestContext(req, nil, nil) + if err != nil { + t.Fatal(err) + } + defer remove() + harLogger, err := har.NewLogger(filepath.Join(t.TempDir(), "traffic.har"), time.Hour) + if err != nil { + t.Fatal(err) + } + defer harLogger.Close() + l := &Logger{harLogger: harLogger, asyncqueue: make(chan logTransaction, 2)} + if err := l.LogRequest(req, types.UserData{}); err != nil { + t.Fatal(err) + } + // Reusing the same Martian request ID causes a HAR duplicate-ID error. + if err := l.LogRequest(req, types.UserData{}); err != nil { + t.Fatalf("HAR error escaped logger: %v", err) + } + got, err := io.ReadAll(req.Body) + req.Body.Close() + if err != nil || string(got) != "payload" { + t.Fatalf("forwarded body = %q, %v", got, err) + } + if len(l.asyncqueue) != 2 { + t.Fatal("HAR error prevented normal request logging") + } +} + +type trailerBody struct { + trailer http.Header + started chan struct{} +} + +func (b *trailerBody) Read([]byte) (int, error) { + close(b.started) + for i := 0; i < 1000; i++ { + b.trailer.Set(fmt.Sprintf("X-Trailer-%d", i), "value") + runtime.Gosched() + } + return 0, io.EOF +} +func (*trailerBody) Close() error { return nil } + +func TestLogEarlyResponseDuringRequestTrailers(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "http://example.test/", nil) + req.Trailer = make(http.Header) + original := &trailerBody{trailer: req.Trailer, started: make(chan struct{})} + req.Body = original + l := &Logger{asyncqueue: make(chan logTransaction, 101)} + if err := l.LogRequest(req, types.UserData{}); err != nil { + t.Fatal(err) + } + readDone := make(chan error, 1) + go func() { _, err := io.Copy(io.Discard, req.Body); readDone <- err }() + <-original.started + for i := 0; i < 100; i++ { + resp := &http.Response{Request: req, Body: http.NoBody} + if err := l.LogResponse(resp, types.UserData{}); err != nil { + t.Fatal(err) + } + } + if err := <-readDone; err != nil { + t.Fatal(err) + } + req.Body.Close() + for i := 0; i < 101; i++ { + transaction := <-l.asyncqueue + _, trailer, err := transaction.requestBody.snapshot() + if err != nil || trailer.Get("X-Trailer-999") != "value" { + t.Fatalf("captured trailer missing, error = %v", err) + } + } +} diff --git a/proxy.go b/proxy.go index 75174ac0..181ce38e 100644 --- a/proxy.go +++ b/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() @@ -369,7 +366,12 @@ func (p *Proxy) MatchReplaceRequest(req *http.Request) error { req.Method = requestNew.Method req.Header = requestNew.Header req.Body = requestNew.Body + req.ContentLength = requestNew.ContentLength + req.TransferEncoding = requestNew.TransferEncoding + req.Trailer = requestNew.Trailer + req.GetBody = nil req.URL = requestNew.URL + return nil } diff --git a/proxy_test.go b/proxy_test.go new file mode 100644 index 00000000..6b0ab904 --- /dev/null +++ b/proxy_test.go @@ -0,0 +1,143 @@ +package proxify + +import ( + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/projectdiscovery/martian/v3" + "github.com/projectdiscovery/proxify/pkg/logger" + "github.com/projectdiscovery/proxify/pkg/logger/elastic" + "github.com/projectdiscovery/proxify/pkg/logger/kafka" + "github.com/projectdiscovery/proxify/pkg/types" +) + +func TestProxyForwardsPOSTBody(t *testing.T) { + for _, chunked := range []bool{false, true} { + name := "content-length" + if chunked { + name = "chunked" + } + t.Run(name, func(t *testing.T) { + payload := strings.Repeat("POST payload\x00\r\n", 8192) + prefixReceived := make(chan struct{}, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + prefix := make([]byte, 1024) + if _, err := io.ReadFull(r.Body, prefix); err != nil { + t.Errorf("read upstream prefix: %v", err) + } + if chunked { + prefixReceived <- struct{}{} + } + rest, err := io.ReadAll(r.Body) + body := append(prefix, rest...) + if err != nil { + t.Errorf("read upstream body: %v", err) + } + if string(body) != payload { + t.Errorf("upstream body differs: got %d bytes, want %d", len(body), len(payload)) + } + if chunked { + if len(r.TransferEncoding) != 1 || r.TransferEncoding[0] != "chunked" { + t.Errorf("TransferEncoding = %v", r.TransferEncoding) + } + } else if r.ContentLength != int64(len(payload)) { + t.Errorf("ContentLength = %d, want %d", r.ContentLength, len(payload)) + } + _, _ = w.Write([]byte("complete response")) + })) + defer upstream.Close() + l := logger.NewLogger(&logger.OptionsLogger{Elastic: &elastic.Options{}, Kafka: &kafka.Options{}, Verbosity: types.VerbositySilent}) + defer l.Close() + p := &Proxy{options: &Options{}, logger: l} + hp := martian.NewProxy() + transport := &http.Transport{} + defer transport.CloseIdleConnections() + hp.SetRoundTripper(transport) + hp.SetRequestModifier(p) + hp.SetResponseModifier(p) + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + done := make(chan struct{}) + go func() { defer close(done); _ = hp.Serve(listener) }() + defer func() { listener.Close(); hp.Close(); <-done }() + proxyURL, err := url.Parse("http://" + listener.Addr().String()) + if err != nil { + t.Fatal(err) + } + clientTransport := &http.Transport{Proxy: http.ProxyURL(proxyURL)} + defer clientTransport.CloseIdleConnections() + client := &http.Client{Transport: clientTransport, Timeout: 5 * time.Second} + for i := 0; i < 3; i++ { + req, err := http.NewRequest(http.MethodPost, upstream.URL, strings.NewReader(payload)) + if err != nil { + t.Fatal(err) + } + var writeDone chan error + if chunked { + reader, writer := io.Pipe() + req.Body = reader + req.GetBody = nil + req.ContentLength = -1 + writeDone = make(chan error, 1) + go func() { + _, err := io.WriteString(writer, payload[:1024]) + if err == nil { + // The proxy must forward the prefix before the upload finishes. + select { + case <-prefixReceived: + _, err = io.WriteString(writer, payload[1024:]) + case <-time.After(3 * time.Second): + err = fmt.Errorf("upstream did not receive prefix before EOF") + } + } + _ = writer.CloseWithError(err) + writeDone <- err + }() + } + resp, err := client.Do(req) + if writeDone != nil { + if writeErr := <-writeDone; writeErr != nil { + t.Error(writeErr) + } + } + if err != nil { + t.Fatal(err) + } + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + if err != nil { + t.Fatal(err) + } + if string(body) != "complete response" { + t.Fatalf("response body = %q", body) + } + } + }) + } +} + +func TestMatchReplaceRequestUpdatesFraming(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "http://example.test/", strings.NewReader("old")) + req.Header.Set("Content-Length", "3") + p := &Proxy{options: &Options{RequestMatchReplaceDSL: []string{`replace(replace(request, 'old', 'longer'), 'Content-Length: 3', 'Content-Length: 6')`}}} + if err := p.MatchReplaceRequest(req); err != nil { + t.Fatal(err) + } + if req.ContentLength != 6 { + t.Fatalf("ContentLength = %d, want 6", req.ContentLength) + } + body, err := io.ReadAll(req.Body) + req.Body.Close() + if err != nil || string(body) != "longer" { + t.Fatalf("body = %q, %v", body, err) + } +} From 5cdaf33b624e199bea3375c8c903900a67d22bf1 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Sun, 13 Sep 2026 23:30:30 +0200 Subject: [PATCH 2/2] harden logging --- pkg/logger/body.go | 26 ++++++++- pkg/logger/har/har.go | 8 +++ pkg/logger/logger.go | 118 ++++++++++++++++++++++++++------------ pkg/logger/logger_test.go | 112 ++++++++++++++++++++++++++++++------ proxy.go | 7 +-- proxy_test.go | 116 ++++++++++++++++++++++++++++++++++++- 6 files changed, 324 insertions(+), 63 deletions(-) diff --git a/pkg/logger/body.go b/pkg/logger/body.go index f40bfdad..a8339c60 100644 --- a/pkg/logger/body.go +++ b/pkg/logger/body.go @@ -2,9 +2,16 @@ package logger import ( "bytes" + "errors" "io" "net/http" "sync" + "time" +) + +var ( + errCaptureTimeout = errors.New("body capture timed out") + bodyCaptureWait = 30 * time.Second ) // bodyCapture observes forwarding reads. The logger only accesses data after done @@ -94,7 +101,22 @@ func (c *bodyCapture) finish(err error) { } func (c *bodyCapture) snapshot() (io.ReadCloser, http.Header, error) { - <-c.done + timer := time.NewTimer(bodyCaptureWait) + defer timer.Stop() + + select { + case <-c.done: + case <-timer.C: + c.mu.Lock() + c.finish(errCaptureTimeout) + c.mu.Unlock() + } + + c.mu.Lock() + data := append([]byte(nil), c.data...) + trailer := c.capturedTrailer.Clone() + err := c.err + c.mu.Unlock() - return io.NopCloser(bytes.NewReader(c.data)), c.capturedTrailer.Clone(), c.err + return io.NopCloser(bytes.NewReader(data)), trailer, err } diff --git a/pkg/logger/har/har.go b/pkg/logger/har/har.go index bbe5c3db..7201d836 100644 --- a/pkg/logger/har/har.go +++ b/pkg/logger/har/har.go @@ -75,6 +75,14 @@ func (l *Logger) ModifyResponse(resp *http.Response) error { return l.martianHarLogger.ModifyResponse(resp) } +func (l *Logger) RecordRequest(id string, req *http.Request) error { + return l.martianHarLogger.RecordRequest(id, req) +} + +func (l *Logger) RecordResponse(id string, resp *http.Response) error { + return l.martianHarLogger.RecordResponse(id, resp) +} + func (l *Logger) Flush() error { return l.writer.append(l.martianHarLogger.ExportAndReset()) } diff --git a/pkg/logger/logger.go b/pkg/logger/logger.go index c2cd9820..cc5c4a95 100644 --- a/pkg/logger/logger.go +++ b/pkg/logger/logger.go @@ -12,6 +12,7 @@ import ( "github.com/asaskevich/govalidator" "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/martian/v3" "github.com/projectdiscovery/proxify/pkg/logger/elastic" "github.com/projectdiscovery/proxify/pkg/logger/file" "github.com/projectdiscovery/proxify/pkg/logger/har" @@ -27,6 +28,7 @@ const ( dataWithNewLine = "%s\n" dataWithoutNewLine = "%s" LoggerConfigFilename = "export-config.yaml" + maxLoggedRequestBody = 4096 maxLoggedResponseBody = 4096 ) @@ -53,6 +55,7 @@ type Logger struct { Store []Store sWriter OutputFileWriter // sWriter is the structured writer harLogger *har.Logger + writeDone chan struct{} } // NewLogger instance @@ -110,7 +113,11 @@ func NewLogger(options *OptionsLogger) *Logger { } } - go logger.AsyncWrite() + logger.writeDone = make(chan struct{}) + go func() { + defer close(logger.writeDone) + logger.AsyncWrite() + }() return logger } @@ -118,24 +125,29 @@ func NewLogger(options *OptionsLogger) *Logger { type logTransaction struct { types.HTTPTransaction requestBody, responseBody *bodyCapture + martianID string } type requestBodyKey struct{} +func requestMartianID(req *http.Request) string { + if req == nil { + return "" + } + ctx := martian.NewContext(req) + if ctx == nil { + return "" + } + return ctx.ID() +} + // LogRequest and user data func (l *Logger) LogRequest(req *http.Request, userdata types.UserData) error { if req == nil { return nil } - // HAR requires the original request pointer for Martian context lookup. - if l.harLogger != nil { - if err := l.harLogger.ModifyRequest(req); err != nil { - gologger.Error().Msgf("Could not modify HAR request: %s\n", err) - } - } - - body := captureBody(req.Body, req.Trailer, 0) + body := captureBody(req.Body, req.Trailer, maxLoggedRequestBody) if req.Body != nil && req.Body != http.NoBody { req.Body = body } @@ -146,6 +158,7 @@ func (l *Logger) LogRequest(req *http.Request, userdata types.UserData) error { l.asyncqueue <- logTransaction{ HTTPTransaction: types.HTTPTransaction{Userdata: userdata, Request: loggedReq}, requestBody: body, + martianID: requestMartianID(req), } return nil @@ -157,12 +170,6 @@ func (l *Logger) LogResponse(resp *http.Response, userdata types.UserData) error return nil } - if l.harLogger != nil { - if err := l.harLogger.ModifyResponse(resp); err != nil { - gologger.Error().Msgf("Could not modify HAR response: %s\n", err) - } - } - loggedResp := new(http.Response) *loggedResp = *resp loggedResp.Header = resp.Header.Clone() @@ -170,10 +177,12 @@ func (l *Logger) LogResponse(resp *http.Response, userdata types.UserData) error loggedResp.Body = http.NoBody var requestBody, responseBody *bodyCapture + var martianID string if resp.Request != nil { // An early response can arrive while the upload is still filling trailers. // The worker uses the capture's final trailer copy instead. + martianID = requestMartianID(resp.Request) request := *resp.Request request.Trailer = nil request.Body = http.NoBody @@ -193,11 +202,33 @@ func (l *Logger) LogResponse(resp *http.Response, userdata types.UserData) error HTTPTransaction: types.HTTPTransaction{Userdata: userdata, Response: loggedResp, Request: loggedResp.Request}, requestBody: requestBody, responseBody: responseBody, + martianID: martianID, } return nil } +func (l *Logger) writeHAR(transaction logTransaction, httpData types.HTTPTransaction) { + if l.harLogger == nil || transaction.martianID == "" { + return + } + + // Request logs own the HAR request entry so a later response cannot duplicate the ID. + if httpData.Response == nil { + if httpData.Request == nil { + return + } + if err := l.harLogger.RecordRequest(transaction.martianID, httpData.Request); err != nil { + gologger.Error().Msgf("Could not modify HAR request: %s\n", err) + } + return + } + + if err := l.harLogger.RecordResponse(transaction.martianID, httpData.Response); err != nil { + gologger.Error().Msgf("Could not modify HAR response: %s\n", err) + } +} + // AsyncWrite data func (l *Logger) AsyncWrite() { for transaction := range l.asyncqueue { @@ -207,9 +238,10 @@ func (l *Logger) AsyncWrite() { if err != nil { gologger.Warning().Msgf("Error capturing request body: %s", err) } - - httpData.Request.Body = body - httpData.Request.Trailer = trailer + if httpData.Request != nil { + httpData.Request.Body = body + httpData.Request.Trailer = trailer + } } if transaction.responseBody != nil { @@ -217,11 +249,14 @@ func (l *Logger) AsyncWrite() { if err != nil { gologger.Warning().Msgf("Error capturing response body: %s", err) } - - httpData.Response.Body = body - httpData.Response.Trailer = trailer + if httpData.Response != nil { + httpData.Response.Body = body + httpData.Response.Trailer = trailer + } } + l.writeHAR(transaction, httpData) + if httpData.Request == nil { // we can't do anything without request continue @@ -319,6 +354,11 @@ func (l *Logger) AsyncWrite() { // Close logger instance func (l *Logger) Close() { + close(l.asyncqueue) + if l.writeDone != nil { + <-l.writeDone + } + if l.harLogger != nil { if err := l.harLogger.Close(); err != nil { gologger.Error().Msgf("Could not close HAR logger: %s\n", err) @@ -330,34 +370,36 @@ func (l *Logger) Close() { gologger.Error().Msgf("Could not close OutputFileWriter: %s", err) } } - - close(l.asyncqueue) } // debugLogRequest logs the request to the console if debugging is enabled func (l *Logger) debugLogRequest(reqdump []byte, req *http.Request) { - if l.options.Verbosity >= types.VerbosityVeryVerbose { - contentType := req.Header.Get("Content-Type") - b, _ := io.ReadAll(req.Body) - req.Body = io.NopCloser(bytes.NewReader(b)) + if l.options == nil || l.options.Verbosity < types.VerbosityVeryVerbose { + return + } - if isASCIICheckRequired(contentType) && !govalidator.IsPrintableASCII(string(b)) { - reqdump, _ = httputil.DumpRequest(req, false) - } + contentType := req.Header.Get("Content-Type") + b, _ := io.ReadAll(req.Body) + req.Body = io.NopCloser(bytes.NewReader(b)) - gologger.Silent().Msgf("%s", string(reqdump)) + if isASCIICheckRequired(contentType) && !govalidator.IsPrintableASCII(string(b)) { + reqdump, _ = httputil.DumpRequest(req, false) } + + gologger.Silent().Msgf("%s", string(reqdump)) } // debugLogResponse logs the response to the console if debugging is enabled func (l *Logger) debugLogResponse(respChain *pdhttpUtils.ResponseChain) error { - if l.options.Verbosity >= types.VerbosityVeryVerbose { - contentType := respChain.Response().Header.Get("Content-Type") - if isASCIICheckRequired(contentType) && !govalidator.IsPrintableASCII(conversion.String(respChain.Body().Bytes())) { - gologger.Silent().Msgf("%s", respChain.Headers().String()) - } else { - gologger.Silent().Msgf("%s", respChain.FullResponse().String()) - } + if l.options == nil || l.options.Verbosity < types.VerbosityVeryVerbose { + return nil + } + + contentType := respChain.Response().Header.Get("Content-Type") + if isASCIICheckRequired(contentType) && !govalidator.IsPrintableASCII(conversion.String(respChain.Body().Bytes())) { + gologger.Silent().Msgf("%s", respChain.Headers().String()) + } else { + gologger.Silent().Msgf("%s", respChain.FullResponse().String()) } return nil diff --git a/pkg/logger/logger_test.go b/pkg/logger/logger_test.go index 21fd403d..4b17131d 100644 --- a/pkg/logger/logger_test.go +++ b/pkg/logger/logger_test.go @@ -21,6 +21,13 @@ import ( "github.com/projectdiscovery/proxify/pkg/types" ) +func closeErr(t *testing.T, c io.Closer) { + t.Helper() + if err := c.Close(); err != nil { + t.Errorf("Close: %v", err) + } +} + func TestLogRequestPreservesForwardedBody(t *testing.T) { for _, chunked := range []bool{false, true} { t.Run(fmt.Sprintf("chunked=%t", chunked), func(t *testing.T) { @@ -44,7 +51,7 @@ func TestLogRequestPreservesForwardedBody(t *testing.T) { if err != nil { t.Fatal(err) } - defer req.Body.Close() + defer closeErr(t, req.Body) l := &Logger{options: &OptionsLogger{Verbosity: types.VerbosityVeryVerbose}, asyncqueue: make(chan logTransaction, 1)} if err := l.LogRequest(req, types.UserData{}); err != nil { t.Fatal(err) @@ -53,14 +60,14 @@ func TestLogRequestPreservesForwardedBody(t *testing.T) { close(l.asyncqueue) done := make(chan struct{}) go func() { defer close(done); l.AsyncWrite() }() - defer func() { req.Body.Close(); <-done }() + defer func() { closeErr(t, req.Body); <-done }() req.RequestURI = "" client := &http.Client{Timeout: 3 * time.Second} resp, err := client.Do(req) if err != nil { t.Fatal(err) } - resp.Body.Close() + closeErr(t, resp.Body) }) } } @@ -76,7 +83,7 @@ func TestLogResponseUsesCapturedRequest(t *testing.T) { if _, err := io.Copy(io.Discard, req.Body); err != nil { t.Fatal(err) } - req.Body.Close() + closeErr(t, req.Body) resp := &http.Response{Request: req, Body: io.NopCloser(strings.NewReader("response payload")), Header: make(http.Header)} if err := l.LogResponse(resp, types.UserData{}); err != nil { t.Fatal(err) @@ -96,7 +103,7 @@ func TestLogResponseUsesCapturedRequest(t *testing.T) { if err != nil || string(got) != "response payload" { t.Fatalf("live response = %q, %v", got, err) } - resp.Body.Close() + closeErr(t, resp.Body) loggedBody, _, err := responseLog.responseBody.snapshot() if err != nil { t.Fatal(err) @@ -138,7 +145,7 @@ func TestLogRequestReadErrorPreservesFailure(t *testing.T) { if string(got) != "partial" || !errors.Is(err, io.ErrUnexpectedEOF) { t.Fatalf("forwarded body = %q, %v", got, err) } - req.Body.Close() + closeErr(t, req.Body) if !body.closed { t.Fatal("original body was not closed") } @@ -170,7 +177,7 @@ func TestLogResponsePreservesBodyBeyondLogLimit(t *testing.T) { } logged := <-l.asyncqueue got, err := io.ReadAll(resp.Body) - resp.Body.Close() + closeErr(t, resp.Body) if err != nil || string(got) != payload { t.Fatalf("live body length = %d, error = %v", len(got), err) } @@ -222,7 +229,7 @@ func TestHARLoggingPreservesMartianContext(t *testing.T) { t.Fatal(err) } got, err := io.ReadAll(req.Body) - req.Body.Close() + closeErr(t, req.Body) if err != nil || string(got) != "request payload" { t.Fatalf("forwarded request = %q, %v", got, err) } @@ -232,10 +239,12 @@ func TestHARLoggingPreservesMartianContext(t *testing.T) { t.Fatal(err) } got, err = io.ReadAll(resp.Body) - resp.Body.Close() + closeErr(t, resp.Body) if err != nil || string(got) != "response payload" { t.Fatalf("forwarded response = %q, %v", got, err) } + close(l.asyncqueue) + l.AsyncWrite() if err := harLogger.Close(); err != nil { t.Fatal(err) } @@ -250,7 +259,7 @@ func TestHARLoggingPreservesMartianContext(t *testing.T) { func TestBodyCaptureCloseUnblocksRead(t *testing.T) { reader, writer := io.Pipe() - defer writer.Close() + defer closeErr(t, writer) body := captureBody(reader, nil, 0) done := make(chan struct{}) go func() { defer close(done); _, _ = io.Copy(io.Discard, body) }() @@ -267,9 +276,9 @@ func TestBodyCaptureCloseUnblocksRead(t *testing.T) { func TestBodyCaptureStreamsBeforeEOF(t *testing.T) { reader, writer := io.Pipe() - defer writer.Close() + defer closeErr(t, writer) body := captureBody(reader, nil, 0) - defer body.Close() + defer closeErr(t, body) wrote := make(chan error, 1) go func() { _, err := io.WriteString(writer, "prefix"); wrote <- err }() got := make([]byte, len("prefix")) @@ -287,7 +296,7 @@ func TestBodyCaptureStreamsBeforeEOF(t *testing.T) { t.Fatal("capture completed before EOF") default: } - writer.Close() + closeErr(t, writer) if _, err := io.Copy(io.Discard, body); err != nil { t.Fatal(err) } @@ -313,7 +322,7 @@ func TestHARLoggingErrorDoesNotStopRequestLogging(t *testing.T) { if err != nil { t.Fatal(err) } - defer harLogger.Close() + defer closeErr(t, harLogger) l := &Logger{harLogger: harLogger, asyncqueue: make(chan logTransaction, 2)} if err := l.LogRequest(req, types.UserData{}); err != nil { t.Fatal(err) @@ -323,7 +332,7 @@ func TestHARLoggingErrorDoesNotStopRequestLogging(t *testing.T) { t.Fatalf("HAR error escaped logger: %v", err) } got, err := io.ReadAll(req.Body) - req.Body.Close() + closeErr(t, req.Body) if err != nil || string(got) != "payload" { t.Fatalf("forwarded body = %q, %v", got, err) } @@ -368,7 +377,7 @@ func TestLogEarlyResponseDuringRequestTrailers(t *testing.T) { if err := <-readDone; err != nil { t.Fatal(err) } - req.Body.Close() + closeErr(t, req.Body) for i := 0; i < 101; i++ { transaction := <-l.asyncqueue _, trailer, err := transaction.requestBody.snapshot() @@ -377,3 +386,74 @@ func TestLogEarlyResponseDuringRequestTrailers(t *testing.T) { } } } + +func TestLogRequestCapsLoggedBody(t *testing.T) { + payload := strings.Repeat("request", 2048) + req := httptest.NewRequest(http.MethodPost, "http://example.test/", strings.NewReader(payload)) + l := &Logger{asyncqueue: make(chan logTransaction, 1)} + if err := l.LogRequest(req, types.UserData{}); err != nil { + t.Fatal(err) + } + logged := <-l.asyncqueue + got, err := io.ReadAll(req.Body) + closeErr(t, req.Body) + if err != nil || string(got) != payload { + t.Fatalf("live body length = %d, error = %v", len(got), err) + } + body, _, err := logged.requestBody.snapshot() + if err != nil { + t.Fatal(err) + } + got, err = io.ReadAll(body) + if err != nil || len(got) != maxLoggedRequestBody { + t.Fatalf("logged body length = %d, error = %v", len(got), err) + } +} + +func TestBodyCaptureSnapshotTimesOut(t *testing.T) { + original := bodyCaptureWait + bodyCaptureWait = 50 * time.Millisecond + t.Cleanup(func() { bodyCaptureWait = original }) + + reader, writer := io.Pipe() + defer func() { _ = writer.Close() }() + body := captureBody(reader, nil, 0) + defer closeErr(t, body) + + started := time.Now() + snapshot, _, err := body.snapshot() + if !errors.Is(err, errCaptureTimeout) { + t.Fatalf("snapshot error = %v", err) + } + if time.Since(started) > time.Second { + t.Fatal("snapshot blocked on unread body") + } + got, err := io.ReadAll(snapshot) + if err != nil || len(got) != 0 { + t.Fatalf("timed out snapshot = %q, %v", got, err) + } +} + +func TestHARLoggingDoesNotReadLiveBody(t *testing.T) { + body := &failingBody{} + req := httptest.NewRequest(http.MethodPost, "http://example.test/", nil) + req.Body = body + req.Header.Set("Content-Type", "text/plain") + _, remove, err := martian.TestContext(req, nil, nil) + if err != nil { + t.Fatal(err) + } + defer remove() + harLogger, err := har.NewLogger(filepath.Join(t.TempDir(), "traffic.har"), time.Hour) + if err != nil { + t.Fatal(err) + } + defer closeErr(t, harLogger) + l := &Logger{harLogger: harLogger, asyncqueue: make(chan logTransaction, 1)} + if err := l.LogRequest(req, types.UserData{}); err != nil { + t.Fatal(err) + } + if body.read || body.closed { + t.Fatal("HAR logging consumed the live request body") + } +} diff --git a/proxy.go b/proxy.go index 181ce38e..0f961ca7 100644 --- a/proxy.go +++ b/proxy.go @@ -377,9 +377,6 @@ func (p *Proxy) MatchReplaceRequest(req *http.Request) error { // MatchReplaceRequest strings or regex func (p *Proxy) MatchReplaceResponse(resp *http.Response) error { - // // Set Content-Length to zero to allow automatic calculation - resp.ContentLength = -1 - // lazy mode - dump request respdump, err := httputil.DumpResponse(resp, true) if err != nil { @@ -413,10 +410,12 @@ func (p *Proxy) MatchReplaceResponse(resp *http.Response) error { if err != nil { return err } + resp.ContentLength = responseNew.ContentLength + resp.TransferEncoding = responseNew.TransferEncoding + resp.Trailer = responseNew.Trailer if resp.ContentLength == 0 { resp.Header.Del("Content-Length") } - // resp.ContentLength = responseNew.ContentLength return nil } diff --git a/proxy_test.go b/proxy_test.go index 6b0ab904..14b5bf6e 100644 --- a/proxy_test.go +++ b/proxy_test.go @@ -7,6 +7,8 @@ import ( "net/http" "net/http/httptest" "net/url" + "os" + "path/filepath" "strings" "testing" "time" @@ -68,7 +70,13 @@ func TestProxyForwardsPOSTBody(t *testing.T) { } done := make(chan struct{}) go func() { defer close(done); _ = hp.Serve(listener) }() - defer func() { listener.Close(); hp.Close(); <-done }() + defer func() { + if err := listener.Close(); err != nil { + t.Errorf("Close: %v", err) + } + hp.Close() + <-done + }() proxyURL, err := url.Parse("http://" + listener.Addr().String()) if err != nil { t.Fatal(err) @@ -113,7 +121,9 @@ func TestProxyForwardsPOSTBody(t *testing.T) { t.Fatal(err) } body, err := io.ReadAll(resp.Body) - resp.Body.Close() + if err := resp.Body.Close(); err != nil { + t.Errorf("Close: %v", err) + } if err != nil { t.Fatal(err) } @@ -136,8 +146,108 @@ func TestMatchReplaceRequestUpdatesFraming(t *testing.T) { t.Fatalf("ContentLength = %d, want 6", req.ContentLength) } body, err := io.ReadAll(req.Body) - req.Body.Close() + if err := req.Body.Close(); err != nil { + t.Errorf("Close: %v", err) + } if err != nil || string(body) != "longer" { t.Fatalf("body = %q, %v", body, err) } } + +func TestMatchReplaceResponseUpdatesFraming(t *testing.T) { + resp := &http.Response{ + Status: "200 OK", + StatusCode: http.StatusOK, + Proto: "HTTP/1.1", + ProtoMajor: 1, + ProtoMinor: 1, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader("old")), + ContentLength: 3, + } + resp.Header.Set("Content-Length", "3") + p := &Proxy{options: &Options{ResponseMatchReplaceDSL: []string{`replace(replace(response, 'old', 'longer'), 'Content-Length: 3', 'Content-Length: 6')`}}} + if err := p.MatchReplaceResponse(resp); err != nil { + t.Fatal(err) + } + if resp.ContentLength != 6 { + t.Fatalf("ContentLength = %d, want 6", resp.ContentLength) + } + body, err := io.ReadAll(resp.Body) + if err := resp.Body.Close(); err != nil { + t.Errorf("Close: %v", err) + } + if err != nil || string(body) != "longer" { + t.Fatalf("body = %q, %v", body, err) + } +} + +func TestProxyForwardsPOSTBodyWithHAR(t *testing.T) { + payload := strings.Repeat("HAR payload", 1024) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read upstream body: %v", err) + } + if string(got) != payload { + t.Errorf("upstream body length = %d, want %d", len(got), len(payload)) + } + _, _ = w.Write([]byte("complete response")) + })) + defer upstream.Close() + harPath := filepath.Join(t.TempDir(), "traffic.har") + l := logger.NewLogger(&logger.OptionsLogger{Elastic: &elastic.Options{}, Kafka: &kafka.Options{}, Verbosity: types.VerbositySilent, OutputHar: harPath}) + p := &Proxy{options: &Options{}, logger: l} + hp := martian.NewProxy() + transport := &http.Transport{} + defer transport.CloseIdleConnections() + hp.SetRoundTripper(transport) + hp.SetRequestModifier(p) + hp.SetResponseModifier(p) + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + done := make(chan struct{}) + go func() { defer close(done); _ = hp.Serve(listener) }() + defer func() { + if err := listener.Close(); err != nil { + t.Errorf("Close: %v", err) + } + hp.Close() + <-done + }() + proxyURL, err := url.Parse("http://" + listener.Addr().String()) + if err != nil { + t.Fatal(err) + } + clientTransport := &http.Transport{Proxy: http.ProxyURL(proxyURL)} + defer clientTransport.CloseIdleConnections() + client := &http.Client{Transport: clientTransport, Timeout: 5 * time.Second} + req, err := http.NewRequest(http.MethodPost, upstream.URL, strings.NewReader(payload)) + if err != nil { + t.Fatal(err) + } + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + body, err := io.ReadAll(resp.Body) + if closeErr := resp.Body.Close(); closeErr != nil { + t.Errorf("Close: %v", closeErr) + } + if err != nil { + t.Fatal(err) + } + if string(body) != "complete response" { + t.Fatalf("response body = %q", body) + } + l.Close() + data, err := os.ReadFile(harPath) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(data), payload[:32]) { + t.Fatalf("HAR missing request body: %s", data) + } +}