lanhu/proxy
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化连接效率

Browse Source
This commit is contained in:
luorijun
2025-03-07 15:50:57 +08:00
parent 791f20d2d7
commit 053041ae34
13 changed files with 564 additions and 418 deletions
Download Patch File Download Diff File Expand all files Collapse all files

234
client/client.go
View File

@@ -2,26 +2,32 @@ package client
import (
"bufio"
"encoding/binary"
"flag"
"fmt"
"io"
"log/slog"
"net"
"net/http"
"os"
"proxy-server/pkg/utils"
"runtime"
"strconv"
"time"
"github.com/joho/godotenv"
"github.com/pkg/errors"
_ "net/http/pprof"
)
const Version byte = 1
type Config struct {
FrpHost string
FrpCtrlPort uint16
FrpDataPort uint16
FwdPort uint16
RetryInterval int
Name string
FwdHost string
FwdCtrlPort uint
FwdDataPort uint
RetryInterval uint
}
var cfg Config
@@ -32,15 +38,24 @@ var frpDataAddr string
func Start() {
initLog()
initEnv()
initCmd()
initDevEnv()
frpCtrlAddr = net.JoinHostPort(cfg.FrpHost, strconv.Itoa(int(cfg.FrpCtrlPort)))
frpDataAddr = net.JoinHostPort(cfg.FrpHost, strconv.Itoa(int(cfg.FrpDataPort)))
frpCtrlAddr = net.JoinHostPort(cfg.FwdHost, strconv.Itoa(int(cfg.FwdCtrlPort)))
frpDataAddr = net.JoinHostPort(cfg.FwdHost, strconv.Itoa(int(cfg.FwdDataPort)))
// 性能监控
go func() {
runtime.SetBlockProfileRate(1)
err := http.ListenAndServe(":6060", nil)
if err != nil {
slog.Error("性能监控服务启动失败", "err", err)
}
}()
// 建立控制通道
for {
slog.Info("建立控制通道", "addr", frpCtrlAddr)
err := control()
err := ctrl()
if err != nil {
slog.Error("建立控制通道失败", err)
slog.Info(fmt.Sprintf("%d 秒后重试", cfg.RetryInterval))
@@ -49,7 +64,9 @@ func Start() {
}
}
func control() error {
func ctrl() error {
slog.Info("建立控制通道", "addr", frpCtrlAddr)
conn, err := net.Dial("tcp", frpCtrlAddr)
if err != nil {
return errors.Wrap(err, "连接失败")
@@ -59,19 +76,50 @@ func control() error {
reader := bufio.NewReader(conn)
// 请求转发端口
slog.Info("注册转发端口", "port", cfg.FwdPort)
portBuf := make([]byte, 2)
binary.BigEndian.PutUint16(portBuf, cfg.FwdPort)
_, err = conn.Write(portBuf)
_, err = conn.Write([]byte{Version})
if err != nil {
return errors.Wrap(err, "注册转发端口失败")
return errors.Wrap(err, "发送版本号失败")
}
// 发送客户端名称
nameLen := byte(len(cfg.Name))
nameBuf := make([]byte, 1+nameLen)
nameBuf[0] = nameLen
copy(nameBuf[1:], cfg.Name)
_, err = conn.Write(nameBuf)
if err != nil {
return errors.Wrap(err, "发送 name 失败")
}
// 等待服务端响应
respBuf, err := reader.ReadByte()
if err != nil {
return errors.Wrap(err, "接收响应失败")
}
if respBuf != 1 {
return errors.New("服务端响应失败")
} else {
slog.Info("成功建立连接")
}
// 等待用户连接
// 读写失败后退出重连，防止后续数据读写顺序错位导致卡死控制通道
slog.Info("等待用户连接")
for {
tagLen, err := utils.ReadByte(reader)
// 接收 dst
dstLen, err := reader.ReadByte()
if err != nil {
return errors.Wrap(err, "接收 dstLen 失败")
}
dstBuf, err := utils.ReadBuffer(reader, int(dstLen))
if err != nil {
return errors.Wrap(err, "接收 dstBuf 失败")
}
addr := string(dstBuf)
// 接收 tag
tagLen, err := reader.ReadByte()
if err != nil {
return errors.Wrap(err, "接收 tagLen 失败")
}
@@ -81,9 +129,8 @@ func control() error {
}
// 建立数据通道
slog.Info("收到用户连接，建立数据通道", "tag", string(tagBuf))
go func() {
err := data(tagLen, tagBuf)
err := data(addr, tagBuf)
if err != nil {
slog.Error("建立数据通道失败", err)
}
@@ -91,121 +138,86 @@ func control() error {
}
}
func data(tagLen byte, tagBuf []byte) error {
timerAll := time.Now()
func data(addr string, tag []byte) error {
// 向服务端建立连接
src, err := net.Dial("tcp", frpDataAddr)
if err != nil {
return errors.Wrap(err, "连接失败")
}
defer utils.Close(src)
// 发送 tag
slog.Info("准备代理流量")
writeBuf := make([]byte, 1+tagLen)
writeBuf[0] = tagLen
copy(writeBuf[1:], tagBuf)
_, err = src.Write(writeBuf)
if err != nil {
return errors.Wrap(err, "发送 tag 失败")
return errors.Wrap(err, "连接服务端失败")
}
// 接收目标地址
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)
tagLen := byte(len(tag))
tagBuf := make([]byte, 2+tagLen)
tagBuf[1] = tagLen
copy(tagBuf[2:], tag)
// 数据转发
slog.Info("向目标 " + addr + " 建立连接")
dest, err := net.Dial("tcp", addr)
// 向目标地址建立连接
dst, err := net.Dial("tcp", addr)
if err != nil {
return errors.Wrap(err, "连接失败")
tagBuf[0] = 0
} else {
tagBuf[0] = 1
}
// 发送连接状态
_, err = src.Write(tagBuf)
if err != nil {
utils.Close(src)
if dst != nil {
utils.Close(dst)
}
return errors.Wrap(err, "发送连接状态失败")
}
if tagBuf[0] == 0 {
utils.Close(src)
if dst != nil {
utils.Close(dst)
}
return errors.New("目标地址连接失败")
}
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)
defer utils.Close(dst)
_, err := io.Copy(dst, src)
if err != nil && !errors.Is(err, net.ErrClosed) {
slog.Error("上行流量代理失败", "err", err)
errCh <- err
return
} else {
slog.Debug("上行流量代理结束")
}
slog.Debug("上行流量", "bytes", written)
errCh <- nil
}()
go func() {
written, err := io.Copy(src, dest)
defer utils.Close(src)
_, err := io.Copy(src, dst)
if err != nil && !errors.Is(err, net.ErrClosed) {
slog.Error("下行流量代理失败", "err", err)
errCh <- err
return
} else {
slog.Debug("下行流量代理结束")
}
slog.Debug("下行流量", "bytes", written)
errCh <- nil
}()
<-errCh
slog.Debug("代理流量结束", "time", time.Since(timer))
slog.Debug("数据通道结束", "time", time.Since(timerAll))
return nil
}
func initEnv() {
func initLog() {
slog.SetLogLoggerLevel(slog.LevelDebug)
}
func initCmd() {
flag.StringVar(&cfg.Name, "n", "", "客户端名称")
flag.StringVar(&cfg.FwdHost, "h", "", "转发服务器地址")
flag.UintVar(&cfg.FwdCtrlPort, "c", 18080, "转发服务器控制通道端口")
flag.UintVar(&cfg.FwdDataPort, "d", 18081, "转发服务器数据通道端口")
flag.UintVar(&cfg.RetryInterval, "r", 5, "重试间隔时间")
flag.Parse()
if cfg.Name == "" {
slog.Error("客户端名称不能为空")
flag.Usage()
os.Exit(1)
}
}
func initDevEnv() {
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.LevelInfo)
cfg.FwdHost = os.Getenv("FWD_HOST")
}