From 0227154cb202345fe556f234dd3c1cec0eea4bbe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matczuk?= <michal.matczuk@saucelabs.com>
Date: Tue, 28 May 2024 14:39:44 +0200
Subject: [PATCH] martian(proxy): don't close connection on connect response

Fixes #813
---
 internal/martian/proxy_conn.go | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/internal/martian/proxy_conn.go b/internal/martian/proxy_conn.go
index c246d9df..b5b1002e 100644
--- a/internal/martian/proxy_conn.go
+++ b/internal/martian/proxy_conn.go
@@ -424,9 +424,19 @@ func (p *proxyConn) writeResponse(res *http.Response) error {
 		defer p.conn.SetWriteDeadline(time.Time{})
 	}
 
-	if req.Close || p.closing() {
+	if p.closing() {
 		res.Close = true
+	} else {
+		if req.Close {
+			res.Close = true
+		}
+		// Support CONNECT over HTTP/1.0.
+		// If connect is successful, the connection should not be closed.
+		if req.Method == http.MethodConnect && res.StatusCode/100 == 2 {
+			res.Close = false
+		}
 	}
+
 	if res.Close {
 		res.Header.Add("Connection", "close")
 	}