limingjie/OmniSocketGo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat:多跳（B->D->C->A）功能

Browse Source
This commit is contained in:
nnbcccscdscdsc
2026-03-27 23:03:00 +08:00
parent 5be3ff670f
commit 8e2bd0ffc6
8 changed files with 455 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
cmd/internal/peer/client.go
View File

@@ -25,6 +25,7 @@ type clientOptions struct {
udpLinuxTimestamping bool
bindIP string
bindDevice string
kcpDialAddress string
}
// Option 用于配置 Client 的可选行为，例如时延日志。
@@ -73,6 +74,13 @@ func WithBindDevice(device string) Option {
}
}
// WithKCPDialAddress 指定 KCP 实际拨号使用的 UDP 地址，可用于通过 relay 连接逻辑上的 server。
func WithKCPDialAddress(addr string) Option {
return func(options *clientOptions) {
options.kcpDialAddress = addr
}
}
// WithUDPLinuxTimestamping controls whether UDP clients enable Linux timestamping.
func WithUDPLinuxTimestamping(enabled bool) Option {
return func(options *clientOptions) {

7
cmd/internal/peer/kcp_client.go
View File

@@ -32,8 +32,13 @@ func DialKCP(serverAddr, peerID string, opts ...Option) (*KCPClient, error) {
options.logger = latencylog.NoopLogger{}
}
dialAddr := serverAddr
if options.kcpDialAddress != "" {
dialAddr = options.kcpDialAddress
}
session, err := transport.DialKCPSession(
serverAddr,
dialAddr,
options.bindIP,
options.bindDevice,
options.kcpPacketDebugLogger,

90
cmd/internal/peer/kcp_client_test.go
View File

@@ -267,6 +267,96 @@ func TestKCPClientsExchangeMessagesAcrossRelayedServers(t *testing.T) {
}
}
func TestKCPClientsExchangeMessagesViaUDPRelayToSingleHub(t *testing.T) {
hub := server.NewKCPHub()
serverAddr, cleanupHub := startRealKCPHubServer(t, hub)
defer cleanupHub()
remoteAddr, err := net.ResolveUDPAddr("udp", serverAddr)
if err != nil {
t.Fatalf("ResolveUDPAddr(server) error = %v", err)
}
baseRelayConn, err := net.ListenPacket("udp", "127.0.0.1:0")
if err != nil {
t.Fatalf("ListenPacket(relay) error = %v", err)
}
relayConn := &countingPacketConn{PacketConn: baseRelayConn}
relay, err := server.NewUDPRelay(relayConn, remoteAddr)
if err != nil {
_ = relayConn.Close()
t.Fatalf("NewUDPRelay() error = %v", err)
}
var relayWG sync.WaitGroup
relayWG.Add(1)
go func() {
defer relayWG.Done()
if serveErr := relay.Serve(); serveErr != nil {
t.Errorf("relay.Serve() error = %v", serveErr)
}
}()
defer func() {
_ = relayConn.Close()
relayWG.Wait()
}()
peerA, err := DialKCP(serverAddr, "peer-a", WithKCPDialAddress(relayConn.LocalAddr().String()))
if err != nil {
t.Fatalf("DialKCP(peer-a via relay) error = %v", err)
}
defer func() { _ = peerA.Close() }()
peerB, err := DialKCP(serverAddr, "peer-b")
if err != nil {
t.Fatalf("DialKCP(peer-b direct) error = %v", err)
}
defer func() { _ = peerB.Close() }()
waitFor(t, func() bool { return hub.HasPeer("peer-a") && hub.HasPeer("peer-b") }, "both peers to be registered on the single hub")
if err := peerB.SendText("peer-a", "hello via udp relay"); err != nil {
t.Fatalf("peerB.SendText() error = %v", err)
}
gotAtA, err := peerA.Receive()
if err != nil {
t.Fatalf("peerA.Receive() error = %v", err)
}
wantAtA := protocol.Message{
Type: protocol.MessageTypeText,
ID: 1,
From: "peer-b",
To: "peer-a",
Body: []byte("hello via udp relay"),
}
if !reflect.DeepEqual(gotAtA, wantAtA) {
t.Fatalf("peerA received %+v, want %+v", gotAtA, wantAtA)
}
if err := peerA.SendText("peer-b", "hello back through relay"); err != nil {
t.Fatalf("peerA.SendText() error = %v", err)
}
gotAtB, err := peerB.Receive()
if err != nil {
t.Fatalf("peerB.Receive() error = %v", err)
}
wantAtB := protocol.Message{
Type: protocol.MessageTypeText,
ID: 1,
From: "peer-a",
To: "peer-b",
Body: []byte("hello back through relay"),
}
if !reflect.DeepEqual(gotAtB, wantAtB) {
t.Fatalf("peerB received %+v, want %+v", gotAtB, wantAtB)
}
if got := relayConn.WriteCount(); got == 0 {
t.Fatal("relay should have forwarded packets for peer-a session")
}
}
func TestKCPHubPrefersLocalPeerBeforeRelay(t *testing.T) {
fixture := startRelayedKCPHubs(t)
defer fixture.cleanup()

127
cmd/internal/server/udp_relay.go Normal file
View File

@@ -0,0 +1,127 @@
package server
import (
"fmt"
"log"
"net"
"sync"
"omnisocketgo/cmd/internal/transport"
)
// UDPRelay 负责在固定远端与多个客户端之间双向透明转发 KCP UDP datagram。
type UDPRelay struct {
conn net.PacketConn
remote *net.UDPAddr
mu sync.RWMutex
clients map[uint32]*net.UDPAddr
}
// NewUDPRelay 创建一个绑定到给定 PacketConn 的透明 UDP relay。
func NewUDPRelay(conn net.PacketConn, remote *net.UDPAddr) (*UDPRelay, error) {
if conn == nil {
return nil, fmt.Errorf("server: nil udp relay conn")
}
if remote == nil {
return nil, fmt.Errorf("server: nil udp relay remote")
}
return &UDPRelay{
conn: conn,
remote: cloneUDPAddr(remote),
clients: make(map[uint32]*net.UDPAddr),
}, nil
}
// Serve 持续双向转发原始 UDP datagram，不解析业务消息。
func (r *UDPRelay) Serve() error {
buffer := make([]byte, 64*1024)
for {
n, addr, err := r.conn.ReadFrom(buffer)
if err != nil {
if isExpectedRelayServeExit(err) {
return nil
}
return fmt.Errorf("server: udp relay read packet: %w", err)
}
udpAddr, ok := addr.(*net.UDPAddr)
if !ok {
log.Printf("udp relay dropped packet from non-udp addr %T", addr)
continue
}
payload := append([]byte(nil), buffer[:n]...)
if sameUDPAddr(udpAddr, r.remote) {
if err := r.forwardRemotePacket(payload); err != nil {
log.Printf("udp relay failed forwarding remote packet from %s: %v", udpAddr, err)
}
continue
}
if err := r.forwardClientPacket(udpAddr, payload); err != nil {
log.Printf("udp relay failed forwarding client packet from %s: %v", udpAddr, err)
}
}
}
func (r *UDPRelay) forwardClientPacket(addr *net.UDPAddr, payload []byte) error {
convID, ok := transport.ParseKCPConversationID(payload)
if !ok {
return fmt.Errorf("missing kcp conversation id")
}
r.mu.Lock()
r.clients[convID] = cloneUDPAddr(addr)
r.mu.Unlock()
if _, err := r.conn.WriteTo(payload, r.remote); err != nil {
return fmt.Errorf("write conv %d to remote %s: %w", convID, r.remote, err)
}
return nil
}
func (r *UDPRelay) forwardRemotePacket(payload []byte) error {
convID, ok := transport.ParseKCPConversationID(payload)
if !ok {
return fmt.Errorf("missing kcp conversation id")
}
r.mu.RLock()
clientAddr := cloneUDPAddr(r.clients[convID])
r.mu.RUnlock()
if clientAddr == nil {
return fmt.Errorf("unknown client for conv %d", convID)
}
if _, err := r.conn.WriteTo(payload, clientAddr); err != nil {
return fmt.Errorf("write conv %d to client %s: %w", convID, clientAddr, err)
}
return nil
}
func cloneUDPAddr(addr *net.UDPAddr) *net.UDPAddr {
if addr == nil {
return nil
}
ipCopy := make([]byte, len(addr.IP))
copy(ipCopy, addr.IP)
return &net.UDPAddr{
IP: ipCopy,
Port: addr.Port,
Zone: addr.Zone,
}
}
func sameUDPAddr(left, right *net.UDPAddr) bool {
if left == nil || right == nil {
return left == right
}
if left.Port != right.Port || left.Zone != right.Zone {
return false
}
return left.IP.Equal(right.IP)
}

103
cmd/internal/server/udp_relay_test.go Normal file
View File

@@ -0,0 +1,103 @@
package server
import (
"encoding/binary"
"net"
"sync"
"testing"
"time"
)
func TestUDPRelayRoutesPacketsByKCPConversationID(t *testing.T) {
remote, err := net.ListenPacket("udp", "127.0.0.1:0")
if err != nil {
t.Fatalf("ListenPacket(remote) error = %v", err)
}
defer remote.Close()
relayConn, err := net.ListenPacket("udp", "127.0.0.1:0")
if err != nil {
t.Fatalf("ListenPacket(relay) error = %v", err)
}
relay, err := NewUDPRelay(relayConn, remote.LocalAddr().(*net.UDPAddr))
if err != nil {
_ = relayConn.Close()
t.Fatalf("NewUDPRelay() error = %v", err)
}
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
if serveErr := relay.Serve(); serveErr != nil {
t.Errorf("relay.Serve() error = %v", serveErr)
}
}()
defer func() {
_ = relayConn.Close()
wg.Wait()
}()
client1, err := net.ListenPacket("udp", "127.0.0.1:0")
if err != nil {
t.Fatalf("ListenPacket(client1) error = %v", err)
}
defer client1.Close()
client2, err := net.ListenPacket("udp", "127.0.0.1:0")
if err != nil {
t.Fatalf("ListenPacket(client2) error = %v", err)
}
defer client2.Close()
relayAddr := relayConn.LocalAddr()
sendPacket(t, client1, relayAddr, buildRelayTestPacket(1, []byte("client-one")))
assertPacketReceived(t, remote, buildRelayTestPacket(1, []byte("client-one")))
sendPacket(t, client2, relayAddr, buildRelayTestPacket(2, []byte("client-two")))
assertPacketReceived(t, remote, buildRelayTestPacket(2, []byte("client-two")))
sendPacket(t, remote, relayAddr, buildRelayTestPacket(2, []byte("reply-two")))
assertPacketReceived(t, client2, buildRelayTestPacket(2, []byte("reply-two")))
sendPacket(t, remote, relayAddr, buildRelayTestPacket(1, []byte("reply-one")))
assertPacketReceived(t, client1, buildRelayTestPacket(1, []byte("reply-one")))
}
func buildRelayTestPacket(convID uint32, body []byte) []byte {
packet := make([]byte, 4+len(body))
binary.LittleEndian.PutUint32(packet[:4], convID)
copy(packet[4:], body)
return packet
}
func sendPacket(t *testing.T, conn net.PacketConn, addr net.Addr, payload []byte) {
t.Helper()
if err := conn.SetWriteDeadline(time.Now().Add(2 * time.Second)); err != nil {
t.Fatalf("SetWriteDeadline() error = %v", err)
}
if _, err := conn.WriteTo(payload, addr); err != nil {
t.Fatalf("WriteTo(%s) error = %v", addr, err)
}
}
func assertPacketReceived(t *testing.T, conn net.PacketConn, want []byte) {
t.Helper()
if err := conn.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil {
t.Fatalf("SetReadDeadline() error = %v", err)
}
buffer := make([]byte, 1024)
n, _, err := conn.ReadFrom(buffer)
if err != nil {
t.Fatalf("ReadFrom() error = %v", err)
}
got := buffer[:n]
if string(got) != string(want) {
t.Fatalf("packet = %v, want %v", got, want)
}
}

9
cmd/internal/transport/kcp_packet_metadata.go
View File

@@ -21,6 +21,15 @@ func parseKCPConversationID(packet []byte) *uint32 {
return &conv
}
// ParseKCPConversationID 从原始 KCP UDP datagram 中提取 conv ID。
func ParseKCPConversationID(packet []byte) (uint32, bool) {
conv := parseKCPConversationID(packet)
if conv == nil {
return 0, false
}
return *conv, true
}
func parseKCPPacketSegments(packet []byte) ([]KCPPacketDebugSegment, bool) {
if len(packet) == 0 {
return nil, false

8
cmd/kcppeer/main.go
View File

@@ -17,6 +17,7 @@ import (
func main() {
peerID := flag.String("id", "peer-a", "peer identity")
serverAddr := flag.String("server", "127.0.0.1:9002", "KCP server address")
relayVia := flag.String("relay-via", "", "optional UDP relay address used to reach the KCP server")
targetPeer := flag.String("to", "", "optional target peer for one outgoing message")
text := flag.String("text", "", "optional text to send after connecting")
filePath := flag.String("file", "", "optional file path to send after connecting")
@@ -66,6 +67,9 @@ func main() {
if *bindDevice != "" {
clientOptions = append(clientOptions, peerpkg.WithBindDevice(*bindDevice))
}
if *relayVia != "" {
clientOptions = append(clientOptions, peerpkg.WithKCPDialAddress(*relayVia))
}
client, err := peerpkg.DialKCP(*serverAddr, *peerID, clientOptions...)
if err != nil {
@@ -73,7 +77,11 @@ func main() {
}
defer client.Close()
if *relayVia != "" {
log.Printf("connected to %s via relay %s as %s (KCP)", *serverAddr, *relayVia, client.ID())
} else {
log.Printf("connected to %s as %s (KCP)", *serverAddr, client.ID())
}
receiveErr := make(chan error, 1)
go func() {

152
cmd/kcpserver/main.go
View File

@@ -5,6 +5,7 @@ import (
"log"
"net"
"strings"
"time"
kcp "github.com/xtaci/kcp-go/v5"
@@ -13,94 +14,124 @@ import (
"omnisocketgo/cmd/internal/transport"
)
const (
kcpServerModeHub = "hub"
kcpServerModeRelay = "relay"
)
func main() {
listenAddr := flag.String("listen", ":9002", "KCP server listen address")
mode := flag.String("mode", kcpServerModeHub, "kcpserver mode: hub or relay")
listenAddr := flag.String("listen", ":9002", "listen address; KCP listener in hub mode, UDP relay listener in relay mode")
bindDevice := flag.String("bind-device", "", "optional Linux network device used when listening")
logPath := flag.String("latency-log", "", "optional JSONL file path for latency timestamp logs")
kcpTimestampDebugLogPath := flag.String("kcp-ts-debug-log", "", "optional JSONL file path for KCP packet kernel timestamp debug records")
kcpSessionStatsLogPath := flag.String("kcp-session-stats-log", "", "optional JSONL file path for KCP session stats records")
kcpSessionStatsInterval := flag.String("kcp-session-stats-interval", transport.DefaultKCPSessionStatsInterval.String(), "sampling interval for KCP session stats, for example 100ms")
relayListenAddr := flag.String("relay-listen", "", "optional raw UDP relay listen address")
relayPeerAddr := flag.String("relay-peer", "", "optional fixed raw UDP relay peer address")
relayLearnPeer := flag.Bool("relay-learn-peer", false, "learn the relay peer address from the first inbound relay packet")
relayListenAddr := flag.String("relay-listen", "", "deprecated alias for -listen in relay mode")
relayRemoteAddr := flag.String("relay-remote", "", "fixed remote UDP address used in relay mode")
relayPeerAddr := flag.String("relay-peer", "", "deprecated alias for -relay-remote")
flag.Parse()
var relayRemoteFlagSet bool
var relayPeerFlagSet bool
var relayListenFlagSet bool
flag.Visit(func(f *flag.Flag) {
switch f.Name {
case "relay-listen":
relayListenFlagSet = true
case "relay-remote":
relayRemoteFlagSet = true
case "relay-peer":
relayPeerFlagSet = true
}
})
switch {
case relayRemoteFlagSet && relayPeerFlagSet && *relayRemoteAddr != *relayPeerAddr:
log.Fatal("flags -relay-remote and -relay-peer must match when both are set")
case *relayRemoteAddr == "" && *relayPeerAddr != "":
*relayRemoteAddr = *relayPeerAddr
}
if relayPeerFlagSet {
log.Printf("warning: flag -relay-peer is deprecated; use -relay-remote instead")
}
if relayListenFlagSet {
if *relayListenAddr == "" {
log.Fatal("flag -relay-listen must not be empty when set")
}
if *mode != kcpServerModeRelay {
log.Fatal("flag -relay-listen may only be used in relay mode")
}
if *listenAddr != ":9002" && *listenAddr != *relayListenAddr {
log.Fatal("flags -listen and -relay-listen must match when both are set in relay mode")
}
*listenAddr = *relayListenAddr
log.Printf("warning: flag -relay-listen is deprecated; use -listen with -mode=relay instead")
}
statsInterval, err := transport.ParseKCPSessionStatsInterval(*kcpSessionStatsInterval)
if err != nil {
log.Fatalf("parse -kcp-session-stats-interval=%q: %v", *kcpSessionStatsInterval, err)
}
switch *mode {
case kcpServerModeHub:
if *relayRemoteAddr != "" {
log.Fatal("flag -relay-remote may only be used in relay mode")
}
runHubServer(*listenAddr, *bindDevice, *logPath, *kcpTimestampDebugLogPath, *kcpSessionStatsLogPath, statsInterval)
case kcpServerModeRelay:
if *bindDevice != "" {
log.Fatal("flag -bind-device is not supported in relay mode")
}
if *relayRemoteAddr == "" {
log.Fatal("flag -relay-remote is required in relay mode")
}
runUDPRelayServer(*listenAddr, *relayRemoteAddr)
default:
log.Fatalf("unsupported -mode=%q; want %q or %q", *mode, kcpServerModeHub, kcpServerModeRelay)
}
}
func runHubServer(listenAddr, bindDevice, logPath, packetDebugLogPath, sessionStatsLogPath string, statsInterval time.Duration) {
hubOptions := make([]server.KCPOption, 0, 2)
if *logPath != "" {
logger, err := latencylog.NewJSONLLogger(*logPath)
if logPath != "" {
logger, err := latencylog.NewJSONLLogger(logPath)
if err != nil {
log.Fatalf("create latency logger %s: %v", *logPath, err)
log.Fatalf("create latency logger %s: %v", logPath, err)
}
defer logger.Close()
hubOptions = append(hubOptions, server.WithKCPLogger(logger))
}
var packetLogger transport.KCPPacketDebugLogger
if *kcpTimestampDebugLogPath != "" {
logger, err := transport.NewJSONLKCPPacketDebugLogger(*kcpTimestampDebugLogPath)
if packetDebugLogPath != "" {
logger, err := transport.NewJSONLKCPPacketDebugLogger(packetDebugLogPath)
if err != nil {
log.Fatalf("create kcp packet debug logger %s: %v", *kcpTimestampDebugLogPath, err)
log.Fatalf("create kcp packet debug logger %s: %v", packetDebugLogPath, err)
}
defer logger.Close()
packetLogger = logger
}
if *kcpSessionStatsLogPath != "" {
logger, err := transport.NewJSONLKCPSessionStatsLogger(*kcpSessionStatsLogPath)
if sessionStatsLogPath != "" {
logger, err := transport.NewJSONLKCPSessionStatsLogger(sessionStatsLogPath)
if err != nil {
log.Fatalf("create kcp session stats logger %s: %v", *kcpSessionStatsLogPath, err)
log.Fatalf("create kcp session stats logger %s: %v", sessionStatsLogPath, err)
}
defer logger.Close()
hubOptions = append(hubOptions, server.WithKCPSessionStatsLogger(logger, statsInterval))
}
listener, packetConn, err := transport.ListenKCPSessions(*listenAddr, *bindDevice, packetLogger, latencylog.NodeRoleServer, "hub")
listener, packetConn, err := transport.ListenKCPSessions(listenAddr, bindDevice, packetLogger, latencylog.NodeRoleServer, "hub")
if err != nil {
log.Fatalf("listen kcp on %s: %v", *listenAddr, err)
log.Fatalf("listen kcp on %s: %v", listenAddr, err)
}
defer packetConn.Close()
defer listener.Close()
hub := server.NewKCPHub(hubOptions...)
if *relayPeerAddr != "" && *relayListenAddr == "" {
log.Fatal("flag -relay-listen is required when -relay-peer is set")
}
if *relayLearnPeer && *relayListenAddr == "" {
log.Fatal("flag -relay-listen is required when -relay-learn-peer is set")
}
if *relayListenAddr != "" {
relayConn, err := net.ListenPacket("udp", *relayListenAddr)
if err != nil {
log.Fatalf("listen relay udp on %s: %v", *relayListenAddr, err)
}
defer relayConn.Close()
var relayPeer net.Addr
if *relayPeerAddr != "" {
relayPeer, err = net.ResolveUDPAddr("udp", *relayPeerAddr)
if err != nil {
log.Fatalf("resolve relay peer %s: %v", *relayPeerAddr, err)
}
}
hub.SetRelaySocket(relayConn, relayPeer, *relayLearnPeer)
go func() {
if serveErr := hub.ServeRelay(); serveErr != nil {
log.Printf("kcp relay loop ended: %v", serveErr)
}
}()
log.Printf("kcp relay listening on %s", relayConn.LocalAddr())
if relayPeer != nil {
log.Printf("kcp relay peer configured as %s", relayPeer)
}
}
log.Printf("kcp server listening on %s", listener.Addr())
log.Printf("kcp hub listening on %s", listener.Addr())
for {
session, err := listener.AcceptKCP()
@@ -119,3 +150,26 @@ func main() {
}(session)
}
}
func runUDPRelayServer(listenAddr, remoteAddr string) {
conn, err := net.ListenPacket("udp", listenAddr)
if err != nil {
log.Fatalf("listen udp relay on %s: %v", listenAddr, err)
}
defer conn.Close()
remote, err := net.ResolveUDPAddr("udp", remoteAddr)
if err != nil {
log.Fatalf("resolve relay remote %s: %v", remoteAddr, err)
}
relay, err := server.NewUDPRelay(conn, remote)
if err != nil {
log.Fatalf("create udp relay: %v", err)
}
log.Printf("udp relay listening on %s and forwarding to %s", conn.LocalAddr(), remote)
if err := relay.Serve(); err != nil {
log.Fatalf("udp relay stopped: %v", err)
}
}