feat:KCP协议

2026-03-24 21:09:06 +08:00
parent 290ba18962
commit be013b701b
20 changed files with 2284 additions and 16 deletions
--- a/cmd/internal/peer/client.go
+++ b/cmd/internal/peer/client.go
@@ -18,6 +18,7 @@ var dialServer = dialServerWithOptions
 type clientOptions struct {
 	logger                 latencylog.Logger
 	txTimestampDebugLogger transport.TXTimestampDebugLogger
+	kcpPacketDebugLogger   transport.KCPPacketDebugLogger
 	bindIP                 string
 	bindDevice             string
 }
@@ -39,6 +40,13 @@ func WithTXTimestampDebugLogger(logger transport.TXTimestampDebugLogger) Option
 	}
 }

+// WithKCPPacketDebugLogger 为 KCP UDP packet timestamp 调试日志注入记录器。
+func WithKCPPacketDebugLogger(logger transport.KCPPacketDebugLogger) Option {
+	return func(options *clientOptions) {
+		options.kcpPacketDebugLogger = logger
+	}
+}
+
 // WithBindIP 指定拨号时使用的本地源 IP。
 func WithBindIP(ip string) Option {
 	return func(options *clientOptions) {
--- a/cmd/internal/peer/kcp_client.go
+++ b/cmd/internal/peer/kcp_client.go
@@ -0,0 +1,184 @@
+package peer
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"sync/atomic"
+
+	"omnisocketgo/cmd/internal/latencylog"
+	"omnisocketgo/cmd/internal/protocol"
+	"omnisocketgo/cmd/internal/transport"
+)
+
+// KCPClient 表示一个通过 KCP 连接到 server 的 peer。
+type KCPClient struct {
+	id     string
+	conn   *transport.KCPConn
+	logger latencylog.Logger
+
+	nextID uint64
+}
+
+// DialKCP 通过 KCP 连接到 server，并发送 register 消息完成身份注册。
+func DialKCP(serverAddr, peerID string, opts ...Option) (*KCPClient, error) {
+	options := clientOptions{
+		logger: latencylog.NoopLogger{},
+	}
+	for _, opt := range opts {
+		opt(&options)
+	}
+	if options.logger == nil {
+		options.logger = latencylog.NoopLogger{}
+	}
+
+	session, err := transport.DialKCPSession(
+		serverAddr,
+		options.bindIP,
+		options.bindDevice,
+		options.kcpPacketDebugLogger,
+		latencylog.NodeRolePeer,
+		peerID,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("peer: dial kcp server: %w", err)
+	}
+
+	conn, err := transport.NewKCPConn(
+		session,
+		transport.WithKCPLogger(options.logger, latencylog.NodeRolePeer, peerID),
+	)
+	if err != nil {
+		_ = session.Close()
+		return nil, fmt.Errorf("peer: create kcp transport conn: %w", err)
+	}
+
+	client := &KCPClient{
+		id:     peerID,
+		conn:   conn,
+		logger: options.logger,
+	}
+
+	if err := conn.Send(protocol.Message{
+		Type: protocol.MessageTypeRegister,
+		From: peerID,
+		To:   protocol.ServerPeerID,
+	}); err != nil {
+		_ = conn.Close()
+		return nil, fmt.Errorf("peer: register with kcp server: %w", err)
+	}
+
+	return client, nil
+}
+
+// ID 返回当前 client 的 peer 标识。
+func (c *KCPClient) ID() string {
+	return c.id
+}
+
+// SendText 向目标 peer 发送一条文本消息。
+func (c *KCPClient) SendText(to, body string) error {
+	msg := protocol.Message{
+		Type: protocol.MessageTypeText,
+		ID:   c.nextMessageID(),
+		From: c.id,
+		To:   to,
+	}
+	latencylog.LogMessageEvent(c.logger, latencylog.NodeRolePeer, c.id, latencylog.EventAAppPrepBegin, msg)
+	msg.Body = []byte(body)
+	return c.conn.Send(msg)
+}
+
+// SendFile 向目标 peer 发送一条文件消息。
+func (c *KCPClient) SendFile(to, fileName string, body []byte) error {
+	msg := protocol.Message{
+		Type:     protocol.MessageTypeFile,
+		ID:       c.nextMessageID(),
+		From:     c.id,
+		To:       to,
+		FileName: fileName,
+	}
+	latencylog.LogMessageEvent(c.logger, latencylog.NodeRolePeer, c.id, latencylog.EventAAppPrepBegin, msg)
+
+	bodyCopy := make([]byte, len(body))
+	copy(bodyCopy, body)
+	msg.Body = bodyCopy
+
+	return c.conn.Send(msg)
+}
+
+// SendFilePath 从本地文件读取内容并发送给目标 peer。
+func (c *KCPClient) SendFilePath(to, path string) error {
+	msg := protocol.Message{
+		Type:     protocol.MessageTypeFile,
+		ID:       c.nextMessageID(),
+		From:     c.id,
+		To:       to,
+		FileName: filepath.Base(path),
+	}
+	latencylog.LogMessageEvent(c.logger, latencylog.NodeRolePeer, c.id, latencylog.EventAAppPrepBegin, msg)
+
+	body, err := os.ReadFile(path)
+	if err != nil {
+		return fmt.Errorf("peer: read file %s: %w", path, err)
+	}
+	msg.Body = body
+
+	return c.conn.Send(msg)
+}
+
+// Receive 读取一条来自 server 的消息。
+func (c *KCPClient) Receive() (protocol.Message, error) {
+	msg, err := c.conn.Receive()
+	if err != nil {
+		return protocol.Message{}, fmt.Errorf("peer: receive from kcp server: %w", err)
+	}
+
+	latencylog.LogMessageEvent(c.logger, latencylog.NodeRolePeer, c.id, latencylog.EventBAppRecv, msg)
+	return msg, nil
+}
+
+// ReceiveLoop 持续接收 server 消息并交给 handler 处理。
+func (c *KCPClient) ReceiveLoop(handler func(protocol.Message) error) error {
+	return c.conn.ReceiveLoop(func(msg protocol.Message) error {
+		switch msg.Type {
+		case protocol.MessageTypeText, protocol.MessageTypeFile, protocol.MessageTypeError:
+			latencylog.LogMessageEvent(c.logger, latencylog.NodeRolePeer, c.id, latencylog.EventBAppRecv, msg)
+			return handler(msg)
+		default:
+			return fmt.Errorf("peer: unexpected message type from kcp server: %s", msg.Type)
+		}
+	})
+}
+
+// PersistMessage 将收到的业务消息写入本地磁盘。
+func (c *KCPClient) PersistMessage(msg protocol.Message, inboxDir string) (string, error) {
+	if !latencylog.IsBusinessMessage(msg) {
+		return "", fmt.Errorf("peer: cannot persist message type %s", msg.Type)
+	}
+	if inboxDir == "" {
+		return "", fmt.Errorf("peer: inbox directory is required")
+	}
+
+	latencylog.LogMessageEvent(c.logger, latencylog.NodeRolePeer, c.id, latencylog.EventBPersistBegin, msg)
+	if err := os.MkdirAll(inboxDir, 0o755); err != nil {
+		return "", fmt.Errorf("peer: create inbox dir %s: %w", inboxDir, err)
+	}
+
+	path, err := persistMessageToDisk(msg, inboxDir)
+	if err != nil {
+		return "", err
+	}
+
+	latencylog.LogMessageEvent(c.logger, latencylog.NodeRolePeer, c.id, latencylog.EventBPersistEnd, msg)
+	return path, nil
+}
+
+// Close 关闭与 server 的 KCP 会话。
+func (c *KCPClient) Close() error {
+	return c.conn.Close()
+}
+
+func (c *KCPClient) nextMessageID() uint64 {
+	return atomic.AddUint64(&c.nextID, 1)
+}
--- a/cmd/internal/peer/kcp_client_test.go
+++ b/cmd/internal/peer/kcp_client_test.go
@@ -0,0 +1,263 @@
+package peer
+
+import (
+	"os"
+	"path/filepath"
+	"reflect"
+	"strings"
+	"sync"
+	"testing"
+
+	kcp "github.com/xtaci/kcp-go/v5"
+
+	"omnisocketgo/cmd/internal/latencylog"
+	"omnisocketgo/cmd/internal/protocol"
+	"omnisocketgo/cmd/internal/server"
+	"omnisocketgo/cmd/internal/transport"
+)
+
+func TestKCPDialRegistersPeer(t *testing.T) {
+	hub := server.NewKCPHub()
+	serverAddr, cleanup := startRealKCPHubServer(t, hub)
+	defer cleanup()
+
+	client, err := DialKCP(serverAddr, "peer-a")
+	if err != nil {
+		t.Fatalf("DialKCP() error = %v", err)
+	}
+	defer func() { _ = client.Close() }()
+
+	waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered")
+}
+
+func TestKCPDialRejectsInvalidBindIP(t *testing.T) {
+	_, err := DialKCP("127.0.0.1:9002", "peer-a", WithBindIP("not-an-ip"))
+	if err == nil {
+		t.Fatal("DialKCP() error = nil, want invalid bind ip error")
+	}
+	if !strings.Contains(err.Error(), `invalid bind ip "not-an-ip"`) {
+		t.Fatalf("DialKCP() error = %v, want invalid bind ip error", err)
+	}
+}
+
+func TestKCPClientsExchangeTextAndFileMessages(t *testing.T) {
+	hub := server.NewKCPHub()
+	serverAddr, cleanup := startRealKCPHubServer(t, hub)
+	defer cleanup()
+
+	peerA, err := DialKCP(serverAddr, "peer-a")
+	if err != nil {
+		t.Fatalf("DialKCP(peer-a) error = %v", err)
+	}
+	defer func() { _ = peerA.Close() }()
+
+	peerB, err := DialKCP(serverAddr, "peer-b")
+	if err != nil {
+		t.Fatalf("DialKCP(peer-b) 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")
+
+	received := make(chan protocol.Message, 2)
+	receiveErr := make(chan error, 1)
+	go func() {
+		for i := 0; i < 2; i++ {
+			msg, err := peerB.Receive()
+			if err != nil {
+				receiveErr <- err
+				return
+			}
+			received <- msg
+		}
+		receiveErr <- nil
+	}()
+
+	if err := peerA.SendText("peer-b", "hello over kcp"); err != nil {
+		t.Fatalf("SendText() error = %v", err)
+	}
+	fileBody := []byte{0x01, 0x02, 0x03}
+	if err := peerA.SendFile("peer-b", "payload.bin", fileBody); err != nil {
+		t.Fatalf("SendFile() error = %v", err)
+	}
+
+	if err := <-receiveErr; err != nil {
+		t.Fatalf("peerB.Receive() error = %v", err)
+	}
+
+	gotFirst := <-received
+	wantFirst := protocol.Message{
+		Type: protocol.MessageTypeText,
+		ID:   1,
+		From: "peer-a",
+		To:   "peer-b",
+		Body: []byte("hello over kcp"),
+	}
+	if !reflect.DeepEqual(gotFirst, wantFirst) {
+		t.Fatalf("first message mismatch: got %+v want %+v", gotFirst, wantFirst)
+	}
+
+	gotSecond := <-received
+	wantSecond := protocol.Message{
+		Type:     protocol.MessageTypeFile,
+		ID:       2,
+		From:     "peer-a",
+		To:       "peer-b",
+		FileName: "payload.bin",
+		Body:     fileBody,
+	}
+	if !reflect.DeepEqual(gotSecond, wantSecond) {
+		t.Fatalf("second message mismatch: got %+v want %+v", gotSecond, wantSecond)
+	}
+}
+
+func TestKCPClientReceivesServerErrorForUnknownTarget(t *testing.T) {
+	hub := server.NewKCPHub()
+	serverAddr, cleanup := startRealKCPHubServer(t, hub)
+	defer cleanup()
+
+	client, err := DialKCP(serverAddr, "peer-a")
+	if err != nil {
+		t.Fatalf("DialKCP() error = %v", err)
+	}
+	defer func() { _ = client.Close() }()
+
+	waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered")
+
+	if err := client.SendText("missing-peer", "hello"); err != nil {
+		t.Fatalf("SendText() error = %v", err)
+	}
+
+	got, err := client.Receive()
+	if err != nil {
+		t.Fatalf("Receive() error = %v", err)
+	}
+	if got.Type != protocol.MessageTypeError {
+		t.Fatalf("got type %s, want %s", got.Type, protocol.MessageTypeError)
+	}
+	if string(got.Body) != "unknown target: missing-peer" {
+		t.Fatalf("error body = %q, want unknown target message", got.Body)
+	}
+}
+
+func TestKCPClientsExchangeMessagesWithLatencyLogs(t *testing.T) {
+	hub := server.NewKCPHub()
+	serverAddr, cleanup := startRealKCPHubServer(t, hub)
+	defer cleanup()
+
+	peerALogger := &recordingLogger{}
+	peerA, err := DialKCP(serverAddr, "peer-a", WithLogger(peerALogger))
+	if err != nil {
+		t.Fatalf("DialKCP(peer-a) error = %v", err)
+	}
+	defer func() { _ = peerA.Close() }()
+
+	peerBLogger := &recordingLogger{}
+	peerB, err := DialKCP(serverAddr, "peer-b", WithLogger(peerBLogger))
+	if err != nil {
+		t.Fatalf("DialKCP(peer-b) error = %v", err)
+	}
+	defer func() { _ = peerB.Close() }()
+
+	inboxDir := t.TempDir()
+
+	waitFor(t, func() bool { return hub.HasPeer("peer-a") && hub.HasPeer("peer-b") }, "both peers to be registered")
+
+	if err := peerA.SendText("peer-b", "hello"); err != nil {
+		t.Fatalf("SendText() error = %v", err)
+	}
+	textMsg, err := peerB.Receive()
+	if err != nil {
+		t.Fatalf("peerB.Receive(text) error = %v", err)
+	}
+	if _, err := peerB.PersistMessage(textMsg, inboxDir); err != nil {
+		t.Fatalf("peerB.PersistMessage(text) error = %v", err)
+	}
+
+	filePath := filepath.Join(t.TempDir(), "payload.bin")
+	if err := os.WriteFile(filePath, []byte{0x01, 0x02, 0x03}, 0o644); err != nil {
+		t.Fatalf("os.WriteFile() error = %v", err)
+	}
+	if err := peerA.SendFilePath("peer-b", filePath); err != nil {
+		t.Fatalf("SendFilePath() error = %v", err)
+	}
+	fileMsg, err := peerB.Receive()
+	if err != nil {
+		t.Fatalf("peerB.Receive(file) error = %v", err)
+	}
+	if _, err := peerB.PersistMessage(fileMsg, inboxDir); err != nil {
+		t.Fatalf("peerB.PersistMessage(file) error = %v", err)
+	}
+
+	waitFor(t, func() bool { return len(peerALogger.Events()) == 6 }, "peer-a latency events")
+	waitFor(t, func() bool { return len(peerBLogger.Events()) == 6 }, "peer-b latency events")
+
+	assertEventSequencesByMessage(t, peerALogger.Events(), map[uint64][]string{
+		1: {latencylog.EventAAppPrepBegin, latencylog.EventSendHandoffBegin, latencylog.EventSendHandoffEnd},
+		2: {latencylog.EventAAppPrepBegin, latencylog.EventSendHandoffBegin, latencylog.EventSendHandoffEnd},
+	})
+	assertEventSequencesByMessage(t, peerBLogger.Events(), map[uint64][]string{
+		1: {latencylog.EventBAppRecv, latencylog.EventBPersistBegin, latencylog.EventBPersistEnd},
+		2: {latencylog.EventBAppRecv, latencylog.EventBPersistBegin, latencylog.EventBPersistEnd},
+	})
+}
+
+func startRealKCPHubServer(t *testing.T, hub *server.KCPHub) (string, func()) {
+	t.Helper()
+
+	listener, packetConn, err := transport.ListenKCPSessions("127.0.0.1:0", "", nil, latencylog.NodeRoleServer, "hub")
+	if err != nil {
+		t.Fatalf("ListenKCPSessions() error = %v", err)
+	}
+
+	var (
+		wg   sync.WaitGroup
+		stop = make(chan struct{})
+	)
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for {
+			session, acceptErr := listener.AcceptKCP()
+			if acceptErr != nil {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				if strings.Contains(acceptErr.Error(), "closed") {
+					return
+				}
+				t.Errorf("AcceptKCP() error = %v", acceptErr)
+				return
+			}
+
+			wg.Add(1)
+			go func(sess *kcp.UDPSession) {
+				defer wg.Done()
+				if serveErr := hub.ServeSession(sess); serveErr != nil && !isExpectedKCPHubServeExit(serveErr) {
+					t.Logf("hub.ServeSession() ended with %v", serveErr)
+				}
+			}(session)
+		}
+	}()
+
+	cleanup := func() {
+		close(stop)
+		_ = listener.Close()
+		_ = packetConn.Close()
+		wg.Wait()
+	}
+
+	return listener.Addr().String(), cleanup
+}
+
+func isExpectedKCPHubServeExit(err error) bool {
+	if err == nil {
+		return true
+	}
+
+	message := err.Error()
+	return strings.Contains(message, "closed") || strings.Contains(message, "broken pipe") || strings.Contains(message, "io: read/write on closed pipe")
+}