重新规划网关与节点的交互协议，实现统一命令位的识别和处理

client/client.go

@@ -2,20 +2,22 @@ package client
 
 import (
 	"bufio"
+	"context"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
 	"net"
+	_ "net/http/pprof"
+	"os"
+	"os/signal"
 	"proxy-server/client/core"
 	"proxy-server/client/env"
 	"proxy-server/client/geo"
 	"proxy-server/client/report"
 	"proxy-server/pkg/utils"
 	"time"
-
-	"errors"
-	_ "net/http/pprof"
 )
 
 func Start() error {
@@ -31,7 +33,7 @@ func Start() error {
 	slog.Debug("获取节点归属地...")
 	err = geo.Query()
 	if err != nil {
-		slog.Error("获取归属地失败", "err", err)
+		return fmt.Errorf("获取节点归属地失败: %w", err)
 	}
 
 	// 注册节点
@@ -41,27 +43,38 @@ func Start() error {
 		return fmt.Errorf("注册节点失败: %w", err)
 	}
 
-	// 性能监控
-	// go func() {
-	// 	runtime.SetBlockProfileRate(1)
-	// 	err := http.ListenAndServe(":7070", nil)
-	// 	if err != nil {
-	// 		slog.Error("性能监控服务启动失败", "err", err)
-	// 	}
-	// }()
-
 	// 建立控制通道
-	for {
-		err := ctrl(id, host)
-		if err != nil {
+	var ctx, cancel = signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
+	defer cancel()
+
+	go func() {
+		for {
+			err = ctrl(ctx, id, host)
+			if err == nil {
+				return
+			}
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
 			slog.Error("建立控制通道失败", "err", err)
 			slog.Info(fmt.Sprintf("%d 秒后重试", core.RetryInterval))
 			time.Sleep(time.Duration(core.RetryInterval) * time.Second)
 		}
+	}()
+
+	// 下线节点
+	slog.Debug("下线节点...")
+	err = report.Offline()
+	if err != nil {
+		slog.Error("下线节点失败", "err", err)
 	}
+
+	return ctx.Err()
 }
 
-func ctrl(id int32, host string) error {
+func ctrl(ctx context.Context, id int32, host string) error {
 	ctrlAddr := net.JoinHostPort(host, fmt.Sprintf("%d", core.FwdCtrlPort))
 	dataAddr := net.JoinHostPort(host, fmt.Sprintf("%d", core.FwdDataPort))
 
@@ -71,103 +84,94 @@ func ctrl(id int32, host string) error {
 		return errors.New("连接失败")
 	}
 	defer utils.Close(conn)
+	var reader = bufio.NewReader(conn)
 
-	// 发送客户端信息
-	var buf = make([]byte, 4)
-	_, err = binary.Encode(buf, binary.BigEndian, id)
+	// 发送节点连接命令
+	err = sendOpen(reader, conn, id)
 	if err != nil {
-		return fmt.Errorf("编码客户端 ID 失败: %w", err)
-	}
-	_, err = conn.Write(buf)
-	if err != nil {
-		return fmt.Errorf("发送客户端 ID 失败: %w", err)
+		return fmt.Errorf("发送节点信息失败: %w", err)
 	}
 
-	// 等待服务端响应
-	reader := bufio.NewReader(conn)
-	respBuf, err := reader.ReadByte()
-	if err != nil {
-		return errors.New("接收响应失败")
-	}
-	if respBuf != 1 {
-		return errors.New("服务端响应失败")
-	} else {
-		slog.Info("成功建立连接")
-	}
+	// 异步定时发送心跳
+	go func() {
+		ticker := time.NewTicker(time.Duration(core.HeartbeatInterval) * time.Second)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case tick := <-ticker.C:
+				err := sendPing(reader, conn)
+				if err != nil {
+					slog.Error("发送心跳失败", "time", tick, "err", err)
+				}
+			}
+		}
+	}()
 
 	// 等待用户连接
 	// 读写失败后退出重连，防止后续数据读写顺序错位导致卡死控制通道
 	slog.Info("等待用户连接")
-	for {
-
-		// 接收 dst
-		dstLen, err := reader.ReadByte()
-		if err != nil {
-			return errors.New("接收 dstLen 失败")
-		}
-		dstBuf, err := utils.ReadBuffer(reader, int(dstLen))
-		if err != nil {
-			return errors.New("接收 dstBuf 失败")
-		}
-		addr := string(dstBuf)
-
-		// 接收 tag
-		tagLen, err := reader.ReadByte()
-		if err != nil {
-			return errors.New("接收 tagLen 失败")
-		}
-		tagBuf, err := utils.ReadBuffer(reader, int(tagLen))
-		if err != nil {
-			return errors.New("接收 tagBuf 失败")
-		}
-
-		// 建立数据通道
-		go func() {
-			err := data(dataAddr, addr, tagBuf)
+	for loop := true; loop; {
+		select {
+		case <-ctx.Done():
+			loop = false
+		default:
+			// 接收 dst
+			tag, addr, err := onConn(reader)
 			if err != nil {
-				slog.Error("建立数据通道失败", "err", err)
+				return fmt.Errorf("接收连接命令失败: %w", err)
 			}
-		}()
+
+			// 建立数据通道
+			go func() {
+				err := data(dataAddr, addr, tag)
+				if err != nil {
+					slog.Error("建立数据通道失败", "err", err)
+				}
+			}()
+		}
 	}
+
+	// 发送关闭连接（不 return err，否则会重新连接）
+	err = sendClose(reader, conn)
+	if err != nil {
+		slog.Error("发送关闭连接失败", "err", err)
+	}
+
+	return nil
 }
 
-func data(dataAddr string, dest string, tag []byte) error {
+func data(proxy string, dest string, tag [16]byte) error {
+
+	// 向目标地址建立连接
+	var result = 1
+	var dstErr error
+	dst, err := net.Dial("tcp", dest)
+	if err != nil {
+		dstErr = fmt.Errorf("连接目标地址失败: %w", dstErr)
+		result = 0
+	}
+	defer utils.Close(dst)
 
 	// 向服务端建立连接
-	src, err := net.Dial("tcp", dataAddr)
+	src, err := net.Dial("tcp", proxy)
 	if err != nil {
 		return errors.New("连接服务端失败")
 	}
-
-	tagLen := byte(len(tag))
-	tagBuf := make([]byte, 2+tagLen)
-	tagBuf[1] = tagLen
-	copy(tagBuf[2:], tag)
-
-	// 向目标地址建立连接
-	dst, dstErr := net.Dial("tcp", dest)
-	if dstErr != nil {
-		tagBuf[0] = 0
-	} else {
-		tagBuf[0] = 1
-	}
+	defer utils.Close(src)
 
 	// 发送连接状态
-	_, err = src.Write(tagBuf)
+	var buf = make([]byte, 17)
+	copy(buf[0:16], tag[:])
+	buf[16] = byte(result)
+	_, err = src.Write(buf)
 	if err != nil {
-		utils.Close(src)
-		if dst != nil {
-			utils.Close(dst)
-		}
 		return errors.New("发送连接状态失败")
 	}
 
-	if tagBuf[0] == 0 {
-		utils.Close(src)
-		if dst != nil {
-			utils.Close(dst)
-		}
-		return errors.New("连接目标地址失败")
+	if result == 0 {
+		return dstErr
 	}
 
 	go func() {
@@ -186,3 +190,91 @@ func data(dataAddr string, dest string, tag []byte) error {
 	}()
 	return nil
 }
+
+func sendOpen(reader io.Reader, writer io.Writer, id int32) error {
+
+	// 发送打开连接
+	var buf = make([]byte, 5)
+	buf[0] = 3
+	binary.BigEndian.PutUint32(buf[1:], uint32(id))
+
+	_, err := writer.Write(buf)
+	if err != nil {
+		return fmt.Errorf("发送打开连接失败: %w", err)
+	}
+
+	// 等待服务端响应
+	respBuf := make([]byte, 1)
+	_, err = io.ReadFull(reader, respBuf)
+	if err != nil {
+		return fmt.Errorf("接收服务端响应失败: %w", err)
+	}
+	if respBuf[0] != 1 {
+		return errors.New("服务端响应失败")
+	}
+
+	return nil
+}
+
+func sendClose(reader io.Reader, writer io.Writer) error {
+	// 发送关闭连接
+	_, err := writer.Write([]byte{4})
+	if err != nil {
+		return err
+	}
+
+	// 等待服务端响应
+	respBuf := make([]byte, 1)
+	_, err = io.ReadFull(reader, respBuf)
+	if err != nil {
+		return fmt.Errorf("接收服务端响应失败: %w", err)
+	}
+	if respBuf[0] != 1 {
+		return errors.New("服务端响应失败")
+	}
+
+	return nil
+}
+
+func sendPing(reader io.Reader, writer io.Writer) error {
+	_, err := writer.Write([]byte{2})
+	if err != nil {
+		return err
+	}
+
+	// 等待服务端响应
+	respBuf := make([]byte, 1)
+	_, err = io.ReadFull(reader, respBuf)
+	if err != nil {
+		return fmt.Errorf("接收服务端响应失败: %w", err)
+	}
+	if respBuf[0] != 1 {
+		return errors.New("服务端响应失败")
+	}
+
+	return nil
+}
+
+func onConn(reader io.Reader) (tag [16]byte, addr string, err error) {
+	var buf = make([]byte, 1+16+2)
+	_, err = io.ReadFull(reader, buf)
+	if err != nil {
+		return [16]byte{}, "", err
+	}
+
+	if buf[0] != 5 {
+		return [16]byte{}, "", errors.New("命令错误")
+	}
+
+	tag = [16]byte(buf[1:17])
+
+	var addrLen = binary.BigEndian.Uint16(buf[17:19])
+	var addrBuf = make([]byte, addrLen)
+	_, err = io.ReadFull(reader, addrBuf)
+	if err != nil {
+		return [16]byte{}, "", err
+	}
+
+	addr = string(addrBuf)
+	return tag, addr, nil
+}

client/core/consts.go

@@ -6,3 +6,4 @@ const FwdCtrlPort uint = 18080
 const FwdDataPort uint = 18081
 
 const RetryInterval uint = 5
+const HeartbeatInterval uint = 30

client/report/report.go

@@ -67,3 +67,7 @@ func Online(prov, city, isp string) (id int32, host string, err error) {
 
 	return respBody.Id, respBody.Host, nil
 }
+
+func Offline() error {
+	return nil
+}