使用自定义流程实现网关协议解析转发功能

cmd/client/main.go

@@ -0,0 +1,156 @@
+package main
+
+import (
+	"bufio"
+	"encoding/binary"
+	"github.com/joho/godotenv"
+	"io"
+	"log/slog"
+	"net"
+	"os"
+	"proxy-server/pkg/utils"
+	"strconv"
+)
+
+func main() {
+	slog.SetLogLoggerLevel(slog.LevelDebug)
+
+	// 初始化环境变量
+	err := godotenv.Load()
+	if err != nil {
+		slog.Debug("没有本地环境变量文件")
+	}
+
+	// 建立控制连接
+	for {
+		slog.Info("与服务端建立控制连接")
+		frpHost := os.Getenv("FRP_SERVER")
+		frpPort := os.Getenv("FRP_PORT")
+		frpAddr := net.JoinHostPort(frpHost, frpPort)
+		slog.Info("frpAddr", frpAddr)
+		conn, err := net.Dial("tcp", frpAddr)
+		if err != nil {
+			slog.Error("dial error", err)
+			return
+		}
+
+		reader := bufio.NewReader(conn)
+
+		// 请求转发端口
+		slog.Info("请求转发端口")
+		fwdPortStr := os.Getenv("FWD_PORT")
+		fwdPort, err := strconv.ParseUint(fwdPortStr, 10, 16)
+		if err != nil {
+			slog.Error("parse error", err)
+			return
+		}
+		portBuf := make([]byte, 2)
+		binary.BigEndian.PutUint16(portBuf, uint16(fwdPort))
+		_, err = conn.Write(portBuf)
+		if err != nil {
+			slog.Error("write error", err)
+			return
+		}
+
+		// 读取目标地址
+		for {
+			slog.Info("等待建立数据通道命令")
+
+			tagLen, err := reader.ReadByte()
+			if err != nil {
+				slog.Error("read error", err)
+				return
+			}
+
+			tagBuf := make([]byte, tagLen)
+			_, err = io.ReadFull(conn, tagBuf)
+			if err != nil {
+				slog.Error("read error", err)
+				return
+			}
+
+			// 建立数据通道
+			go dataTun(tagLen, tagBuf)
+		}
+	}
+}
+
+func dataTun(tagLen byte, tagBuff []byte) {
+
+	slog.Info("建立数据通道")
+	src, err := net.Dial("tcp", "localhost:18081")
+	if err != nil {
+		slog.Error("建立数据通道失败", err)
+		return
+	}
+	defer func() {
+		err := src.Close()
+		if err != nil {
+			slog.Error("close error", err)
+		}
+	}()
+
+	// 发送 tag
+	_, err = src.Write([]byte{tagLen})
+	if err != nil {
+		slog.Error("write error", err)
+		return
+	}
+	_, err = src.Write(tagBuff)
+	if err != nil {
+		slog.Error("write error", err)
+		return
+	}
+
+	// 接收目标地址
+	slog.Info("等待目标地址")
+
+	addrLen, err := utils.ReadByte(src)
+	if err != nil {
+		slog.Error("接收目标地址失败", err)
+		return
+	}
+	addrBuf, err := utils.ReadBuffer(src, int(addrLen))
+	if err != nil {
+		slog.Error("接收目标地址失败", err)
+		return
+	}
+	addr := string(addrBuf)
+
+	// 数据转发
+	slog.Info("向 " + addr + " 建立连接")
+	dest, err := net.Dial("tcp", addr)
+	if err != nil {
+		slog.Error("与目标地址连接建立失败", err)
+		return
+	}
+	defer func() {
+		err = dest.Close()
+		if err != nil {
+			slog.Error("close error", err)
+		}
+	}()
+
+	slog.Info("开始数据转发 " + src.RemoteAddr().String() + " <-> " + dest.RemoteAddr().String())
+
+	errCh := make(chan error, 2)
+	go func() {
+		_, err := io.Copy(src, dest)
+		if err != nil {
+			slog.Error("copy error f2t", err)
+			errCh <- err
+			return
+		}
+		errCh <- nil
+	}()
+	go func() {
+		_, err := io.Copy(dest, src)
+		if err != nil {
+			slog.Error("copy error t2f", err)
+			errCh <- err
+			return
+		}
+		errCh <- nil
+	}()
+	<-errCh
+}