feat:KCP协议
This commit is contained in:
nnbcccscdscdsc
263
cmd/internal/peer/kcp_client_test.go
Normal file
263
cmd/internal/peer/kcp_client_test.go
Normal file
@@ -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")
|
||||
}
|
||||
Reference in New Issue
Block a user