proxy/edge/edge.go

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/geo"
	"proxy-server/edge/report"
	"proxy-server/utils"
	"time"
)

func Start() error {

	// 初始化环境变量
	slog.Debug("初始化环境变量...")
	err := env.Init()
	if err != nil {
		return fmt.Errorf("初始化环境变量失败: %w", err)
	}

	// 获取归属地
	slog.Debug("获取节点归属地...")
	err = geo.Query()
	if err != nil {
		return fmt.Errorf("获取节点归属地失败: %w", err)
	}

	// 注册节点
	slog.Debug("注册节点...")
	id, host, err := report.Online(geo.Prov, geo.City, geo.Isp)
	if err != nil {
		return fmt.Errorf("注册节点失败: %w", err)
	}

	// 连接到网关
	var ctx, cancel = signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
	defer cancel()

	var errCh = make(chan error)
	go func() {
		for {
			err = ctrl(ctx, id, host)
			if err == nil {
				errCh <- nil
				return
			}

			select {
			case <-ctx.Done():
				return
			default:
				slog.Error("建立控制通道失败", "err", err)
				slog.Info(fmt.Sprintf("%d 秒后重试", core.RetryInterval))
			}

			select {
			case <-ctx.Done():
				return
			case <-time.After(time.Duration(core.RetryInterval) * time.Second):
			}
		}
	}()

	// 等待退出
	select {
	case err := <-errCh:
		if err != nil {
			slog.Error("控制通道发生错误", "err", err)
		}
	}

	// 下线节点
	slog.Debug("下线节点...")
	err = report.Offline()
	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)

	// 发送节点连接命令
	slog.Debug("发送节点连接命令")
	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)
				}
			}
		}
	}()

	// 异步等待连接命令
	slog.Info("等待用户连接")
	var cmdCh = make(chan ConnCmd)
	var errCh = make(chan error)
	go func() {
		for {
			cmd, err := reader.ReadByte()
			if err != nil {
				switch {
				case errors.Is(err, net.ErrClosed):
					err = fmt.Errorf("控制通道关闭: %w", err)
				case errors.Is(err, io.EOF):
					err = fmt.Errorf("网关关闭了控制通道: %w", err)
				default:
					err = fmt.Errorf("读取命令失败: %w", err)
				}
				errCh <- err
				return
			}

			switch cmd {
			case 1:
				// 忽略网关响应的 pong 命令
			case 5:
				tag, addr, err := onConn(reader)
				if err != nil {
					slog.Error("接收连接命令失败", "err", err)
					return
				}
				cmdCh <- ConnCmd{
					Tag:  tag,
					Addr: addr,
				}
			}
		}
	}()

	// 等待建立数据通道
	for loop := true; loop; {
		select {
		case <-ctx.Done():
			loop = false
		case err = <-errCh:
			slog.Error("读取控制命令失败", "err", err)
			loop = false
		case cmd := <-cmdCh:
			slog.Debug("建立数据通道", "tag", cmd.Tag, "addr", cmd.Addr)
			go func() {
				err := data(dataAddr, cmd.Addr, cmd.Tag)
				if err != nil {
					slog.Error("建立数据通道失败", "err", err)
				}
			}()
		}
	}

	// 发送关闭连接（不 return err，否则会重新连接）
	slog.Debug("发送关闭连接")
	err = sendClose(conn)
	if err != nil {
		slog.Error("发送关闭连接失败", "err", err)
	}

	return nil
}

func data(proxy string, dest string, tag [16]byte) error {

	// 向目标地址建立连接
	var result = 1
	var dstErr error
	dst, err := net.Dial("tcp", dest)
	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) (tag [16]byte, addr string, err error) {
	var buf = make([]byte, 16+2)
	_, err = io.ReadFull(reader, buf)
	if err != nil {
		return [16]byte{}, "", err
	}

	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 [16]byte{}, "", err
	}

	addr = string(addrBuf)
	return tag, addr, nil
}

type ConnCmd struct {
	Tag  [16]byte
	Addr string
}