Initial commit
This commit is contained in:
@@ -0,0 +1,99 @@
|
||||
package transport
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// ConnManager 连接管理器
|
||||
// 维护 peer_key → net.Conn 的映射关系
|
||||
type ConnManager struct {
|
||||
conns map[string]net.Conn
|
||||
mu sync.RWMutex
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewConnManager 创建连接管理器
|
||||
func NewConnManager(logger *zap.Logger) *ConnManager {
|
||||
return &ConnManager{
|
||||
conns: make(map[string]net.Conn),
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// Add 添加连接
|
||||
func (m *ConnManager) Add(peerKey string, conn net.Conn) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
// 如果已存在，先关闭旧连接
|
||||
if oldConn, ok := m.conns[peerKey]; ok {
|
||||
oldConn.Close()
|
||||
m.logger.Debug("关闭旧连接", zap.String("peer_key", peerKey))
|
||||
}
|
||||
|
||||
m.conns[peerKey] = conn
|
||||
m.logger.Info("添加新连接",
|
||||
zap.String("peer_key", peerKey),
|
||||
zap.String("remote_addr", conn.RemoteAddr().String()))
|
||||
}
|
||||
|
||||
// Get 获取连接
|
||||
func (m *ConnManager) Get(peerKey string) (net.Conn, bool) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
conn, ok := m.conns[peerKey]
|
||||
return conn, ok
|
||||
}
|
||||
|
||||
// Remove 移除连接
|
||||
func (m *ConnManager) Remove(peerKey string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if conn, ok := m.conns[peerKey]; ok {
|
||||
conn.Close()
|
||||
delete(m.conns, peerKey)
|
||||
m.logger.Info("移除连接", zap.String("peer_key", peerKey))
|
||||
}
|
||||
}
|
||||
|
||||
// Count 获取连接数量
|
||||
func (m *ConnManager) Count() int {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return len(m.conns)
|
||||
}
|
||||
|
||||
// CloseAll 关闭所有连接
|
||||
func (m *ConnManager) CloseAll() {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
for peerKey, conn := range m.conns {
|
||||
conn.Close()
|
||||
m.logger.Debug("关闭连接", zap.String("peer_key", peerKey))
|
||||
}
|
||||
|
||||
m.conns = make(map[string]net.Conn)
|
||||
}
|
||||
|
||||
// List 列出所有连接（返回副本）
|
||||
func (m *ConnManager) List() map[string]net.Conn {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
result := make(map[string]net.Conn)
|
||||
for k, v := range m.conns {
|
||||
result[k] = v
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// GetAll 获取所有连接（同 List，为了兼容）
|
||||
func (m *ConnManager) GetAll() map[string]net.Conn {
|
||||
return m.List()
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package transport
|
||||
|
||||
// ProtocolPlugin 协议插件接口
|
||||
// relay.go 通过这个接口适配不同协议，不感知具体协议细节
|
||||
type ProtocolPlugin interface {
|
||||
// IsControlPacket 判断是否为控制包
|
||||
// 控制包用于建连协商，需要透传到对端
|
||||
IsControlPacket(packet []byte) bool
|
||||
|
||||
// IsDataPacket 判断是否为数据包
|
||||
// 数据包包含路由标识，需要查表转发
|
||||
IsDataPacket(packet []byte) bool
|
||||
|
||||
// ExtractRouteID 从数据包中提取路由标识
|
||||
// 返回的 route_id 用于查找对应的本地端口
|
||||
ExtractRouteID(packet []byte) (uint32, error)
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
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))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user