Prevent stale labels, flaky websocket tests, and speed misreports

Bound Leo IP sends and stale response delivery by timeout, bootstrap speed-mode metadata lookups without JSON banner noise, respect --speed=false pre-dispatch, validate speed source binding, cap probe response reads, avoid caching empty nali labels, treat server-only speed metadata as healthy when no client IP is available, make websocket async tests event-driven, restart Leo receive parsing reliably, close websocket receive channels exactly once, send Leo requests on the waited websocket, back off failed Geo/RDNS metadata lookups, refresh MTR API info cache timestamps after lookup, report upload response read errors, avoid launching latency probes after cancellation, and keep speed JSON tests tolerant of degraded latency samples.

Author: tsosunchia

cmd/cmd.go

@@ -751,9 +751,13 @@ func applyDN42Mode(enabled bool, dataOrigin *string, disableMaptrace *bool) {
 	if !enabled {
 		return
 	}
+	applyDN42DataOrigin(dataOrigin)
+	*disableMaptrace = true
+}
+
+func applyDN42DataOrigin(dataOrigin *string) {
 	config.InitConfig()
 	*dataOrigin = "DN42"
-	*disableMaptrace = true
 }
 
 func prepareRuntimeEnvironment(ctx context.Context, dn42 bool, dataOrigin *string, disableMaptrace *bool, powProvider *string, asyncLeo bool) *wshandle.WsConn {
@@ -1254,8 +1258,8 @@ func Execute() {
 		return
 	}
 	if *speedMode {
-		// `--speed` should have been handled before the main parser runs. If we reached
-		// here, return a defensive error to avoid silently falling through.
+		// maybeRunSpeedMode should handle --speed before parser.Parse runs. This
+		// branch is a safety net for parser edge cases such as a "--" terminator.
 		fmt.Fprintln(os.Stderr, "internal error: speed mode dispatch failed")
 		os.Exit(1)
 	}

cmd/nali_mode.go

@@ -218,8 +218,9 @@ func runNaliMode(ctx context.Context, opts naliRunOptions) error {
 	restoreFastIPOutput := setFastIPOutputSuppression(true)
 	defer restoreFastIPOutput()
 
-	disableMaptrace := false
-	applyDN42Mode(opts.dn42, &opts.data, &disableMaptrace)
+	if opts.dn42 {
+		applyDN42DataOrigin(&opts.data)
+	}
 	leoWs := initLeoWebsocket(ctx, &opts.data, &opts.pow, false)
 	defer closeLeoWebsocket(leoWs)
 

cmd/nali_mode_test.go

@@ -92,6 +92,29 @@ func TestValidateNaliModeOptions(t *testing.T) {
 	}
 }
 
+func TestValidateNaliModeOptionsRejectsStartupConflictsFromInputs(t *testing.T) {
+	tests := []struct {
+		name  string
+		input naliModeOptionInputs
+		want  string
+	}{
+		{name: "deploy", input: naliModeOptionInputs{deploy: true}, want: "--deploy"},
+		{name: "from", input: naliModeOptionInputs{from: "Germany"}, want: "--from"},
+		{name: "init", input: naliModeOptionInputs{init: true}, want: "--init"},
+		{name: "listen", input: naliModeOptionInputs{listen: "127.0.0.1:8080"}, want: "--listen"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.input.parser = argparse.NewParser("nexttrace", "")
+			err := validateNaliModeOptions(buildNaliModeOptions(tt.input))
+			if err == nil || !strings.Contains(err.Error(), tt.want) {
+				t.Fatalf("validateNaliModeOptions() error = %v, want conflict %q", err, tt.want)
+			}
+		})
+	}
+}
+
 func TestBuildNaliModeOptionsDetectsExplicitDefaults(t *testing.T) {
 	parser := argparse.NewParser("nexttrace", "")
 	parser.Int("q", "queries", &argparse.Options{Default: 3})

cmd/speed_mode.go

@@ -16,8 +16,17 @@ import (
 	speedconfig "github.com/nxtrace/NTrace-core/internal/speedtest/config"
 	speedrender "github.com/nxtrace/NTrace-core/internal/speedtest/render"
 	speedrunner "github.com/nxtrace/NTrace-core/internal/speedtest/runner"
+	"github.com/nxtrace/NTrace-core/wshandle"
 )
 
+type speedMetadataBackend interface {
+	Close()
+}
+
+var newSpeedMetadataBackend = func(ctx context.Context) speedMetadataBackend {
+	return wshandle.NewWithContextAsync(ctx)
+}
+
 func registerSpeedFlag(parser *argparse.Parser) *bool {
 	return registerSpeedFlagWithAvailability(parser, enableSpeed)
 }
@@ -49,9 +58,12 @@ func containsSpeedFlag(args []string) bool {
 		if arg == "--" {
 			return false
 		}
-		if arg == "--speed" || strings.HasPrefix(arg, "--speed=") {
+		if arg == "--speed" {
 			return true
 		}
+		if v, ok := strings.CutPrefix(arg, "--speed="); ok {
+			return v == "" || strings.EqualFold(v, "true") || v == "1"
+		}
 	}
 	return false
 }
@@ -69,6 +81,10 @@ func runSpeedMode(rawArgs []string, stdout, stderr io.Writer) int {
 
 	rootCtx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
 	defer stop()
+	restoreFastIPOutput := suppressSpeedMetadataOutput(cfg)
+	defer restoreFastIPOutput()
+	metadataBackend := startSpeedMetadataBackend(rootCtx, cfg)
+	defer closeSpeedMetadataBackend(metadataBackend)
 
 	var bus *speedrender.Bus
 	isTTY := speedrender.IsTTY()
@@ -80,9 +96,9 @@ func runSpeedMode(rawArgs []string, stdout, stderr io.Writer) int {
 	} else {
 		bus = speedrender.NewBus(speedrender.NewPlainRenderer(stderr))
 	}
+	defer bus.Close()
 
 	res := speedrunner.Run(rootCtx, cfg, bus, isTTY)
-	bus.Close()
 	if cfg.OutputJSON {
 		enc := json.NewEncoder(stdout)
 		enc.SetEscapeHTML(false)
@@ -93,3 +109,23 @@ func runSpeedMode(rawArgs []string, stdout, stderr io.Writer) int {
 	}
 	return res.ExitCode
 }
+
+func suppressSpeedMetadataOutput(cfg *speedconfig.Config) func() {
+	if cfg == nil || cfg.NoMetadata || !cfg.OutputJSON {
+		return func() {}
+	}
+	return setFastIPOutputSuppression(true)
+}
+
+func startSpeedMetadataBackend(ctx context.Context, cfg *speedconfig.Config) speedMetadataBackend {
+	if cfg == nil || cfg.NoMetadata {
+		return nil
+	}
+	return newSpeedMetadataBackend(ctx)
+}
+
+func closeSpeedMetadataBackend(backend speedMetadataBackend) {
+	if backend != nil {
+		backend.Close()
+	}
+}

cmd/speed_mode_test.go

@@ -2,6 +2,8 @@ package cmd
 
 import (
 	"bytes"
+	"context"
+	"crypto/x509"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -13,10 +15,12 @@ import (
 
 	"github.com/akamensky/argparse"
 
+	speedconfig "github.com/nxtrace/NTrace-core/internal/speedtest/config"
 	"github.com/nxtrace/NTrace-core/internal/speedtest/netx"
 	"github.com/nxtrace/NTrace-core/internal/speedtest/provider/apple"
 	"github.com/nxtrace/NTrace-core/internal/speedtest/provider/cloudflare"
 	"github.com/nxtrace/NTrace-core/internal/speedtest/result"
+	"github.com/nxtrace/NTrace-core/util"
 )
 
 func TestRegisterSpeedFlagWithAvailabilityEnabledAddsSingleHelpEntry(t *testing.T) {
@@ -79,6 +83,70 @@ func TestRunSpeedModeRejectsUnexpectedTarget(t *testing.T) {
 	}
 }
 
+type fakeSpeedMetadataBackend struct {
+	closed bool
+}
+
+func (f *fakeSpeedMetadataBackend) Close() {
+	f.closed = true
+}
+
+func TestStartSpeedMetadataBackendHonorsNoMetadata(t *testing.T) {
+	calls := 0
+	prev := newSpeedMetadataBackend
+	newSpeedMetadataBackend = func(context.Context) speedMetadataBackend {
+		calls++
+		return &fakeSpeedMetadataBackend{}
+	}
+	defer func() { newSpeedMetadataBackend = prev }()
+
+	if backend := startSpeedMetadataBackend(context.Background(), &speedconfig.Config{NoMetadata: true}); backend != nil {
+		t.Fatalf("backend = %#v, want nil when metadata is disabled", backend)
+	}
+	if calls != 0 {
+		t.Fatalf("metadata backend calls = %d, want 0", calls)
+	}
+}
+
+func TestStartSpeedMetadataBackendClosesWhenEnabled(t *testing.T) {
+	fake := &fakeSpeedMetadataBackend{}
+	prev := newSpeedMetadataBackend
+	newSpeedMetadataBackend = func(context.Context) speedMetadataBackend {
+		return fake
+	}
+	defer func() { newSpeedMetadataBackend = prev }()
+
+	backend := startSpeedMetadataBackend(context.Background(), &speedconfig.Config{})
+	if backend != fake {
+		t.Fatalf("backend = %#v, want fake backend", backend)
+	}
+	closeSpeedMetadataBackend(backend)
+	if !fake.closed {
+		t.Fatal("metadata backend was not closed")
+	}
+}
+
+func TestSuppressSpeedMetadataOutputOnlyForJSONMetadata(t *testing.T) {
+	orig := util.SuppressFastIPOutput
+	defer func() { util.SuppressFastIPOutput = orig }()
+
+	util.SuppressFastIPOutput = false
+	restore := suppressSpeedMetadataOutput(&speedconfig.Config{OutputJSON: true})
+	if !util.SuppressFastIPOutput {
+		t.Fatal("SuppressFastIPOutput = false, want true for JSON metadata")
+	}
+	restore()
+	if util.SuppressFastIPOutput {
+		t.Fatal("SuppressFastIPOutput = true after restore, want false")
+	}
+
+	restore = suppressSpeedMetadataOutput(&speedconfig.Config{OutputJSON: true, NoMetadata: true})
+	if util.SuppressFastIPOutput {
+		t.Fatal("SuppressFastIPOutput = true, want unchanged when metadata is disabled")
+	}
+	restore()
+}
+
 func TestRunSpeedModeAppleJSON(t *testing.T) {
 	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
@@ -100,16 +168,22 @@ func TestRunSpeedModeAppleJSON(t *testing.T) {
 
 	restoreApple := apple.SetBaseForTest(srv.URL)
 	defer restoreApple()
-	restoreRoots := netx.SetExtraRootCAsForTest(srv.Client().Transport.(*http.Transport).TLSClientConfig.RootCAs)
+	restoreRoots := netx.SetExtraRootCAsForTest(testServerRootCAs(srv))
 	defer restoreRoots()
 
 	u, _ := url.Parse(srv.URL)
 	var stdout, stderr bytes.Buffer
 	rc := runSpeedMode([]string{"--speed", "--json", "--no-metadata", "--max", "64KiB", "--timeout", "1500", "--threads", "2", "--latency-count", "2", "--endpoint", u.Hostname()}, &stdout, &stderr)
 	if rc != 0 && rc != 2 {
-		t.Fatalf("runSpeedMode(apple) rc = %d, stderr=%s", rc, stderr.String())
+		t.Fatalf("runSpeedMode(apple) rc = %d, want 0 or degraded 2, stderr=%s", rc, stderr.String())
+	}
+	res := assertPureJSONSpeedResult(t, stdout.Bytes(), "apple")
+	if res.ExitCode != rc {
+		t.Fatalf("JSON exit_code = %d, want rc %d", res.ExitCode, rc)
+	}
+	if rc == 2 && !res.Degraded {
+		t.Fatal("Degraded = false, want true when rc=2")
 	}
-	assertPureJSONSpeedResult(t, stdout.Bytes(), "apple")
 }
 
 func TestRunSpeedModeCloudflareJSON(t *testing.T) {
@@ -136,18 +210,24 @@ func TestRunSpeedModeCloudflareJSON(t *testing.T) {
 
 	restoreCF := cloudflare.SetBaseForTest(srv.URL)
 	defer restoreCF()
-	restoreRoots := netx.SetExtraRootCAsForTest(srv.Client().Transport.(*http.Transport).TLSClientConfig.RootCAs)
+	restoreRoots := netx.SetExtraRootCAsForTest(testServerRootCAs(srv))
 	defer restoreRoots()
 
 	var stdout, stderr bytes.Buffer
 	rc := runSpeedMode([]string{"--speed", "--speed-provider", "cloudflare", "--json", "--no-metadata", "--max", "64KiB", "--timeout", "1500", "--threads", "2", "--latency-count", "2", "--non-interactive"}, &stdout, &stderr)
 	if rc != 0 && rc != 2 {
-		t.Fatalf("runSpeedMode(cloudflare) rc = %d, stderr=%s", rc, stderr.String())
+		t.Fatalf("runSpeedMode(cloudflare) rc = %d, want 0 or degraded 2, stderr=%s", rc, stderr.String())
+	}
+	res := assertPureJSONSpeedResult(t, stdout.Bytes(), "cloudflare")
+	if res.ExitCode != rc {
+		t.Fatalf("JSON exit_code = %d, want rc %d", res.ExitCode, rc)
+	}
+	if rc == 2 && !res.Degraded {
+		t.Fatal("Degraded = false, want true when rc=2")
 	}
-	assertPureJSONSpeedResult(t, stdout.Bytes(), "cloudflare")
 }
 
-func assertPureJSONSpeedResult(t *testing.T, data []byte, providerName string) {
+func assertPureJSONSpeedResult(t *testing.T, data []byte, providerName string) result.RunResult {
 	t.Helper()
 	if bytes.Contains(data, []byte("preferred API IP")) {
 		t.Fatalf("speed JSON output should not contain text noise:\n%s", data)
@@ -162,11 +242,28 @@ func assertPureJSONSpeedResult(t *testing.T, data []byte, providerName string) {
 	if len(res.Rounds) != 4 {
 		t.Fatalf("len(Rounds) = %d, want 4", len(res.Rounds))
 	}
+	return res
+}
+
+func testServerRootCAs(srv *httptest.Server) *x509.CertPool {
+	if transport, ok := srv.Client().Transport.(*http.Transport); ok && transport.TLSClientConfig != nil && transport.TLSClientConfig.RootCAs != nil {
+		return transport.TLSClientConfig.RootCAs
+	}
+	pool := x509.NewCertPool()
+	pool.AddCert(srv.Certificate())
+	return pool
 }
 
 func TestContainsSpeedFlagSupportsAssignedFormAndRespectsTerminator(t *testing.T) {
-	if !containsSpeedFlag([]string{"--speed=true"}) {
-		t.Fatal("containsSpeedFlag(--speed=true) = false, want true")
+	for _, arg := range []string{"--speed=true", "--speed=True", "--speed=1", "--speed="} {
+		if !containsSpeedFlag([]string{arg}) {
+			t.Fatalf("containsSpeedFlag(%s) = false, want true", arg)
+		}
+	}
+	for _, arg := range []string{"--speed=false", "--speed=0", "--speed=no"} {
+		if containsSpeedFlag([]string{arg}) {
+			t.Fatalf("containsSpeedFlag(%s) = true, want false", arg)
+		}
 	}
 	if containsSpeedFlag([]string{"--", "--speed"}) {
 		t.Fatal("containsSpeedFlag should ignore --speed after terminator")

internal/nali/nali.go

@@ -216,7 +216,7 @@ func (a *Annotator) lookupLabel(ctx context.Context, ip string) string {
 			shouldCache = true
 		}
 	}
-	if shouldCache {
+	if shouldCache && label != "" {
 		a.storeLabel(ip, label)
 	}
 	return label

internal/nali/nali_test.go

@@ -167,6 +167,30 @@ func TestAnnotateLineDoesNotCacheLookupErrors(t *testing.T) {
 	}
 }
 
+func TestAnnotateLineDoesNotCacheEmptyGeo(t *testing.T) {
+	calls := 0
+	a := New(Config{
+		Lang: "en",
+		Source: func(ip string, timeout time.Duration, lang string, maptrace bool) (*ipgeo.IPGeoData, error) {
+			calls++
+			if calls == 1 {
+				return &ipgeo.IPGeoData{}, nil
+			}
+			return &ipgeo.IPGeoData{CountryEn: "ok"}, nil
+		},
+	})
+
+	if got := a.AnnotateLine(context.Background(), "A 8.8.8.8"); got != "A 8.8.8.8" {
+		t.Fatalf("first AnnotateLine() = %q, want original", got)
+	}
+	if got := a.AnnotateLine(context.Background(), "A 8.8.8.8"); got != "A 8.8.8.8 [ok]" {
+		t.Fatalf("second AnnotateLine() = %q, want annotated", got)
+	}
+	if calls != 2 {
+		t.Fatalf("lookup calls = %d, want 2", calls)
+	}
+}
+
 func TestFindIPSpans(t *testing.T) {
 	line := "IP:1.1.1.1 [2001:db8::1]:443"
 	spans := FindIPSpans(line)