DebTik/internal/bwtest/tcp_frames_test.go

package bwtest

import (
	"bufio"
	"bytes"
	"context"
	"encoding/binary"
	"encoding/json"
	"net"
	"sync"
	"testing"
	"time"
)

func TestConcurrentTCPFramesRemainDistinct(t *testing.T) {
	a, b := net.Pipe()
	defer a.Close()
	defer b.Close()
	w := tcpFrameWriter{conn: a}
	var wg sync.WaitGroup
	for _, kind := range []byte{tcpFrameData, tcpFrameControl} {
		wg.Add(1)
		go func(kind byte) {
			defer wg.Done()
			for range 20 {
				if _, err := w.write(kind, bytes.Repeat([]byte{kind}, 1024)); err != nil {
					return
				}
			}
		}(kind)
	}
	counts := map[byte]int{}
	buf := make([]byte, tcpFrameMax)
	_ = b.SetReadDeadline(time.Now().Add(3 * time.Second))
	for range 40 {
		kind, payload, err := readTCPFrame(b, buf)
		if err != nil {
			t.Fatal(err)
		}
		if !bytes.Equal(payload, bytes.Repeat([]byte{kind}, 1024)) {
			t.Fatal("interleaved control and data")
		}
		counts[kind]++
	}
	wg.Wait()
	if counts[tcpFrameData] != 20 || counts[tcpFrameControl] != 20 {
		t.Fatal(counts)
	}
}
func TestTCPFrameBounds(t *testing.T) {
	for _, n := range []uint32{0, tcpFrameMax + 1} {
		h := make([]byte, 5)
		h[0] = tcpFrameData
		binary.BigEndian.PutUint32(h[1:], n)
		if _, _, err := readTCPFrame(bytes.NewReader(h), make([]byte, tcpFrameMax)); err == nil {
			t.Fatal("invalid frame length accepted")
		}
	}
}
func testServer(t *testing.T) *BwTestServer {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	port := ln.Addr().(*net.TCPAddr).Port
	ln.Close()
	s := NewServer(ServerConfig{Enabled: true, ListenAddr: "127.0.0.1", Port: port})
	if err = s.Start(); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(s.Stop)
	return s
}
func TestBandwidthTCPAllDirections(t *testing.T) {
	s := testServer(t)
	for _, direction := range []string{"tx", "rx", "both"} {
		t.Run(direction, func(t *testing.T) {
			c := NewClient()
			t.Cleanup(c.Stop)
			err := c.Run(context.Background(), SessionRequest{Target: "127.0.0.1", Port: s.config.Port, Protocol: "tcp", Direction: direction, BandwidthLimit: 2_000_000, Duration: 2})
			if err != nil {
				t.Fatal(err)
			}
			deadline := time.Now().Add(5 * time.Second)
			for c.GetStatus().Running && time.Now().Before(deadline) {
				time.Sleep(10 * time.Millisecond)
			}
			st := c.GetStatus()
			if st.Running || st.Error != "" || st.Summary == nil {
				t.Fatalf("test did not finish: %+v", st)
			}
			if direction != "rx" && st.Summary.TotalTxBytes < 65536 {
				t.Fatalf("no transmitted data: %+v", st.Summary)
			}
			if direction != "tx" && st.Summary.TotalRxBytes < 65536 {
				t.Fatalf("no received data/reports: %+v", st.Summary)
			}
		})
	}
}
func TestBandwidthStopInterruptsThrottledWriter(t *testing.T) {
	s := testServer(t)
	c := NewClient()
	if err := c.Run(context.Background(), SessionRequest{Target: "127.0.0.1", Port: s.config.Port, Protocol: "tcp", Direction: "both", BandwidthLimit: 1000, Duration: 30}); err != nil {
		t.Fatal(err)
	}
	time.Sleep(150 * time.Millisecond)
	c.Stop()
	deadline := time.Now().Add(time.Second)
	for c.GetStatus().Running && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if c.GetStatus().Running {
		t.Fatal("Stop blocked behind rate limiter")
	}
}
func TestBandwidthRejectsUnframedTCP(t *testing.T) {
	s := testServer(t)
	conn, err := net.Dial("tcp", s.tcpListener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(time.Second))
	b, _ := json.Marshal(SessionHandshake{Type: "request", Protocol: "tcp"})
	conn.Write(append(b, '\n'))
	line, err := bufio.NewReader(conn).ReadBytes('\n')
	if err != nil {
		t.Fatal(err)
	}
	var response SessionHandshake
	if err = json.Unmarshal(line, &response); err != nil || response.Type != "reject" {
		t.Fatalf("legacy TCP silently accepted: %s %v", line, err)
	}
}
func TestRateLimiterCanSendFrameLargerThanBucket(t *testing.T) {
	limiter := NewRateLimiter(4_000_000)
	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer cancel()
	if !limiter.WaitContext(ctx, 65536) {
		t.Fatal("64KB frame can never fit nominal 50KB bucket")
	}
	ctx2, cancel2 := context.WithCancel(context.Background())
	cancel2()
	if limiter.WaitContext(ctx2, 65536) {
		t.Fatal("ignored cancellation")
	}
}