From 89349038aa41efd847ccf2113f8b3716bd27d221 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matczuk?= <michal.matczuk@saucelabs.com>
Date: Fri, 18 Nov 2022 13:09:25 +0100
Subject: [PATCH] On CONNECT check downstream proxy scheme and allow only
 HTTP(S)

The connect() function only supports HTTP CONNECT proxies [1].
This patch validates that the downstream proxy is an HTTP proxy.
It also refactors the function to prepare it for the patches that follow.

[1] https://www.rfc-editor.org/rfc/rfc7231#section-4.3.6
---
 proxy.go      | 40 +++++++++++++++++++++++++---------------
 proxy_test.go |  3 ++-
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/proxy.go b/proxy.go
index 0344c17d..f60c0292 100644
--- a/proxy.go
+++ b/proxy.go
@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"crypto/tls"
 	"errors"
+	"fmt"
 	"io"
 	"net"
 	"net/http"
@@ -637,33 +638,42 @@ func (p *Proxy) connect(req *http.Request) (*http.Response, net.Conn, error) {
 		proxyURL = u
 	}
 
-	if proxyURL != nil {
-		log.Debugf("martian: CONNECT with downstream proxy: %s", proxyURL.Host)
+	if proxyURL == nil {
+		log.Debugf("martian: CONNECT to host directly: %s", req.URL.Host)
 
-		conn, err := p.dial("tcp", proxyURL.Host)
+		conn, err := p.dial("tcp", req.URL.Host)
 		if err != nil {
 			return nil, nil, err
 		}
-		pbw := bufio.NewWriter(conn)
-		pbr := bufio.NewReader(conn)
 
-		req.Write(pbw)
-		pbw.Flush()
+		return proxyutil.NewResponse(200, nil, req), conn, nil
+	}
 
-		res, err := http.ReadResponse(pbr, req)
-		if err != nil {
-			return nil, nil, err
-		}
+	switch proxyURL.Scheme {
+	case "http", "https":
+		return p.connectHTTP(req, proxyURL)
+	default:
+		return nil, nil, fmt.Errorf("martian: unsupported proxy scheme: %s", proxyURL.Scheme)
+	}
+}
+
+func (p *Proxy) connectHTTP(req *http.Request, proxyURL *url.URL) (*http.Response, net.Conn, error) {
+	log.Debugf("martian: CONNECT with downstream HTTP proxy: %s", proxyURL.Host)
 
-		return res, conn, nil
+	conn, err := p.dial("tcp", proxyURL.Host)
+	if err != nil {
+		return nil, nil, err
 	}
+	pbw := bufio.NewWriter(conn)
+	pbr := bufio.NewReader(conn)
 
-	log.Debugf("martian: CONNECT to host directly: %s", req.URL.Host)
+	req.Write(pbw)
+	pbw.Flush()
 
-	conn, err := p.dial("tcp", req.URL.Host)
+	res, err := http.ReadResponse(pbr, req)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	return proxyutil.NewResponse(200, nil, req), conn, nil
+	return res, conn, nil
 }
diff --git a/proxy_test.go b/proxy_test.go
index 244cc8b6..67db8a31 100644
--- a/proxy_test.go
+++ b/proxy_test.go
@@ -807,7 +807,8 @@ func TestIntegrationConnectDownstreamProxy(t *testing.T) {
 
 	// Set upstream proxy's downstream proxy to the host:port of the first proxy.
 	upstream.SetDownstreamProxy(&url.URL{
-		Host: dl.Addr().String(),
+		Scheme: "http",
+		Host:   dl.Addr().String(),
 	})
 
 	go upstream.Serve(ul)