martian: use proxyConn.writeResponse() to write successful CONNECT re…

…sponses

Add writeConnectOKResponse() and modify proxyConn.writeResponse() to use it.

internal/martian/proxy_conn.go

@@ -267,14 +267,7 @@ func (p *proxyConn) handleUpgradeResponse(res *http.Response) error {
 }
 
 func (p *proxyConn) tunnel(name string, res *http.Response, crw io.ReadWriteCloser) error {
-	if _, err := p.brw.WriteString(connectResponse); err != nil {
-		err := fmt.Errorf("got error while writing response back to client: %w", err)
-		p.traceWroteResponse(res, err)
-		return err
-	}
-	if err := p.brw.Flush(); err != nil {
-		err := fmt.Errorf("got error while flushing response back to client: %w", err)
-		p.traceWroteResponse(res, err)
+	if err := p.writeResponse(res); err != nil {
 		return err
 	}
 	if err := drainBuffer(crw, p.brw.Reader); err != nil {
@@ -431,12 +424,15 @@ func (p *proxyConn) writeResponse(res *http.Response) error {
 	}
 
 	var err error
-	if req.Method == http.MethodHead && res.Body == http.NoBody {
+	switch {
+	case req.Method == http.MethodConnect && res.StatusCode/100 == 2:
+		err = writeConnectOKResponse(p.brw.Writer)
+	case req.Method == http.MethodHead && res.Body == http.NoBody:
 		// The http package is misbehaving when writing a HEAD response.
 		// See https://github.com/golang/go/issues/62015 for details.
 		// This works around the issue by writing the response manually.
 		err = writeHeadResponse(p.brw.Writer, res)
-	} else {
+	default:
 		// Add support for Server Sent Events - relay HTTP chunks and flush after each chunk.
 		// This is safe for events that are smaller than the buffer io.Copy uses (32KB).
 		// If the event is larger than the buffer, the event will be split into multiple chunks.

internal/martian/proxy_connect.go

@@ -184,7 +184,12 @@ func newConnectResponse(req *http.Request) *http.Response {
 	}
 }
 
-const connectResponse = "HTTP/1.1 200 OK\r\n\r\n"
+var connectOKResponse = []byte("HTTP/1.1 200 OK\r\n\r\n")
+
+func writeConnectOKResponse(w io.Writer) error {
+	_, err := w.Write(connectOKResponse)
+	return err
+}
 
 const terminateTLSHeader = "X-Martian-Terminate-Tls"
 

internal/martian/proxy_handler.go

@@ -171,16 +171,13 @@ func (p proxyHandler) tunnel(name string, rw http.ResponseWriter, req *http.Requ
 		}
 		defer conn.Close()
 
-		if _, err := brw.WriteString(connectResponse); err != nil {
-			err := fmt.Errorf("got error while writing response back to client: %w", err)
-			p.traceWroteResponse(res, err)
-			return err
-		}
-		if err := brw.Flush(); err != nil {
-			err := fmt.Errorf("got error while flushing response back to client: %w", err)
-			p.traceWroteResponse(res, err)
-			return err
+		pc := proxyConn{
+			Proxy: p.Proxy,
+			brw:   brw,
+			conn:  conn,
 		}
+		pc.writeResponse(res)
+
 		if err := drainBuffer(crw, brw.Reader); err != nil {
 			err := fmt.Errorf("got error while draining buffer: %w", err)
 			p.traceWroteResponse(res, err)