proxy/cmd/client/main.go

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

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)
}