Prioritize VPN gateway control packets

agents/rap-node-agent/internal/vpnruntime/gateway_test.go

@@ -0,0 +1,92 @@
+package vpnruntime
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+)
+
+type recordingGatewayTransport struct {
+	mu      sync.Mutex
+	batches [][][]byte
+}
+
+func (t *recordingGatewayTransport) SendGatewayPacketBatch(ctx context.Context, packets [][]byte) error {
+	copied := make([][]byte, len(packets))
+	for i, packet := range packets {
+		copied[i] = append([]byte(nil), packet...)
+	}
+	t.mu.Lock()
+	t.batches = append(t.batches, copied)
+	t.mu.Unlock()
+	return nil
+}
+
+func (t *recordingGatewayTransport) ReceiveGatewayPacketBatch(ctx context.Context, timeout time.Duration) ([][]byte, error) {
+	return nil, ctx.Err()
+}
+
+func (t *recordingGatewayTransport) firstBatch() [][]byte {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	if len(t.batches) == 0 {
+		return nil
+	}
+	return t.batches[0]
+}
+
+func TestGatewayUploadPrioritizesTCPControlPackets(t *testing.T) {
+	transport := &recordingGatewayTransport{}
+	gateway := &Gateway{Transport: transport, VPNConnectionID: "vpn-1"}
+	priorityPackets := make(chan []byte, 1)
+	packets := make(chan []byte, 1)
+
+	normal := testIPv4TCPPacket([4]byte{101, 32, 118, 25}, [4]byte{10, 77, 0, 2}, 443, 37566)
+	priority := testIPv4TCPPacket([4]byte{192, 168, 200, 95}, [4]byte{10, 77, 0, 2}, 3389, 51000)
+	priority[33] = 0x12
+
+	packets <- normal
+	priorityPackets <- priority
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan error, 1)
+	go func() {
+		done <- gateway.uploadGatewayPackets(ctx, priorityPackets, packets)
+	}()
+	defer func() {
+		cancel()
+		<-done
+	}()
+
+	deadline := time.After(time.Second)
+	for {
+		if batch := transport.firstBatch(); len(batch) == 1 {
+			if string(batch[0]) != string(priority) {
+				t.Fatalf("first uploaded packet = %#v, want priority packet", batch[0])
+			}
+			return
+		}
+		select {
+		case <-deadline:
+			t.Fatal("timed out waiting for first gateway upload batch")
+		default:
+			time.Sleep(time.Millisecond)
+		}
+	}
+}
+
+func TestIsTCPControlPacket(t *testing.T) {
+	packet := testIPv4TCPPacket([4]byte{192, 168, 200, 95}, [4]byte{10, 77, 0, 2}, 3389, 51000)
+	if isTCPControlPacket(packet) {
+		t.Fatal("packet without control flags was classified as control")
+	}
+	packet[33] = 0x12
+	if !isTCPControlPacket(packet) {
+		t.Fatal("tcp syn-ack was not classified as control")
+	}
+	packet[9] = 17
+	if isTCPControlPacket(packet) {
+		t.Fatal("udp packet was classified as tcp control")
+	}
+}