From b70d2bc9cc97540f950e0aad33595e7e2f454570 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Wed, 10 Aug 2022 17:16:03 +0800
Subject: [PATCH] fix https can use https_proxy env

---
 proxy/connection.go  | 72 ++++++++++++++++++++++++++++++++++++++++++--
 proxy/interceptor.go |  2 +-
 2 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/proxy/connection.go b/proxy/connection.go
index e808f75..cb99e9d 100644
--- a/proxy/connection.go
+++ b/proxy/connection.go
@@ -1,11 +1,16 @@
 package proxy
 
 import (
+	"bufio"
 	"context"
 	"crypto/tls"
 	"encoding/json"
+	"errors"
 	"net"
 	"net/http"
+	"net/url"
+	"strings"
+	"time"
 
 	uuid "github.com/satori/go.uuid"
 	log "github.com/sirupsen/logrus"
@@ -138,13 +143,24 @@ func (connCtx *ConnContext) initHttpServerConn() {
 	connCtx.ServerConn = serverConn
 }
 
-func (connCtx *ConnContext) initServerTcpConn() error {
+func (connCtx *ConnContext) initServerTcpConn(req *http.Request) error {
 	log.Debugln("in initServerTcpConn")
 	ServerConn := newServerConn()
 	connCtx.ServerConn = ServerConn
 	ServerConn.Address = connCtx.pipeConn.host
 
-	plainConn, err := (&net.Dialer{}).DialContext(context.Background(), "tcp", ServerConn.Address)
+	// test is use proxy
+	clientReq := &http.Request{URL: &url.URL{Scheme: "https", Host: ServerConn.Address}}
+	proxyUrl, err := http.ProxyFromEnvironment(clientReq)
+	if err != nil {
+		return err
+	}
+	var plainConn net.Conn
+	if proxyUrl != nil {
+		plainConn, err = getProxyConn(proxyUrl, ServerConn.Address)
+	} else {
+		plainConn, err = (&net.Dialer{}).DialContext(context.Background(), "tcp", ServerConn.Address)
+	}
 	if err != nil {
 		return err
 	}
@@ -161,13 +177,63 @@ func (connCtx *ConnContext) initServerTcpConn() error {
 	return nil
 }
 
+// connect proxy when set https_proxy env
+// ref: http/transport.go dialConn func
+func getProxyConn(proxyUrl *url.URL, address string) (net.Conn, error) {
+	conn, err := (&net.Dialer{}).DialContext(context.Background(), "tcp", proxyUrl.Host)
+	if err != nil {
+		return nil, err
+	}
+	connectReq := &http.Request{
+		Method: "CONNECT",
+		URL:    &url.URL{Opaque: address},
+		Host:   address,
+	}
+	connectCtx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
+	defer cancel()
+	didReadResponse := make(chan struct{}) // closed after CONNECT write+read is done or fails
+	var resp *http.Response
+	// Write the CONNECT request & read the response.
+	go func() {
+		defer close(didReadResponse)
+		err = connectReq.Write(conn)
+		if err != nil {
+			return
+		}
+		// Okay to use and discard buffered reader here, because
+		// TLS server will not speak until spoken to.
+		br := bufio.NewReader(conn)
+		resp, err = http.ReadResponse(br, connectReq)
+	}()
+	select {
+	case <-connectCtx.Done():
+		conn.Close()
+		<-didReadResponse
+		return nil, connectCtx.Err()
+	case <-didReadResponse:
+		// resp or err now set
+	}
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+	if resp.StatusCode != 200 {
+		_, text, ok := strings.Cut(resp.Status, " ")
+		conn.Close()
+		if !ok {
+			return nil, errors.New("unknown status code")
+		}
+		return nil, errors.New(text)
+	}
+	return conn, nil
+}
+
 func (connCtx *ConnContext) initHttpsServerConn() {
 	if !connCtx.ClientConn.Tls {
 		return
 	}
 	connCtx.ServerConn.client = &http.Client{
 		Transport: &http.Transport{
-			Proxy: http.ProxyFromEnvironment,
 			DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
 				<-connCtx.ServerConn.tlsHandshaked
 				return connCtx.ServerConn.tlsConn, connCtx.ServerConn.tlsHandshakeErr
diff --git a/proxy/interceptor.go b/proxy/interceptor.go
index 60e5130..e2fe0c8 100644
--- a/proxy/interceptor.go
+++ b/proxy/interceptor.go
@@ -140,7 +140,7 @@ func (m *middle) close() error {
 
 func (m *middle) dial(req *http.Request) (net.Conn, error) {
 	pipeClientConn, pipeServerConn := newPipes(req)
-	err := pipeServerConn.connContext.initServerTcpConn()
+	err := pipeServerConn.connContext.initServerTcpConn(req)
 	if err != nil {
 		pipeClientConn.Close()
 		pipeServerConn.Close()