martian: set CONNECT content length before response modifiers

Modify CONNECT response content length before it gets logged.
saucelabs · May 21, 2024 · deee6ea · deee6ea
1 parent bcc2f82
commit deee6ea
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/internal/martian/proxy_conn.go b/internal/martian/proxy_conn.go
@@ -240,6 +240,8 @@ func (p *proxyConn) handleConnectRequest(req *http.Request) error {
 		return p.writeErrorResponse(req, cerr)
 	}
 
+	res.ContentLength = -1
+
 	if err := p.modifyResponse(res); err != nil {
 		log.Debugf(ctx, "error modifying CONNECT response: %v", err)
 		return p.writeErrorResponse(req, err)
@@ -250,8 +252,6 @@ func (p *proxyConn) handleConnectRequest(req *http.Request) error {
 		return p.writeResponse(res)
 	}
 
-	res.ContentLength = -1
-
 	if err := p.tunnel("CONNECT", res, crw); err != nil {
 		log.Errorf(ctx, "CONNECT tunnel: %v", err)
 	}

diff --git a/internal/martian/proxy_test.go b/internal/martian/proxy_test.go
@@ -822,6 +822,9 @@ func TestIntegrationConnect(t *testing.T) { //nolint:tparallel // Subtests share
 		if got, want := res.StatusCode, 200; got != want {
 			t.Fatalf("res.StatusCode: got %d, want %d", got, want)
 		}
+		if res.ContentLength != -1 {
+			t.Fatalf("res.ContentLength: got %d, want -1", res.ContentLength)
+		}
 
 		if !tm.RequestModified() {
 			t.Error("tm.RequestModified(): got false, want true")