feat: 将架构改成支持中间两个server

2026-03-27 19:08:22 +08:00
parent 8cec6a0766
commit 5be3ff670f
4 changed files with 609 additions and 8 deletions
--- a/cmd/internal/server/kcp_hub.go
+++ b/cmd/internal/server/kcp_hub.go
@@ -1,7 +1,11 @@
 package server

 import (
+	"errors"
 	"fmt"
+	"log"
+	"net"
+	"strings"
 	"sync"
 	"time"

@@ -12,6 +16,15 @@ import (
 	"omnisocketgo/cmd/internal/transport"
 )

+const kcpRelayMaxDatagramSize = 60 * 1024
+
+var (
+	errKCPRelayUnavailable   = errors.New("server: kcp relay socket is not configured")
+	errKCPRelayPeerUnknown   = errors.New("server: kcp relay peer address is unknown")
+	errKCPRelayTooLarge      = errors.New("server: kcp relay message too large")
+	errKCPUnknownLocalTarget = errors.New("server: unknown local kcp target")
+)
+
 // KCPOption 用于配置 KCPHub 的可选行为。
 type KCPOption func(*KCPHub)

@@ -37,6 +50,9 @@ type KCPHub struct {
 	logger               latencylog.Logger
 	sessionStatsLogger   transport.KCPSessionStatsLogger
 	sessionStatsInterval time.Duration
+	relaySocket          net.PacketConn
+	relayPeerAddr        net.Addr
+	relayLearnPeer       bool
 }

 // NewKCPHub 创建一个空的 KCP 连接中心。
@@ -54,6 +70,16 @@ func NewKCPHub(opts ...KCPOption) *KCPHub {
 	return hub
 }

+// SetRelaySocket 配置 KCPHub 的原始 UDP relay 信道。
+func (h *KCPHub) SetRelaySocket(conn net.PacketConn, peerAddr net.Addr, learnPeer bool) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	h.relaySocket = conn
+	h.relayPeerAddr = cloneRelayAddr(peerAddr)
+	h.relayLearnPeer = learnPeer
+}
+
 // HasPeer 返回给定 ID 是否已经注册到 hub。
 func (h *KCPHub) HasPeer(peerID string) bool {
 	h.mu.RLock()
@@ -63,6 +89,48 @@ func (h *KCPHub) HasPeer(peerID string) bool {
 	return ok
 }

+// ServeRelay 持续从 relay UDP socket 读取消息，并尝试本地投递。
+func (h *KCPHub) ServeRelay() error {
+	h.mu.RLock()
+	conn := h.relaySocket
+	h.mu.RUnlock()
+
+	if conn == nil {
+		return errKCPRelayUnavailable
+	}
+
+	buffer := make([]byte, kcpRelayMaxDatagramSize)
+	for {
+		n, addr, err := conn.ReadFrom(buffer)
+		if err != nil {
+			if isExpectedRelayServeExit(err) {
+				return nil
+			}
+			return fmt.Errorf("server: relay receive packet: %w", err)
+		}
+
+		if !h.acceptRelayPeer(addr) {
+			log.Printf("kcp relay dropped packet from unexpected peer %s", addr)
+			continue
+		}
+
+		msg, err := protocol.DecodeMessage(buffer[:n])
+		if err != nil {
+			log.Printf("kcp relay dropped invalid packet from %s: %v", addr, err)
+			continue
+		}
+
+		if !isRelayBusinessOrErrorMessage(msg.Type) {
+			log.Printf("kcp relay dropped unsupported message type %s from %s", msg.Type, addr)
+			continue
+		}
+
+		if err := h.deliverRelayedMessage(msg); err != nil {
+			log.Printf("kcp relay delivery for %s -> %s failed: %v", msg.From, msg.To, err)
+		}
+	}
+}
+
 // ServeSession 处理一条新接入的 KCP 会话。
 func (h *KCPHub) ServeSession(session *kcp.UDPSession) error {
 	conn, err := transport.NewKCPConn(
@@ -118,16 +186,24 @@ func (h *KCPHub) handlePeerMessage(peerID string, conn *transport.KCPConn, msg p
 	switch msg.Type {
 	case protocol.MessageTypeText, protocol.MessageTypeFile:
 		msg.From = peerID
-		targetConn, ok := h.lookup(msg.To)
-		if !ok {
-			return sendKCPServerError(conn, peerID, fmt.Sprintf("unknown target: %s", msg.To))
-		}
-		if err := targetConn.Send(msg); err != nil {
-			h.unregister(msg.To, targetConn)
-			_ = targetConn.Close()
+
+		if err := h.deliverToLocalPeer(msg); err == nil {
+			return nil
+		} else if !errors.Is(err, errKCPUnknownLocalTarget) {
 			return sendKCPServerError(conn, peerID, fmt.Sprintf("failed to forward to %s", msg.To))
 		}
-		return nil
+
+		err := h.forwardToRelay(msg)
+		switch {
+		case err == nil:
+			return nil
+		case errors.Is(err, errKCPRelayUnavailable):
+			return sendKCPServerError(conn, peerID, fmt.Sprintf("unknown target: %s", msg.To))
+		case errors.Is(err, errKCPRelayTooLarge):
+			return sendKCPServerError(conn, peerID, "message too large for relay udp")
+		default:
+			return sendKCPServerError(conn, peerID, "failed to relay to remote peer")
+		}
 	case protocol.MessageTypeRegister, protocol.MessageTypeError:
 		if err := sendKCPServerError(conn, peerID, "registered peers can only send text or file messages"); err != nil {
 			return fmt.Errorf("server: send kcp protocol error: %w", err)
@@ -157,6 +233,88 @@ func (h *KCPHub) receivePeerLoop(peerID string, conn *transport.KCPConn) error {
 	}
 }

+func (h *KCPHub) deliverRelayedMessage(msg protocol.Message) error {
+	if err := h.deliverToLocalPeer(msg); err == nil {
+		return nil
+	} else if !errors.Is(err, errKCPUnknownLocalTarget) {
+		if msg.Type == protocol.MessageTypeError {
+			log.Printf("kcp relay dropped undeliverable server error to %s: %v", msg.To, err)
+			return nil
+		}
+		return h.forwardRelayServerError(msg.From, fmt.Sprintf("failed to forward to %s", msg.To))
+	}
+
+	if msg.Type == protocol.MessageTypeError {
+		log.Printf("kcp relay dropped server error for unknown local peer %s", msg.To)
+		return nil
+	}
+
+	return h.forwardRelayServerError(msg.From, fmt.Sprintf("unknown target: %s", msg.To))
+}
+
+func (h *KCPHub) deliverToLocalPeer(msg protocol.Message) error {
+	targetConn, ok := h.lookup(msg.To)
+	if !ok {
+		return fmt.Errorf("%w: %s", errKCPUnknownLocalTarget, msg.To)
+	}
+	if err := targetConn.Send(msg); err != nil {
+		h.unregister(msg.To, targetConn)
+		_ = targetConn.Close()
+		return fmt.Errorf("server: forward to local peer %s: %w", msg.To, err)
+	}
+	return nil
+}
+
+func (h *KCPHub) forwardToRelay(msg protocol.Message) error {
+	payload, err := protocol.EncodeMessage(msg)
+	if err != nil {
+		return fmt.Errorf("server: encode relay message: %w", err)
+	}
+	if len(payload) > kcpRelayMaxDatagramSize {
+		return errKCPRelayTooLarge
+	}
+
+	h.mu.RLock()
+	conn := h.relaySocket
+	peerAddr := cloneRelayAddr(h.relayPeerAddr)
+	h.mu.RUnlock()
+
+	if conn == nil {
+		return errKCPRelayUnavailable
+	}
+	if peerAddr == nil {
+		return errKCPRelayPeerUnknown
+	}
+
+	if _, err := conn.WriteTo(payload, peerAddr); err != nil {
+		return fmt.Errorf("server: relay write to %s: %w", peerAddr, err)
+	}
+	return nil
+}
+
+func (h *KCPHub) forwardRelayServerError(to, message string) error {
+	return h.forwardToRelay(protocol.Message{
+		Type: protocol.MessageTypeError,
+		From: protocol.ServerPeerID,
+		To:   to,
+		Body: []byte(message),
+	})
+}
+
+func (h *KCPHub) acceptRelayPeer(addr net.Addr) bool {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	if h.relayPeerAddr == nil && h.relayLearnPeer {
+		h.relayPeerAddr = cloneRelayAddr(addr)
+		return true
+	}
+	if h.relayPeerAddr == nil {
+		return true
+	}
+	return sameRelayAddr(h.relayPeerAddr, addr)
+}
+
 func (h *KCPHub) lookup(peerID string) (*transport.KCPConn, bool) {
 	h.mu.RLock()
 	defer h.mu.RUnlock()
@@ -183,3 +341,40 @@ func sendKCPServerError(conn *transport.KCPConn, to, message string) error {
 		Body: []byte(message),
 	})
 }
+
+func isRelayBusinessOrErrorMessage(messageType protocol.MessageType) bool {
+	switch messageType {
+	case protocol.MessageTypeText, protocol.MessageTypeFile, protocol.MessageTypeError:
+		return true
+	default:
+		return false
+	}
+}
+
+func isExpectedRelayServeExit(err error) bool {
+	return errors.Is(err, net.ErrClosed) || strings.Contains(err.Error(), "use of closed network connection")
+}
+
+func cloneRelayAddr(addr net.Addr) net.Addr {
+	if addr == nil {
+		return nil
+	}
+	udpAddr, ok := addr.(*net.UDPAddr)
+	if !ok {
+		return addr
+	}
+	ipCopy := make([]byte, len(udpAddr.IP))
+	copy(ipCopy, udpAddr.IP)
+	return &net.UDPAddr{
+		IP:   ipCopy,
+		Port: udpAddr.Port,
+		Zone: udpAddr.Zone,
+	}
+}
+
+func sameRelayAddr(left, right net.Addr) bool {
+	if left == nil || right == nil {
+		return left == right
+	}
+	return left.String() == right.String()
+}