新增代理服务与边缘节点注册功能

client/client.go

@@ -2,99 +2,112 @@ package client
 
 import (
 	"bufio"
-	"flag"
 	"fmt"
 	"io"
 	"log/slog"
 	"net"
-	"net/http"
 	"os"
+	"proxy-server/client/core"
+	"proxy-server/client/geo"
+	"proxy-server/client/report"
 	"proxy-server/pkg/utils"
-	"runtime"
-	"strconv"
 	"time"
 
+	"errors"
 	"github.com/joho/godotenv"
-	"github.com/pkg/errors"
 
 	_ "net/http/pprof"
 )
 
-const Version byte = 1
+var Geo geo.Func = geo.Ipapi
 
-type Config struct {
-	Name          string
-	FwdHost       string
-	FwdCtrlPort   uint
-	FwdDataPort   uint
-	RetryInterval uint
-}
+func Start() error {
 
-var cfg Config
+	// 初始化环境变量
+	slog.SetLogLoggerLevel(slog.LevelDebug)
 
-var frpCtrlAddr string
-var frpDataAddr string
-
-func Start() {
-
-	initLog()
-	initCmd()
-	initDevEnv()
-
-	frpCtrlAddr = net.JoinHostPort(cfg.FwdHost, strconv.Itoa(int(cfg.FwdCtrlPort)))
-	frpDataAddr = net.JoinHostPort(cfg.FwdHost, strconv.Itoa(int(cfg.FwdDataPort)))
+	err := godotenv.Load()
+	if err != nil {
+		slog.Debug("没有本地环境变量文件")
+	} else {
+		online := os.Getenv("ENDPOINT_ONLINE")
+		if online != "" {
+			core.EndpointOnline = online
+		}
+		offline := os.Getenv("ENDPOINT_OFFLINE")
+		if offline != "" {
+			core.EndpointOffline = offline
+		}
+	}
 
 	// 性能监控
-	go func() {
-		runtime.SetBlockProfileRate(1)
-		err := http.ListenAndServe(":7070", nil)
-		if err != nil {
-			slog.Error("性能监控服务启动失败", "err", err)
-		}
-	}()
+	// go func() {
+	// 	runtime.SetBlockProfileRate(1)
+	// 	err := http.ListenAndServe(":7070", nil)
+	// 	if err != nil {
+	// 		slog.Error("性能监控服务启动失败", "err", err)
+	// 	}
+	// }()
+
+	// 获取归属地
+	slog.Debug("获取节点归属地...")
+	prov, city, isp, err := Geo()
+	if err != nil {
+		slog.Error("获取归属地失败", "err", err)
+	}
+
+	// 注册节点
+	slog.Debug("注册节点...")
+	host, err := report.Online(prov, city, isp)
+	if err != nil {
+		slog.Error("节点注册失败", "err", err)
+		return err
+	}
 
 	// 建立控制通道
 	for {
-		err := ctrl()
+		err := ctrl(host)
 		if err != nil {
-			slog.Error("建立控制通道失败", err)
-			slog.Info(fmt.Sprintf("%d 秒后重试", cfg.RetryInterval))
-			time.Sleep(time.Duration(cfg.RetryInterval) * time.Second)
+			slog.Error("建立控制通道失败", "err", err)
+			slog.Info(fmt.Sprintf("%d 秒后重试", core.RetryInterval))
+			time.Sleep(time.Duration(core.RetryInterval) * time.Second)
 		}
 	}
 }
 
-func ctrl() error {
-	slog.Info("建立控制通道", "addr", frpCtrlAddr)
+func ctrl(host string) error {
+	ctrlAddr := net.JoinHostPort(host, fmt.Sprintf("%d", core.FwdCtrlPort))
+	dataAddr := net.JoinHostPort(host, fmt.Sprintf("%d", core.FwdDataPort))
 
-	conn, err := net.Dial("tcp", frpCtrlAddr)
+	slog.Info("建立控制通道", "addr", ctrlAddr)
+	conn, err := net.Dial("tcp", ctrlAddr)
 	if err != nil {
-		return errors.Wrap(err, "连接失败")
+		return errors.New("连接失败")
 	}
 	defer utils.Close(conn)
 
 	reader := bufio.NewReader(conn)
 
 	// 请求转发端口
-	_, err = conn.Write([]byte{Version})
+	_, err = conn.Write([]byte{core.Version})
 	if err != nil {
-		return errors.Wrap(err, "发送版本号失败")
+		return errors.New("发送版本号失败")
 	}
 
 	// 发送客户端名称
-	nameLen := byte(len(cfg.Name))
+	nameLen := byte(len(core.Name))
 	nameBuf := make([]byte, 1+nameLen)
 	nameBuf[0] = nameLen
-	copy(nameBuf[1:], cfg.Name)
+	copy(nameBuf[1:], core.Name)
 	_, err = conn.Write(nameBuf)
 	if err != nil {
-		return errors.Wrap(err, "发送 name 失败")
+		return errors.New("发送 name 失败")
 	}
 
 	// 等待服务端响应
 	respBuf, err := reader.ReadByte()
 	if err != nil {
-		return errors.Wrap(err, "接收响应失败")
+		return errors.New("接收响应失败")
 	}
 	if respBuf != 1 {
 		return errors.New("服务端响应失败")
@@ -110,40 +123,40 @@ func ctrl() error {
 		// 接收 dst
 		dstLen, err := reader.ReadByte()
 		if err != nil {
-			return errors.Wrap(err, "接收 dstLen 失败")
+			return errors.New("接收 dstLen 失败")
 		}
 		dstBuf, err := utils.ReadBuffer(reader, int(dstLen))
 		if err != nil {
-			return errors.Wrap(err, "接收 dstBuf 失败")
+			return errors.New("接收 dstBuf 失败")
 		}
 		addr := string(dstBuf)
 
 		// 接收 tag
 		tagLen, err := reader.ReadByte()
 		if err != nil {
-			return errors.Wrap(err, "接收 tagLen 失败")
+			return errors.New("接收 tagLen 失败")
 		}
 		tagBuf, err := utils.ReadBuffer(reader, int(tagLen))
 		if err != nil {
-			return errors.Wrap(err, "接收 tagBuf 失败")
+			return errors.New("接收 tagBuf 失败")
 		}
 
 		// 建立数据通道
 		go func() {
-			err := data(addr, tagBuf)
+			err := data(dataAddr, addr, tagBuf)
 			if err != nil {
-				slog.Error("建立数据通道失败", err)
+				slog.Error("建立数据通道失败", "err", err)
 			}
 		}()
 	}
 }
 
-func data(addr string, tag []byte) error {
+func data(dataAddr string, dest string, tag []byte) error {
 
 	// 向服务端建立连接
-	src, err := net.Dial("tcp", frpDataAddr)
+	src, err := net.Dial("tcp", dataAddr)
 	if err != nil {
-		return errors.Wrap(err, "连接服务端失败")
+		return errors.New("连接服务端失败")
 	}
 
 	tagLen := byte(len(tag))
@@ -152,7 +165,7 @@ func data(addr string, tag []byte) error {
 	copy(tagBuf[2:], tag)
 
 	// 向目标地址建立连接
-	dst, dstErr := net.Dial("tcp", addr)
+	dst, dstErr := net.Dial("tcp", dest)
 	if dstErr != nil {
 		tagBuf[0] = 0
 	} else {
@@ -166,7 +179,7 @@ func data(addr string, tag []byte) error {
 		if dst != nil {
 			utils.Close(dst)
 		}
-		return errors.Wrap(err, "发送连接状态失败")
+		return errors.New("发送连接状态失败")
 	}
 
 	if tagBuf[0] == 0 {
@@ -174,7 +187,7 @@ func data(addr string, tag []byte) error {
 		if dst != nil {
 			utils.Close(dst)
 		}
-		return errors.Wrap(dstErr, "连接目标地址失败")
+		return errors.New("连接目标地址失败")
 	}
 
 	go func() {
@@ -193,31 +206,3 @@ func data(addr string, tag []byte) error {
 	}()
 	return nil
 }
-
-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.FwdHost = os.Getenv("FWD_HOST")
-}