From 54242cfa8b85188d9625fe6018c8e76951f7ae03 Mon Sep 17 00:00:00 2001
From: tsosunchia <59512455+tsosunchia@users.noreply.github.com>
Date: Wed, 7 Jun 2023 03:15:29 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E6=94=AF=E6=8C=81SOCKS5/HTTP?=
 =?UTF-8?q?=E4=BB=A3=E7=90=86,=E9=80=9A=E8=BF=87=E7=8E=AF=E5=A2=83?=
 =?UTF-8?q?=E5=8F=98=E9=87=8FNEXTTRACE=5FPROXY=E9=85=8D=E7=BD=AE.=20?=
 =?UTF-8?q?=E5=A6=82export=20NEXTTRACE=5FPROXY=3Dsocks5://127.0.0.1:10808?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 go.mod               |  2 +-
 go.sum               |  4 ++--
 pow/pow.go           |  4 ++++
 tracemap/tracemap.go |  4 ++++
 util/util.go         | 14 ++++++++++++++
 wshandle/client.go   |  4 ++++
 6 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index 1896f31..71a32a1 100644
--- a/go.mod
+++ b/go.mod
@@ -9,7 +9,7 @@ require (
 	github.com/spf13/viper v1.16.0
 	github.com/stretchr/testify v1.8.4
 	github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635
-	github.com/tsosunchia/powclient v0.1.1
+	github.com/tsosunchia/powclient v0.1.3
 	golang.org/x/net v0.10.0
 	golang.org/x/sync v0.2.0
 )
diff --git a/go.sum b/go.sum
index 2d02e05..6df12e2 100644
--- a/go.sum
+++ b/go.sum
@@ -200,8 +200,8 @@ github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JT
 github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
 github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
 github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
-github.com/tsosunchia/powclient v0.1.1 h1:Llv7vFNXTvpxrd2JjnAkpGz95gjVu6A2891RQCML3Rc=
-github.com/tsosunchia/powclient v0.1.1/go.mod h1:Pm4MP3QqN74SfNskPpFIEyT+NQrcABGoXkkeRwjlMEE=
+github.com/tsosunchia/powclient v0.1.3 h1:L29HQdyZZ0Vcrn0ZzNIyrTpLOnycm3oHF46sYJ+vkhQ=
+github.com/tsosunchia/powclient v0.1.3/go.mod h1:Pm4MP3QqN74SfNskPpFIEyT+NQrcABGoXkkeRwjlMEE=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
diff --git a/pow/pow.go b/pow/pow.go
index 1e1f2b4..aab7bd5 100644
--- a/pow/pow.go
+++ b/pow/pow.go
@@ -19,6 +19,10 @@ func GetToken(fastIp string, host string, port string) (string, error) {
 	getTokenParams.SNI = host
 	getTokenParams.Host = host
 	getTokenParams.UserAgent = util.UserAgent
+	proxyUrl := util.GetProxy()
+	if proxyUrl != nil {
+		getTokenParams.Proxy = proxyUrl
+	}
 	var err error
 	// 尝试三次RetToken，如果都失败了，异常退出
 	for i := 0; i < 3; i++ {
diff --git a/tracemap/tracemap.go b/tracemap/tracemap.go
index 9a9489f..ad33ddd 100644
--- a/tracemap/tracemap.go
+++ b/tracemap/tracemap.go
@@ -35,6 +35,10 @@ func GetMapUrl(r string) (string, error) {
 			},
 		},
 	}
+	proxyUrl := util.GetProxy()
+	if proxyUrl != nil {
+		client.Transport.(*http.Transport).Proxy = http.ProxyURL(proxyUrl)
+	}
 	req, err := http.NewRequest("POST", tracemapUrl, strings.NewReader(r))
 	if err != nil {
 		return "", errors.New("an issue occurred while connecting to the tracemap API")
diff --git a/util/util.go b/util/util.go
index eaaadc1..27afdce 100644
--- a/util/util.go
+++ b/util/util.go
@@ -6,6 +6,7 @@ import (
 	"github.com/xgadget-lab/nexttrace/config"
 	"log"
 	"net"
+	"net/url"
 	"os"
 	"runtime"
 	"strings"
@@ -182,3 +183,16 @@ func GetHostAndPort() (host string, port string) {
 	}
 	return
 }
+
+func GetProxy() *url.URL {
+	proxyURLStr := GetenvDefault("NEXTTRACE_PROXY", "")
+	if proxyURLStr == "" {
+		return nil
+	}
+	proxyURL, err := url.Parse(proxyURLStr)
+	if err != nil {
+		log.Println("Failed to parse proxy URL:", err)
+		return nil
+	}
+	return proxyURL
+}
diff --git a/wshandle/client.go b/wshandle/client.go
index f82a66a..c02734b 100644
--- a/wshandle/client.go
+++ b/wshandle/client.go
@@ -148,6 +148,10 @@ func (c *WsConn) recreateWsConn() {
 	dialer.TLSClientConfig = &tls.Config{
 		ServerName: host,
 	}
+	proxyUrl := util.GetProxy()
+	if proxyUrl != nil {
+		dialer.Proxy = http.ProxyURL(proxyUrl)
+	}
 	ws, _, err := websocket.DefaultDialer.Dial(u.String(), requestHeader)
 	c.Conn = ws
 	if err != nil {