From f220134e5b4164ec807760b59ee6120b345d0a40 Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Sun, 14 Jul 2024 19:41:13 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20HTTP=20=E5=86=85=E7=BD=91=E7=A9=BF?=
 =?UTF-8?q?=E9=80=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/agent/main.go   | 67 ++++++++++++++++++++++++++++++++++++++++++---
 model/task.go       |  6 ++++
 pkg/monitor/myip.go |  4 +--
 3 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index 0bdb785..e719064 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -52,7 +52,7 @@ type AgentCliParam struct {
 	Version               bool   // 当前版本号
 	IPReportPeriod        uint32 // 上报IP间隔
 	UseIPv6CountryCode    bool   // 默认优先展示IPv6旗帜
-	UseGiteeToUpgrade      bool   // 强制从Gitee获取更新
+	UseGiteeToUpgrade     bool   // 强制从Gitee获取更新
 }
 
 var (
@@ -222,7 +222,7 @@ func run() {
 	// 上报服务器信息
 	go reportState()
 	// 更新IP信息
-	go monitor.UpdateIP(agentConfig.Debug, agentCliParam.UseIPv6CountryCode, agentCliParam.IPReportPeriod)
+	go monitor.UpdateIP(agentCliParam.UseIPv6CountryCode, agentCliParam.IPReportPeriod)
 
 	// 定时检查更新
 	if _, err := semver.Parse(version); err == nil && !agentCliParam.DisableAutoUpdate {
@@ -388,8 +388,10 @@ func doTask(task *pb.Task) {
 	case model.TaskTypeTerminalGRPC:
 		handleTerminalTask(task)
 		return
-	case model.TaskTypeKeepalive:
+	case model.TaskTypeNAT:
+		handleNATTask(task)
 		return
+	case model.TaskTypeKeepalive:
 	default:
 		println("不支持的任务：", task)
 		return
@@ -668,7 +670,7 @@ func handleTerminalTask(task *pb.Task) {
 
 	go func() {
 		for {
-			buf := make([]byte, 1024)
+			buf := make([]byte, 10240)
 			read, err := tty.Read(buf)
 			if err != nil {
 				remoteIO.Send(&pb.IOStreamData{Data: []byte(err.Error())})
@@ -702,6 +704,63 @@ func handleTerminalTask(task *pb.Task) {
 	}
 }
 
+func handleNATTask(task *pb.Task) {
+	var nat model.TaskNAT
+	err := util.Json.Unmarshal([]byte(task.GetData()), &nat)
+	if err != nil {
+		println("NAT 任务解析错误：", err)
+		return
+	}
+
+	remoteIO, err := client.IOStream(context.Background())
+	if err != nil {
+		println("NAT IOStream失败：", err)
+		return
+	}
+
+	// 发送 StreamID
+	if err := remoteIO.Send(&pb.IOStreamData{Data: append([]byte{
+		0xff, 0x05, 0xff, 0x05,
+	}, []byte(nat.StreamID)...)}); err != nil {
+		println("NAT 发送StreamID失败：", err)
+		return
+	}
+
+	conn, err := net.Dial("tcp", nat.Host)
+	if err != nil {
+		println(fmt.Sprintf("NAT Dial %s 失败：%s", nat.Host, err))
+		return
+	}
+
+	defer func() {
+		err := conn.Close()
+		errCloseSend := remoteIO.CloseSend()
+		println("NAT exit", nat.StreamID, err, errCloseSend)
+	}()
+	println("NAT init", nat.StreamID)
+
+	go func() {
+		buf := make([]byte, 10240)
+		for {
+			read, err := conn.Read(buf)
+			if err != nil {
+				remoteIO.Send(&pb.IOStreamData{Data: []byte(err.Error())})
+				remoteIO.CloseSend()
+				return
+			}
+			remoteIO.Send(&pb.IOStreamData{Data: buf[:read]})
+		}
+	}()
+
+	for {
+		var remoteData *pb.IOStreamData
+		if remoteData, err = remoteIO.Recv(); err != nil {
+			return
+		}
+		conn.Write(remoteData.Data)
+	}
+}
+
 func println(v ...interface{}) {
 	util.Println(agentConfig.Debug, v...)
 }
diff --git a/model/task.go b/model/task.go
index 6c8e6a2..c81922f 100644
--- a/model/task.go
+++ b/model/task.go
@@ -10,8 +10,14 @@ const (
 	TaskTypeUpgrade
 	TaskTypeKeepalive
 	TaskTypeTerminalGRPC
+	TaskTypeNAT
 )
 
 type TerminalTask struct {
 	StreamID string
 }
+
+type TaskNAT struct {
+	StreamID string
+	Host     string
+}
diff --git a/pkg/monitor/myip.go b/pkg/monitor/myip.go
index 180ac69..6094be5 100644
--- a/pkg/monitor/myip.go
+++ b/pkg/monitor/myip.go
@@ -55,9 +55,9 @@ var (
 )
 
 // UpdateIP 按设置时间间隔更新IP地址与国家码的缓存
-func UpdateIP(logging bool, useIPv6CountryCode bool, period uint32) {
+func UpdateIP(useIPv6CountryCode bool, period uint32) {
 	for {
-		util.Println(logging, "NEZHA_AGENT>> 正在更新本地缓存IP信息")
+		util.Println(agentConfig.Debug, "正在更新本地缓存IP信息")
 		var primaryIP, secondaryIP geoIP
 		if useIPv6CountryCode {
 			primaryIP = fetchGeoIP(geoIPApiList, true)