diff --git a/pkg/logger/body.go b/pkg/logger/body.go new file mode 100644 index 00000000..a8339c60 --- /dev/null +++ b/pkg/logger/body.go @@ -0,0 +1,122 @@ +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 +// 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) { + 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(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 96131b4b..cc5c4a95 100644 --- a/pkg/logger/logger.go +++ b/pkg/logger/logger.go @@ -1,6 +1,8 @@ package logger import ( + "bytes" + "context" "fmt" "io" "net/http" @@ -10,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" @@ -22,9 +25,11 @@ import ( ) const ( - dataWithNewLine = "%s\n" - dataWithoutNewLine = "%s" - LoggerConfigFilename = "export-config.yaml" + dataWithNewLine = "%s\n" + dataWithoutNewLine = "%s" + LoggerConfigFilename = "export-config.yaml" + maxLoggedRequestBody = 4096 + maxLoggedResponseBody = 4096 ) type OptionsLogger struct { @@ -46,17 +51,18 @@ 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 + writeDone chan struct{} } // NewLogger instance 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) @@ -107,26 +113,52 @@ func NewLogger(options *OptionsLogger) *Logger { } } - go logger.AsyncWrite() + logger.writeDone = make(chan struct{}) + go func() { + defer close(logger.writeDone) + logger.AsyncWrite() + }() return logger } +// logTransaction contains metadata snapshots and bodies captured by forwarding. +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 } - // send to writer channel - l.asyncqueue <- types.HTTPTransaction{ - Userdata: userdata, - Request: req, + body := captureBody(req.Body, req.Trailer, maxLoggedRequestBody) + if req.Body != nil && req.Body != http.NoBody { + req.Body = body } - if l.harLogger != nil { - if err := l.harLogger.ModifyRequest(req); err != nil { - gologger.Error().Msgf("Could not modify HAR request: %s\n", err) - } + *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, + martianID: requestMartianID(req), } return nil @@ -138,29 +170,98 @@ 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, + loggedResp := new(http.Response) + *loggedResp = *resp + loggedResp.Header = resp.Header.Clone() + loggedResp.Trailer = resp.Trailer.Clone() + 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 + loggedResp.Request = request.Clone(request.Context()) + requestBody, _ = resp.Request.Context().Value(requestBodyKey{}).(*bodyCapture) } - if l.harLogger != nil { - if err := l.harLogger.ModifyResponse(resp); err != nil { - gologger.Error().Msgf("Could not modify HAR response: %s\n", err) + // 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, + 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 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) + } + if httpData.Request != nil { + 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) + } + 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 } + // 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 +274,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 +303,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 +313,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 } }() @@ -248,41 +354,54 @@ 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) } } + if l.sWriter != nil { if err := l.sWriter.Close(); err != nil { 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) - if isASCIICheckRequired(contentType) && !govalidator.IsPrintableASCII(string(b)) { - reqdump, _ = httputil.DumpRequest(req, false) - } - gologger.Silent().Msgf("%s", string(reqdump)) + if l.options == nil || l.options.Verbosity < types.VerbosityVeryVerbose { + return + } + + 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)) } // 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 } @@ -300,12 +419,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 +442,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..4b17131d --- /dev/null +++ b/pkg/logger/logger_test.go @@ -0,0 +1,459 @@ +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 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) { + 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 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) + } + // 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() { 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) + } + closeErr(t, resp.Body) + }) + } +} + +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) + } + 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) + } + 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) + } + closeErr(t, resp.Body) + 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) + } + closeErr(t, req.Body) + 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) + closeErr(t, resp.Body) + 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) + closeErr(t, req.Body) + 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) + 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) + } + 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 closeErr(t, writer) + 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 closeErr(t, writer) + body := captureBody(reader, nil, 0) + defer closeErr(t, body) + 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: + } + closeErr(t, writer) + 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 closeErr(t, harLogger) + 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) + closeErr(t, req.Body) + 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) + } + closeErr(t, req.Body) + 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) + } + } +} + +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 75174ac0..0f961ca7 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,15 +366,17 @@ 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 } // 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 { @@ -411,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 new file mode 100644 index 00000000..14b5bf6e --- /dev/null +++ b/proxy_test.go @@ -0,0 +1,253 @@ +package proxify + +import ( + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "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() { + 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} + 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) + if err := resp.Body.Close(); err != nil { + t.Errorf("Close: %v", err) + } + 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) + 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) + } +}