DragonTCP/core/cmd/dragontcp-client/wireselect_test.go

package main

import (
	"sync/atomic"
	"testing"
	"time"
)

func TestProfileCandidatesCoverBothWireFamilies(t *testing.T) {
	selector := &wireSelector{configured: WireAuto}
	candidates := selector.profileCandidates()
	binaryCount, bpCount, xorCount := 0, 0, 0
	seen := make(map[wireChoice]bool, len(candidates))
	firstBytes := make(map[byte]bool, 256)
	coveredMasks := map[string]map[byte]bool{WireBinary: {}, WireBP: {}, WireXOR: {}}
	coveredPadding := map[string]map[uint16]bool{WireBinary: {}, WireBP: {}, WireXOR: {}}
	for _, candidate := range candidates {
		if seen[candidate] {
			t.Fatalf("duplicate candidate: %s", candidate)
		}
		seen[candidate] = true
		if candidate.cover.Enabled {
			coveredMasks[candidate.mode][candidate.mask] = true
			coveredPadding[candidate.mode][candidate.cover.Padding] = true
		}
		switch candidate.mode {
		case WireBinary:
			binaryCount++
			if !candidate.cover.Enabled && candidate.mask&7 != 0 {
				t.Fatalf("ambiguous binary mask: %02x", candidate.mask)
			}
			if candidate.cover.Enabled {
				firstBytes[byte(candidate.cover.ID>>8)] = true
			} else {
				for mode := byte(0); mode <= 4; mode++ {
					firstBytes[mode^candidate.mask] = true
				}
			}
		case WireXOR:
			xorCount++
			if !candidate.cover.Enabled && ('U'^candidate.mask)&7 < 5 {
				t.Fatalf("ambiguous XOR mask: %02x", candidate.mask)
			}
			if candidate.cover.Enabled {
				firstBytes[byte(candidate.cover.ID>>8)] = true
			} else {
				firstBytes['U'^candidate.mask] = true
			}
		case WireBP:
			bpCount++
			if candidate.cover.Enabled && !candidate.cover.Clear {
				t.Fatalf("covered BP profile must use clear payloads: %s", candidate)
			}
			if !candidate.cover.Enabled && candidate.mask != 0 {
				t.Fatalf("direct BP profile must keep a clear header: %s", candidate)
			}
		default:
			t.Fatalf("unknown candidate: %s", candidate)
		}
	}
	if binaryCount != 544 || bpCount != 257 || xorCount != 352 {
		t.Fatalf("profiles B=%d BP=%d X=%d, want B=544 BP=257 X=352", binaryCount, bpCount, xorCount)
	}
	if len(firstBytes) != 256 {
		t.Fatalf("profiles cover %d first-byte values, want 256", len(firstBytes))
	}
	for _, mode := range []string{WireBinary, WireBP, WireXOR} {
		if len(coveredMasks[mode]) != 256 {
			t.Fatalf("mode %s covers %d masks, want 256", mode, len(coveredMasks[mode]))
		}
		if len(coveredPadding[mode]) != 16 {
			t.Fatalf("mode %s covers %d padding lengths, want 16", mode, len(coveredPadding[mode]))
		}
	}
}

func TestManualWireStillDiscoversAllProfilesForThatFamily(t *testing.T) {
	for _, tc := range []struct {
		mode string
		want int
	}{{WireBinary, 544}, {WireBP, 257}, {WireXOR, 352}} {
		selector := &wireSelector{configured: tc.mode}
		candidates := selector.profileCandidates()
		if len(candidates) != tc.want {
			t.Fatalf("mode %s profiles=%d, want %d", tc.mode, len(candidates), tc.want)
		}
		for _, candidate := range candidates {
			if candidate.mode != tc.mode {
				t.Fatalf("mode %s included %s", tc.mode, candidate)
			}
		}
	}
}

func TestClearProfilesAreTriedBeforeLegacyFallbacks(t *testing.T) {
	for _, mode := range []string{WireBinary, WireBP} {
		candidates := (&wireSelector{configured: mode}).profileCandidates()
		if len(candidates) < 2 || !candidates[0].cover.Clear {
			t.Fatalf("mode %s does not prefer a clear profile", mode)
		}
		if candidates[1].cover.Enabled {
			t.Fatalf("mode %s does not fall back immediately to a legacy direct profile", mode)
		}
	}
}

func measureDiscoveryConcurrency(t *testing.T, threads int) int32 {
	t.Helper()
	candidates := make([]wireChoice, 24)
	for i := range candidates {
		candidates[i] = wireChoice{mode: WireBinary, mask: byte(i * 8)}
	}
	var active atomic.Int32
	var maximum atomic.Int32
	var calls atomic.Int32
	selector := &wireSelector{
		configured:        WireAuto,
		candidateOverride: candidates,
		probeThreads:      threads,
		probeDelay:        time.Nanosecond,
		probeOverride: func(wireChoice) bool {
			current := active.Add(1)
			for {
				old := maximum.Load()
				if current <= old || maximum.CompareAndSwap(old, current) {
					break
				}
			}
			calls.Add(1)
			// Keep attempts alive long enough for the globally spaced scheduler
			// to fill every configured worker reliably on slower CI runners.
			time.Sleep(20 * time.Millisecond)
			active.Add(-1)
			return false
		},
	}

	if _, ok := selector.detectLocked(); ok {
		t.Fatal("unexpected working profile")
	}
	if calls.Load() != int32(len(candidates)) {
		t.Fatalf("screened=%d, want %d", calls.Load(), len(candidates))
	}
	return maximum.Load()
}

func TestDiscoveryDefaultsToOneWorker(t *testing.T) {
	if maximum := measureDiscoveryConcurrency(t, 0); maximum != 1 {
		t.Fatalf("maximum concurrent probes=%d, want 1", maximum)
	}
}

func TestDiscoveryHonorsConfiguredWorkers(t *testing.T) {
	if maximum := measureDiscoveryConcurrency(t, 4); maximum != 4 {
		t.Fatalf("maximum concurrent probes=%d, want 4", maximum)
	}
}

func TestFailedManualDiscoveryKeepsPinnedWire(t *testing.T) {
	for _, mode := range []string{WireBinary, WireBP, WireXOR} {
		selector := &wireSelector{
			configured:        mode,
			candidateOverride: []wireChoice{{mode: mode}},
			probeDelay:        time.Nanosecond,
			probeOverride:     func(wireChoice) bool { return false },
		}
		if choice := selector.mode(); choice.mode != mode {
			t.Fatalf("configured=%s fallback=%s", mode, choice.mode)
		}
	}
}