准备测试环境代码

cmd/client/main.go

@@ -1,211 +1,7 @@
 package main
 
-import (
-	"bufio"
-	"encoding/binary"
-	"fmt"
-	"io"
-	"log/slog"
-	"net"
-	"os"
-	"proxy-server/pkg/utils"
-	"strconv"
-	"time"
-
-	"github.com/joho/godotenv"
-	"github.com/pkg/errors"
-)
-
-type Config struct {
-	FrpHost       string
-	FrpCtrlPort   uint16
-	FrpDataPort   uint16
-	FwdPort       uint16
-	RetryInterval int
-}
-
-var cfg Config
-
-var frpCtrlAddr string
-var frpDataAddr string
+import "proxy-server/client"
 
 func main() {
-
-	initLog()
-	initEnv()
-
-	frpCtrlAddr = net.JoinHostPort(cfg.FrpHost, strconv.Itoa(int(cfg.FrpCtrlPort)))
-	frpDataAddr = net.JoinHostPort(cfg.FrpHost, strconv.Itoa(int(cfg.FrpDataPort)))
-
-	// 建立控制通道
-	for {
-		slog.Info("建立控制通道", "addr", frpCtrlAddr)
-		err := control()
-		if err != nil {
-			slog.Error("建立控制通道失败", err)
-			slog.Info(fmt.Sprintf("%d 秒后重试", cfg.RetryInterval))
-			time.Sleep(time.Duration(cfg.RetryInterval) * time.Second)
-		}
-	}
-}
-
-func control() error {
-	conn, err := net.Dial("tcp", frpCtrlAddr)
-	if err != nil {
-		return errors.Wrap(err, "连接失败")
-	}
-	defer utils.Close(conn)
-
-	// 请求转发端口
-	slog.Info("注册转发端口", "port", cfg.FwdPort)
-	portBuf := make([]byte, 2)
-	binary.BigEndian.PutUint16(portBuf, cfg.FwdPort)
-	_, err = conn.Write(portBuf)
-	if err != nil {
-		return errors.Wrap(err, "注册转发端口失败")
-	}
-
-	// 等待用户连接
-	// 读写失败后退出重连，防止后续数据读写顺序错位导致卡死控制通道
-	for {
-		slog.Info("等待用户连接")
-		reader := bufio.NewReader(conn)
-
-		tagLen, err := utils.ReadByte(reader)
-		if err != nil {
-			return errors.Wrap(err, "接收 tagLen 失败")
-		}
-		tagBuf, err := utils.ReadBuffer(reader, int(tagLen))
-		if err != nil {
-			return errors.Wrap(err, "接收 tagBuf 失败")
-		}
-
-		// 建立数据通道
-		go func() {
-			slog.Info("收到用户连接，建立数据通道")
-			err := data(tagLen, tagBuf)
-			if err != nil {
-				slog.Error("建立数据通道失败", err)
-			}
-		}()
-	}
-}
-
-func data(tagLen byte, tagBuf []byte) error {
-	timerAll := time.Now()
-	src, err := net.Dial("tcp", frpDataAddr)
-	if err != nil {
-		return errors.Wrap(err, "连接失败")
-	}
-	defer utils.Close(src)
-
-	// 发送 tag
-	slog.Info("准备代理流量")
-	writeBuf := make([]byte, 1+len(tagBuf))
-	writeBuf[0] = tagLen
-	copy(writeBuf[1:], tagBuf)
-	_, err = src.Write(writeBuf)
-	if err != nil {
-		return errors.Wrap(err, "发送 tag 失败")
-	}
-
-	// 接收目标地址
-	slog.Info("接收目标地址")
-	addrLen, err := utils.ReadByte(src)
-	if err != nil {
-		return errors.Wrap(err, "接收 addrLen 失败")
-	}
-	addrBuf, err := utils.ReadBuffer(src, int(addrLen))
-	if err != nil {
-		return errors.Wrap(err, "接收 addrBuf 失败")
-	}
-	addr := string(addrBuf)
-
-	// 数据转发
-	slog.Info("向目标 " + addr + " 建立连接")
-	dest, err := net.Dial("tcp", addr)
-	if err != nil {
-		return errors.Wrap(err, "连接失败")
-	}
-	defer utils.Close(dest)
-
-	slog.Info("开始代理流量 " + src.RemoteAddr().String() + " <-> " + dest.RemoteAddr().String())
-	timer := time.Now()
-	errCh := make(chan error)
-	go func() {
-		written, err := io.Copy(dest, src)
-		if err != nil && !errors.Is(err, net.ErrClosed) {
-			slog.Error("上行流量代理失败", "err", err)
-			errCh <- err
-			return
-		} else {
-			slog.Info("上行流量代理结束")
-		}
-		slog.Info("上行流量", "bytes", written)
-		errCh <- nil
-	}()
-	go func() {
-		written, err := io.Copy(src, dest)
-		if err != nil && !errors.Is(err, net.ErrClosed) {
-			slog.Error("下行流量代理失败", "err", err)
-			errCh <- err
-			return
-		} else {
-			slog.Info("下行流量代理结束")
-		}
-		slog.Info("下行流量", "bytes", written)
-		errCh <- nil
-	}()
-	<-errCh
-	slog.Info("代理流量结束", "time", time.Since(timer))
-	slog.Info("数据通道结束", "time", time.Since(timerAll))
-	return nil
-}
-
-func initEnv() {
-	err := godotenv.Load()
-	if err != nil {
-		slog.Debug("没有本地环境变量文件")
-	}
-
-	cfg.FrpHost = os.Getenv("FRP_HOST")
-
-	frpCtrlPort, err := strconv.ParseUint(os.Getenv("FRP_CTRL_PORT"), 10, 16)
-	if err != nil {
-		panic("环境变量 FRP_CTRL_PORT 的值不合法 " + err.Error())
-	}
-	cfg.FrpCtrlPort = uint16(frpCtrlPort)
-	if cfg.FrpCtrlPort == 0 {
-		panic("环境变量 FRP_CTRL_PORT 不能为空")
-	}
-
-	frpDataPort, err := strconv.ParseUint(os.Getenv("FRP_DATA_PORT"), 10, 16)
-	if err != nil {
-		panic("环境变量 FRP_DATA_PORT 的值不合法 " + err.Error())
-	}
-	cfg.FrpDataPort = uint16(frpDataPort)
-	if cfg.FrpDataPort == 0 {
-		panic("环境变量 FRP_DATA_PORT 不能为空")
-	}
-
-	fwdPort, err := strconv.ParseUint(os.Getenv("FWD_PORT"), 10, 16)
-	if err != nil {
-		panic("环境变量 FWD_PORT 的值不合法 " + err.Error())
-	}
-	cfg.FwdPort = uint16(fwdPort)
-	if cfg.FwdPort == 0 {
-		panic("环境变量 FWD_PORT 不能为空")
-	}
-
-	cfg.RetryInterval, err = strconv.Atoi(os.Getenv("RETRY_INTERVAL"))
-	if err != nil {
-		panic("环境变量 RETRY_INTERVAL 的值不合法 " + err.Error())
-	}
-	if cfg.RetryInterval == 0 {
-		cfg.RetryInterval = 5
-	}
-}
-
-func initLog() {
-	slog.SetLogLoggerLevel(slog.LevelDebug)
+	client.Start()
 }