Meshray-Manager/core/transport/relay.go

package transport

import (
	"context"
	"net"
	"sync"

	"go.uber.org/zap"
)

// OnDialTrigger 当从本地端口收到包但没有远端连接时触发
type OnDialTrigger func(peerKey string)

// Relay 数据转发器
// 负责从本地端口收包 → 查路由 → 通过 conn 发送
// 从 conn 收包 → 发到本地端口
type Relay struct {
	plugin        ProtocolPlugin
	connMgr       *ConnManager
	localPorts    map[uint32]net.PacketConn // route_id → local_port
	lastAddr      map[uint32]net.Addr       // route_id → last wg source addr
	portMu        sync.RWMutex
	logger        *zap.Logger
	OnDialTrigger OnDialTrigger // 拨号触发回调
}

// NewRelay 创建数据转发器
func NewRelay(plugin ProtocolPlugin, connMgr *ConnManager, logger *zap.Logger) *Relay {
	return &Relay{
		plugin:     plugin,
		connMgr:    connMgr,
		localPorts: make(map[uint32]net.PacketConn),
		lastAddr:   make(map[uint32]net.Addr),
		logger:     logger,
	}
}

// RegisterLocalPort 注册本地端口（用于接收 WG 密文包）
func (r *Relay) RegisterLocalPort(routeID uint32, port net.PacketConn) {
	r.portMu.Lock()
	defer r.portMu.Unlock()

	r.localPorts[routeID] = port
	r.logger.Info("注册本地端口",
		zap.Uint32("route_id", routeID),
		zap.String("addr", port.LocalAddr().String()))
}

// UnregisterLocalPort 注销本地端口
func (r *Relay) UnregisterLocalPort(routeID uint32) {
	r.portMu.Lock()
	defer r.portMu.Unlock()

	if port, ok := r.localPorts[routeID]; ok {
		port.Close()
		delete(r.localPorts, routeID)
		delete(r.lastAddr, routeID)
		r.logger.Info("注销本地端口", zap.Uint32("route_id", routeID))
	}
}

// StartReadFromLocalPort 从本地端口读取 WG 密文包并转发（发送到远端）
func (r *Relay) StartReadFromLocalPort(ctx context.Context, routeID uint32, peerKey string) {
	r.portMu.RLock()
	port, ok := r.localPorts[routeID]
	r.portMu.RUnlock()

	if !ok {
		r.logger.Warn("本地端口未注册", zap.Uint32("route_id", routeID))
		return
	}

	go func() {
		buf := make([]byte, 65535)
		for {
			select {
			case <-ctx.Done():
				return
			default:
				n, addr, err := port.ReadFrom(buf)
				if err != nil {
					// 检查是否是由于关闭引起的错误
					select {
					case <-ctx.Done():
						return
					default:
					}
					r.logger.Debug("读取本地端口失败",
						zap.Uint32("route_id", routeID),
						zap.Error(err))
					return
				}

				// 记录 WG 的来源地址，以便后续把包发回去
				r.portMu.Lock()
				r.lastAddr[routeID] = addr
				r.portMu.Unlock()

				packet := buf[:n]
				r.forwardOutgoing(ctx, packet, peerKey)
			}
		}
	}()

	r.logger.Info("启动本地端口读取协程",
		zap.Uint32("route_id", routeID),
		zap.String("peer_key", peerKey))
}

// StartReadFromRemoteConn 从远端连接读取数据并转发给本地监听端口（接收远端数据）
func (r *Relay) StartReadFromRemoteConn(ctx context.Context, peerKey string, conn net.Conn) {
	if conn == nil {
		return
	}

	go func() {
		buf := make([]byte, 65535)
		for {
			select {
			case <-ctx.Done():
				return
			default:
				n, err := conn.Read(buf)
				if err != nil {
					r.logger.Debug("读取远端连接失败，停止读取协程",
						zap.String("peer_key", peerKey),
						zap.Error(err))
					return
				}

				packet := buf[:n]
				// 远端进来的包，需要根据 packet 里的索引转发给对应的 localPort
				r.forwardIncoming(packet, peerKey)
			}
		}
	}()

	r.logger.Info("启动远端连接读取协程",
		zap.String("peer_key", peerKey),
		zap.String("addr", conn.RemoteAddr().String()))
}

// forwardOutgoing 处理发出去的包（Local -> Remote）
func (r *Relay) forwardOutgoing(ctx context.Context, packet []byte, peerKey string) {
	// 获取或触发建连
	conn, ok := r.connMgr.Get(peerKey)
	if !ok || conn == nil {
		// 没有连接，触发拨号
		if r.OnDialTrigger != nil {
			r.OnDialTrigger(peerKey)
		}
		r.logger.Debug("尚未建立连接，包已丢弃，触发静默拨号", zap.String("peer_key", peerKey))
		return
	}

	// 转发给远端
	_, err := conn.Write(packet)
	if err != nil {
		r.logger.Debug("转发包到远端失败",
			zap.String("peer_key", peerKey),
			zap.Error(err))
	}
}

// forwardIncoming 处理进来的包（Remote -> Local）
func (r *Relay) forwardIncoming(packet []byte, _ string) {
	// 1. 判断是否为控制包/数据包并提取 routeID
	// 无论哪种 WG 包，前几位都是 routeID (receiver index)
	routeID, err := r.plugin.ExtractRouteID(packet)
	if err != nil {
		r.logger.Debug("提取包内索引失败", zap.Error(err))
		return
	}

	// 2. 这里的 routeID 是我们 RegisterLocalPort 时用的 ID
	r.portMu.RLock()
	port, ok := r.localPorts[routeID]
	addr, addrOk := r.lastAddr[routeID]
	r.portMu.RUnlock()

	if !ok || port == nil {
		r.logger.Debug("未找到转发目标的本地端口", zap.Uint32("route_id", routeID))
		return
	}

	if !addrOk || addr == nil {
		// 如果还没收到过 WG 的包，尝试发给 127.0.0.1:0 (通常不会成功，但作为 fallback)
		// 实际上 WG 发送握手包后就会刷新 addr
		addr = &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
	}

	// 3. 转发给本地 WG
	_, err = port.WriteTo(packet, addr)
	if err != nil {
		r.logger.Debug("转发给本地 WG 失败", zap.Uint32("route_id", routeID), zap.Error(err))
	}
}