realrasengan
bb08f71e30
Wire-compatible with C mosh. AES-128-OCB3, SSP transport, fragmentation, protobuf state sync, latch extension fields. Includes ServeRW bridge mode for embedding.
395 lines
9.4 KiB
Go
395 lines
9.4 KiB
Go
package mosh
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/zlib"
|
|
"crypto/rand"
|
|
"encoding/binary"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func newTestPair(t *testing.T) (*Transport, *Transport) {
|
|
t.Helper()
|
|
key := make([]byte, 16)
|
|
rand.Read(key)
|
|
ocbS, err := NewOCB(key)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ocbC, err := NewOCB(key)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
server := NewTransport(ocbS, true)
|
|
client := NewTransport(ocbC, false)
|
|
return server, client
|
|
}
|
|
|
|
func TestTransportBasicExchange(t *testing.T) {
|
|
server, client := newTestPair(t)
|
|
|
|
// Server sends a diff.
|
|
server.SetPending([]byte("hello from server"))
|
|
datagrams := server.Tick()
|
|
if len(datagrams) == 0 {
|
|
t.Fatal("no datagrams")
|
|
}
|
|
|
|
// Client receives.
|
|
var diff []byte
|
|
for _, dg := range datagrams {
|
|
if d := client.Recv(dg); d != nil {
|
|
diff = d
|
|
}
|
|
}
|
|
if !bytes.Equal(diff, []byte("hello from server")) {
|
|
t.Fatalf("diff = %q", diff)
|
|
}
|
|
|
|
// Client sends a diff.
|
|
client.SetPending([]byte("hello from client"))
|
|
datagrams = client.Tick()
|
|
if len(datagrams) == 0 {
|
|
t.Fatal("no datagrams")
|
|
}
|
|
|
|
// Server receives.
|
|
diff = nil
|
|
for _, dg := range datagrams {
|
|
if d := server.Recv(dg); d != nil {
|
|
diff = d
|
|
}
|
|
}
|
|
if !bytes.Equal(diff, []byte("hello from client")) {
|
|
t.Fatalf("diff = %q", diff)
|
|
}
|
|
}
|
|
|
|
func TestTransportReplayRejected(t *testing.T) {
|
|
server, client := newTestPair(t)
|
|
|
|
server.SetPending([]byte("data"))
|
|
datagrams := server.Tick()
|
|
|
|
// First receive works.
|
|
if d := client.Recv(datagrams[0]); d == nil {
|
|
t.Fatal("first receive failed")
|
|
}
|
|
|
|
// Replay should be rejected.
|
|
if d := client.Recv(datagrams[0]); d != nil {
|
|
t.Fatal("replay should be rejected")
|
|
}
|
|
}
|
|
|
|
func TestTransportWrongDirection(t *testing.T) {
|
|
server, _ := newTestPair(t)
|
|
|
|
server.SetPending([]byte("data"))
|
|
datagrams := server.Tick()
|
|
|
|
// Server receiving its own datagram (wrong direction).
|
|
if d := server.Recv(datagrams[0]); d != nil {
|
|
t.Fatal("wrong direction should be rejected")
|
|
}
|
|
}
|
|
|
|
func TestTransportLargePayload(t *testing.T) {
|
|
server, client := newTestPair(t)
|
|
|
|
// Payload larger than one fragment.
|
|
payload := make([]byte, maxFragmentPayload*3+42)
|
|
rand.Read(payload)
|
|
|
|
server.SetPending(payload)
|
|
datagrams := server.Tick()
|
|
if len(datagrams) < 2 {
|
|
t.Fatalf("expected multiple datagrams, got %d", len(datagrams))
|
|
}
|
|
|
|
var diff []byte
|
|
for _, dg := range datagrams {
|
|
if d := client.Recv(dg); d != nil {
|
|
diff = d
|
|
}
|
|
}
|
|
if !bytes.Equal(diff, payload) {
|
|
t.Fatal("large payload mismatch")
|
|
}
|
|
}
|
|
|
|
func TestTransportHeartbeat(t *testing.T) {
|
|
server, _ := newTestPair(t)
|
|
|
|
// Force lastSend to be old enough to trigger heartbeat.
|
|
server.mu.Lock()
|
|
server.lastSend = time.Now().Add(-2 * server.rto)
|
|
server.mu.Unlock()
|
|
|
|
datagrams := server.Tick()
|
|
if len(datagrams) == 0 {
|
|
t.Fatal("expected heartbeat datagram")
|
|
}
|
|
}
|
|
|
|
func TestTransportEmptyTick(t *testing.T) {
|
|
server, _ := newTestPair(t)
|
|
|
|
// No pending diff, recent send, no ack needed — should produce nothing.
|
|
datagrams := server.Tick()
|
|
if len(datagrams) != 0 {
|
|
t.Fatalf("expected no datagrams, got %d", len(datagrams))
|
|
}
|
|
}
|
|
|
|
func TestTransportRTOBounds(t *testing.T) {
|
|
server, _ := newTestPair(t)
|
|
|
|
rto := server.RTO()
|
|
if rto < minRTO || rto > maxRTO {
|
|
t.Fatalf("RTO = %v, not in [%v, %v]", rto, minRTO, maxRTO)
|
|
}
|
|
}
|
|
|
|
func TestTransportMultipleExchanges(t *testing.T) {
|
|
server, client := newTestPair(t)
|
|
|
|
for i := 0; i < 10; i++ {
|
|
payload := make([]byte, 100)
|
|
rand.Read(payload)
|
|
|
|
server.SetPending(payload)
|
|
for _, dg := range server.Tick() {
|
|
client.Recv(dg)
|
|
}
|
|
|
|
client.SetPending(payload)
|
|
for _, dg := range client.Tick() {
|
|
server.Recv(dg)
|
|
}
|
|
}
|
|
|
|
// Verify sequence numbers advanced.
|
|
server.mu.Lock()
|
|
if server.sentNum < 10 {
|
|
t.Fatalf("sentNum = %d", server.sentNum)
|
|
}
|
|
server.mu.Unlock()
|
|
}
|
|
|
|
// TestTransportCompressionBomb verifies that a zlib-compressed payload
|
|
// expanding to >1 MiB is rejected without crashing.
|
|
func TestTransportCompressionBomb(t *testing.T) {
|
|
key := make([]byte, 16)
|
|
rand.Read(key)
|
|
ocbS, err := NewOCB(key)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ocbC, err := NewOCB(key)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
server := NewTransport(ocbS, true)
|
|
client := NewTransport(ocbC, false)
|
|
|
|
// Create a zlib payload that expands to >1 MiB (zeros compress very well).
|
|
bigPayload := make([]byte, 2<<20) // 2 MiB of zeros
|
|
var zbuf bytes.Buffer
|
|
w := zlib.NewWriter(&zbuf)
|
|
w.Write(bigPayload)
|
|
w.Close()
|
|
compressed := zbuf.Bytes()
|
|
|
|
// Wrap in a single fragment (final=true, fragment_num=0, id=1).
|
|
var fragWire []byte
|
|
fragWire = make([]byte, fragmentHeaderSize+len(compressed))
|
|
binary.BigEndian.PutUint64(fragWire[0:], 1) // id
|
|
binary.BigEndian.PutUint16(fragWire[8:], fragmentFinalBit|uint16(0)) // final, frag 0
|
|
copy(fragWire[fragmentHeaderSize:], compressed)
|
|
|
|
// Encrypt as a server->client datagram.
|
|
server.mu.Lock()
|
|
server.seqOut++
|
|
seq := server.seqOut
|
|
dirSeq := dirToClient | (seq & seqMask)
|
|
var dirSeqBytes [8]byte
|
|
binary.BigEndian.PutUint64(dirSeqBytes[:], dirSeq)
|
|
var nonce [12]byte
|
|
copy(nonce[4:], dirSeqBytes[:])
|
|
|
|
ts := uint16(time.Now().UnixMilli() & 0xffff)
|
|
plaintext := make([]byte, 4+len(fragWire))
|
|
binary.BigEndian.PutUint16(plaintext[0:], ts)
|
|
binary.BigEndian.PutUint16(plaintext[2:], 0)
|
|
copy(plaintext[4:], fragWire)
|
|
server.mu.Unlock()
|
|
|
|
tagAndCT := ocbS.Encrypt(nonce[:], plaintext)
|
|
wire := make([]byte, 8+len(tagAndCT))
|
|
copy(wire[:8], dirSeqBytes[:])
|
|
copy(wire[8:], tagAndCT)
|
|
|
|
// Client should reject the bomb — Recv returns nil.
|
|
result := client.Recv(wire)
|
|
if result != nil {
|
|
t.Fatalf("compression bomb should be rejected, got %d bytes", len(result))
|
|
}
|
|
}
|
|
|
|
// TestTransportConcurrentSendRecv launches multiple goroutines sending
|
|
// from the client and receiving on the server simultaneously to verify
|
|
// there are no data races under -race.
|
|
func TestTransportConcurrentSendRecv(t *testing.T) {
|
|
server, client := newTestPair(t)
|
|
|
|
const numGoroutines = 10
|
|
errs := make(chan error, numGoroutines*2)
|
|
|
|
// Launch sender goroutines — each sends a unique payload.
|
|
for i := 0; i < numGoroutines; i++ {
|
|
go func(id int) {
|
|
payload := make([]byte, 64)
|
|
// Fill with unique pattern so corruption is detectable.
|
|
for j := range payload {
|
|
payload[j] = byte(id)
|
|
}
|
|
client.SetPending(payload)
|
|
datagrams := client.Tick()
|
|
if len(datagrams) == 0 {
|
|
errs <- nil // no datagram is fine if another goroutine won the race
|
|
return
|
|
}
|
|
errs <- nil
|
|
}(i)
|
|
}
|
|
|
|
// Launch receiver goroutines — each tries to receive datagrams.
|
|
for i := 0; i < numGoroutines; i++ {
|
|
go func(id int) {
|
|
// Generate a datagram from the client side to receive on server.
|
|
payload := make([]byte, 32)
|
|
for j := range payload {
|
|
payload[j] = byte(id + 100)
|
|
}
|
|
client.SetPending(payload)
|
|
datagrams := client.Tick()
|
|
for _, dg := range datagrams {
|
|
server.Recv(dg)
|
|
}
|
|
errs <- nil
|
|
}(i)
|
|
}
|
|
|
|
// Collect all errors.
|
|
for i := 0; i < numGoroutines*2; i++ {
|
|
if err := <-errs; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestTransportCapsNegotiation(t *testing.T) {
|
|
server, client := newTestPair(t)
|
|
|
|
server.SetCaps([]byte{CapSessionControl})
|
|
client.SetCaps([]byte{CapSessionControl})
|
|
|
|
// Exchange a message so caps are transmitted.
|
|
server.SetPending([]byte("hello"))
|
|
for _, dg := range server.Tick() {
|
|
client.Recv(dg)
|
|
}
|
|
client.SetPending([]byte("world"))
|
|
for _, dg := range client.Tick() {
|
|
server.Recv(dg)
|
|
}
|
|
|
|
// Both should see remote caps now.
|
|
if rc := client.RemoteCaps(); len(rc) == 0 {
|
|
t.Fatal("client has no remote caps")
|
|
}
|
|
if rc := server.RemoteCaps(); len(rc) == 0 {
|
|
t.Fatal("server has no remote caps")
|
|
}
|
|
|
|
// Both should agree on the capability.
|
|
if !server.HasCap(CapSessionControl) {
|
|
t.Fatal("server should have CapSessionControl")
|
|
}
|
|
if !client.HasCap(CapSessionControl) {
|
|
t.Fatal("client should have CapSessionControl")
|
|
}
|
|
|
|
// Test intersection: server has bit 0x03, client has 0x01 → intersect = 0x01.
|
|
server.SetCaps([]byte{0x03})
|
|
client.SetCaps([]byte{0x01})
|
|
|
|
server.SetPending([]byte("a"))
|
|
for _, dg := range server.Tick() {
|
|
client.Recv(dg)
|
|
}
|
|
client.SetPending([]byte("b"))
|
|
for _, dg := range client.Tick() {
|
|
server.Recv(dg)
|
|
}
|
|
|
|
if !server.HasCap(0x01) {
|
|
t.Fatal("server should have bit 0x01")
|
|
}
|
|
if server.HasCap(0x02) {
|
|
t.Fatal("server should not have bit 0x02 (client lacks it)")
|
|
}
|
|
}
|
|
|
|
func TestTransportNoCaps(t *testing.T) {
|
|
server, client := newTestPair(t)
|
|
|
|
// No caps set — HasCap should return false.
|
|
if server.HasCap(CapSessionControl) {
|
|
t.Fatal("should be false with no caps")
|
|
}
|
|
if client.HasCap(CapSessionControl) {
|
|
t.Fatal("should be false with no caps")
|
|
}
|
|
|
|
// Exchange without caps.
|
|
server.SetPending([]byte("data"))
|
|
for _, dg := range server.Tick() {
|
|
client.Recv(dg)
|
|
}
|
|
|
|
// RemoteCaps should still be nil.
|
|
if rc := client.RemoteCaps(); rc != nil {
|
|
t.Fatalf("expected nil remote caps, got %x", rc)
|
|
}
|
|
}
|
|
|
|
// TestTransportHighSequenceNumbers verifies that the transport works
|
|
// correctly when sequence numbers are very large (near 2^62).
|
|
func TestTransportHighSequenceNumbers(t *testing.T) {
|
|
server, client := newTestPair(t)
|
|
|
|
// Set server's outgoing sequence counter to near 2^62.
|
|
server.mu.Lock()
|
|
server.seqOut = (1 << 62) - 1
|
|
server.mu.Unlock()
|
|
|
|
server.SetPending([]byte("high seq test"))
|
|
datagrams := server.Tick()
|
|
if len(datagrams) == 0 {
|
|
t.Fatal("no datagrams produced")
|
|
}
|
|
|
|
var diff []byte
|
|
for _, dg := range datagrams {
|
|
if d := client.Recv(dg); d != nil {
|
|
diff = d
|
|
}
|
|
}
|
|
if !bytes.Equal(diff, []byte("high seq test")) {
|
|
t.Fatalf("diff = %q, want %q", diff, "high seq test")
|
|
}
|
|
}
|