package edge import ( "bufio" "context" "encoding/binary" "errors" "fmt" "io" "log/slog" "net" _ "net/http/pprof" "os" "os/signal" "proxy-server/edge/core" "proxy-server/edge/env" "proxy-server/edge/report" "proxy-server/utils" "time" ) func Start() error { var ctx, cancel = signal.NotifyContext(context.Background(), os.Interrupt, os.Kill) defer cancel() // 初始化环境变量 slog.Debug("初始化环境变量...") err := env.Init() if err != nil { return fmt.Errorf("初始化环境变量失败: %w", err) } // 注册节点 slog.Debug("注册节点...") id, host, err := report.Online() if err != nil { return fmt.Errorf("注册节点失败: %w", err) } // 连接到网关 var errCh = make(chan error) go func() { err = ctrl(ctx, id, host) if err == nil { errCh <- err } }() // 等待退出 select { case <-ctx.Done(): case err := <-errCh: if err != nil { slog.Error("控制通道发生错误", "err", err) } } return ctx.Err() } func ctrl(ctx context.Context, id int32, host string) error { ctrlAddr := net.JoinHostPort(host, fmt.Sprintf("%d", core.FwdCtrlPort)) dataAddr := net.JoinHostPort(host, fmt.Sprintf("%d", core.FwdDataPort)) slog.Info("建立控制通道", "addr", ctrlAddr) conn, err := net.Dial("tcp", ctrlAddr) if err != nil { return errors.New("连接失败") } defer utils.Close(conn) var reader = bufio.NewReader(conn) // 发送开启连接 err = sendOpen(conn, id) if err != nil { return fmt.Errorf("发送节点信息失败: %w", err) } // 异步定时发送心跳 go func() { ticker := time.NewTicker(time.Duration(core.HeartbeatInterval) * time.Second) defer ticker.Stop() for { select { case <-ctx.Done(): return case tick := <-ticker.C: err := sendPing(conn) if err != nil { slog.Error("发送心跳失败", "time", tick, "err", err) } } } }() // 异步读取节点命令 var errCh = make(chan error) go func() { for { // 读取命令 cmd, err := reader.ReadByte() if ok, err := utils.WarpConnErr(err); !ok { errCh <- err return } switch cmd { // pong 命令,忽略 case 1: // 代理命令 case 5: err := onConn(reader, dataAddr) if err != nil { errCh <- fmt.Errorf("处理代理命令失败: %w", err) return } } } }() // 等待建立数据通道 select { case <-ctx.Done(): case err = <-errCh: } // 发送关闭连接 slog.Debug("发送关闭连接") err = sendClose(conn) if err != nil { return fmt.Errorf("发送关闭连接失败: %w", err) } return nil } func data(proxy string, destination string, tag [16]byte) error { slog.Debug("建立数据通道", "tag", tag, "addr", destination) // 向目标地址建立连接 var result = 1 var dstErr error dst, err := net.Dial("tcp", destination) if err != nil { dstErr = fmt.Errorf("连接目标地址失败: %w", dstErr) result = 0 } defer utils.Close(dst) // 向服务端建立连接 src, err := net.Dial("tcp", proxy) if err != nil { return errors.New("连接服务端失败") } defer utils.Close(src) // 发送连接状态 var buf = make([]byte, 17) copy(buf[0:16], tag[:]) buf[16] = byte(result) _, err = src.Write(buf) if err != nil { return errors.New("发送连接状态失败") } if result == 0 { return dstErr } var waitSrc = make(chan error) go func() { _, err := io.Copy(dst, src) switch { case errors.Is(err, net.ErrClosed): slog.Debug("网关连接意外关闭") case err != nil: slog.Error("读取网关数据失败", "err", err) default: slog.Debug("网关数据读取完成") } waitSrc <- err }() var waitDst = make(chan error) go func() { _, err := io.Copy(src, dst) switch { case errors.Is(err, net.ErrClosed): slog.Debug("目标连接意外关闭") case err != nil && !errors.Is(err, net.ErrClosed): slog.Error("读取目标数据失败", "err", err) default: slog.Debug("目标数据读取完成") } waitDst <- err }() // 等待任意一方关闭数据连接 select { case <-waitSrc: case <-waitDst: } return nil } func sendOpen(writer io.Writer, id int32) error { // 发送打开连接 var buf = make([]byte, 5) buf[0] = 3 binary.BigEndian.PutUint32(buf[1:], uint32(id)) _, err := writer.Write(buf) if err != nil { return fmt.Errorf("发送打开连接失败: %w", err) } return nil } func sendClose(writer io.Writer) error { // 发送关闭连接 _, err := writer.Write([]byte{4}) if err != nil { return err } return nil } func sendPing(writer io.Writer) error { _, err := writer.Write([]byte{2}) if err != nil { return err } return nil } func onConn(reader io.Reader, dataAddr string) (err error) { // 读取连接命令 var buf = make([]byte, 16+2) _, err = io.ReadFull(reader, buf) if err != nil { return err } var tag = [16]byte(buf[0:16]) var addrLen = binary.BigEndian.Uint16(buf[16:18]) var addrBuf = make([]byte, addrLen) _, err = io.ReadFull(reader, addrBuf) if err != nil { return err } var addr = string(addrBuf) // 异步建立数据通道 go func() { err := data(dataAddr, addr, tag) if err != nil { slog.Error("建立数据通道失败", "err", err) } }() return nil } type ConnCmd struct { Tag [16]byte Addr string }