ExploreDNS/web/api/handler_internal_test.go

package api

import (
	"context"
	"encoding/json"
	"net"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"
)

// TestStartTraversalJobCap verifies that new traversals are rejected with
// 429 once maxRunning jobs are in flight. The store is pre-filled with
// running jobs so no real traversal is spawned.
func TestStartTraversalJobCap(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	h := newHandler(ctx)
	h.maxRunning = 2

	for i := 0; i < 2; i++ {
		h.st.set(&TraversalJob{ID: strconv.Itoa(i), Status: statusRunning, StartedAt: time.Now()})
	}

	req := httptest.NewRequest("POST", "/api/traverse",
		strings.NewReader(`{"domain":"example.com"}`))
	rec := httptest.NewRecorder()
	h.mux.ServeHTTP(rec, req)

	if rec.Code != 429 {
		t.Fatalf("want 429 with %d running jobs, got %d: %s", h.st.runningCount(), rec.Code, rec.Body.String())
	}

	// A finished job frees a slot: runningCount must drop below the cap.
	done := time.Now()
	if j, ok := h.st.get("0"); ok {
		j.mu.Lock()
		j.Status = statusComplete
		j.DoneAt = &done
		j.mu.Unlock()
	}
	if got := h.st.runningCount(); got != 1 {
		t.Fatalf("runningCount after completion = %d, want 1", got)
	}
}

// TestJobTimeoutReachesTerminalState verifies that a traversal launched
// with an already-expired deadline still drives the job to a terminal
// error state (the TTL cleanup only ever purges finished jobs).
func TestJobTimeoutReachesTerminalState(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond)
	defer cancel()
	<-ctx.Done() // deadline already exceeded

	h := &Handler{st: newStore(), jobTimeout: time.Nanosecond}
	job := &TraversalJob{ID: "t", Status: statusRunning, StartedAt: time.Now(), cancel: cancel}
	h.st.set(job)

	doneCh := make(chan struct{})
	go func() {
		h.runTraversal(ctx, job, "example.com", 1, false)
		close(doneCh)
	}()
	select {
	case <-doneCh:
	case <-time.After(30 * time.Second):
		t.Fatal("runTraversal did not return with an expired context")
	}

	job.mu.RLock()
	defer job.mu.RUnlock()
	if job.DoneAt == nil {
		t.Fatal("job never reached a terminal state")
	}
	if job.Status != statusError || !strings.Contains(job.Error, "timed out") {
		t.Fatalf("want error status mentioning timeout, got status=%q error=%q", job.Status, job.Error)
	}
}

// TestSubscribeSnapshot_NoDuplicates regresses the subscribe/snapshot race:
// subscribers arriving while events are being published must never see the
// same event twice (once from the snapshot replay and once from the channel).
// Events are numbered, so any duplicate breaks strict monotonicity. Slow
// subscribers may legitimately drop events (publishLocked is non-blocking),
// so gaps are not an error.
func TestSubscribeSnapshot_NoDuplicates(t *testing.T) {
	job := &TraversalJob{Status: statusRunning}

	const total = 2000
	const subscribers = 8

	var pub sync.WaitGroup
	pub.Add(1)
	go func() {
		defer pub.Done()
		for i := 0; i < total; i++ {
			ev := ProgressEvent{RefID: strconv.Itoa(i)}
			job.mu.Lock()
			job.Progress = append(job.Progress, ev)
			job.publishLocked(ev)
			job.mu.Unlock()
		}
		job.mu.Lock()
		job.Status = statusComplete
		job.mu.Unlock()
		job.closeSubscribers()
	}()

	var subs sync.WaitGroup
	errs := make(chan string, subscribers)
	for s := 0; s < subscribers; s++ {
		subs.Add(1)
		go func() {
			defer subs.Done()
			sub, past, done := job.subscribeSnapshot()
			defer job.unsubscribe(sub)

			last := -1
			check := func(refid string) {
				n, err := strconv.Atoi(refid)
				if err != nil {
					errs <- "bad refid " + refid
					return
				}
				if n <= last {
					errs <- "event " + refid + " out of order or duplicated after " + strconv.Itoa(last)
					return
				}
				last = n
			}
			for _, ev := range past {
				check(ev.RefID)
			}
			if !done {
				for ev := range sub {
					check(ev.RefID)
				}
			}
		}()
	}

	pub.Wait()
	subs.Wait()
	close(errs)
	for msg := range errs {
		t.Error(msg)
	}
}

// fakeVersionQuerier returns canned version strings without touching the
// network.
type fakeVersionQuerier struct {
	mu       sync.Mutex
	versions map[string]string
	queried  []string
}

func (f *fakeVersionQuerier) Query(_ context.Context, ip net.IP) string {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.queried = append(f.queried, ip.String())
	return f.versions[ip.String()]
}

// TestGetServers_PendingWhileRunning verifies the 202 pending shape while a
// job has not finished fingerprinting (running or just-completed).
func TestGetServers_PendingWhileRunning(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	h := newHandler(ctx)

	now := time.Now()
	jobs := []*TraversalJob{
		{ID: "running", Status: statusRunning, StartedAt: now},
		{ID: "fingerprinting", Status: statusComplete, StartedAt: now, DoneAt: &now},
	}
	for _, j := range jobs {
		h.st.set(j)
	}

	for _, id := range []string{"running", "fingerprinting"} {
		req := httptest.NewRequest("GET", "/api/traverse/"+id+"/servers", nil)
		rec := httptest.NewRecorder()
		h.mux.ServeHTTP(rec, req)

		if rec.Code != 202 {
			t.Fatalf("%s: want 202, got %d: %s", id, rec.Code, rec.Body.String())
		}
		var body map[string]string
		if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
			t.Fatalf("%s: decode: %v", id, err)
		}
		if body["status"] != "pending" {
			t.Fatalf("%s: want status=pending, got %q", id, body["status"])
		}
	}
}

// TestFingerprintServers_StoresServersAndPublishes drives the fingerprint
// step with a fake querier: (server, ip) pairs become sorted job.Servers with
// versions, pseudo/invalid entries are skipped, a {"stage":"servers"} event
// is published, and the endpoint flips from pending to the final list.
func TestFingerprintServers_StoresServersAndPublishes(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	h := newHandler(ctx)
	fake := &fakeVersionQuerier{versions: map[string]string{
		"192.0.2.1": "TestDNS 1.0",
		"192.0.2.2": "",
	}}
	h.fp = fake

	now := time.Now()
	job := &TraversalJob{ID: "j", Status: statusComplete, StartedAt: now, DoneAt: &now}
	h.st.set(job)
	sub, _, _ := job.subscribeSnapshot()
	defer job.unsubscribe(sub)

	h.fingerprintServers(job, map[string][]string{
		"b.example.net": {"192.0.2.2"},
		"a.example.net": {"192.0.2.1", "key:pseudo:entry", "not-an-ip"},
	})

	want := []ServerInfo{
		{Name: "a.example.net", IP: "192.0.2.1", Version: "TestDNS 1.0"},
		{Name: "b.example.net", IP: "192.0.2.2", Version: ""},
	}
	job.mu.RLock()
	got := append([]ServerInfo(nil), job.Servers...)
	done := job.serversDone
	job.mu.RUnlock()
	if !done {
		t.Fatal("serversDone not set")
	}
	if len(got) != len(want) {
		t.Fatalf("servers = %+v, want %+v", got, want)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("servers[%d] = %+v, want %+v", i, got[i], want[i])
		}
	}

	select {
	case ev := <-sub:
		if ev.Stage != "servers" {
			t.Fatalf("published stage = %q, want servers", ev.Stage)
		}
	default:
		t.Fatal("no servers event published")
	}

	// Endpoint now serves the final list.
	req := httptest.NewRequest("GET", "/api/traverse/j/servers", nil)
	rec := httptest.NewRecorder()
	h.mux.ServeHTTP(rec, req)
	if rec.Code != 200 {
		t.Fatalf("want 200, got %d: %s", rec.Code, rec.Body.String())
	}
	var body struct {
		Status  string       `json:"status"`
		Servers []ServerInfo `json:"servers"`
	}
	if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
		t.Fatal(err)
	}
	if body.Status != "complete" || len(body.Servers) != 2 {
		t.Fatalf("body = %+v", body)
	}

	// Snapshot includes servers too.
	req = httptest.NewRequest("GET", "/api/traverse/j", nil)
	rec = httptest.NewRecorder()
	h.mux.ServeHTTP(rec, req)
	if rec.Code != 200 {
		t.Fatalf("snapshot: want 200, got %d", rec.Code)
	}
	var snap struct {
		Servers []ServerInfo `json:"servers"`
	}
	if err := json.NewDecoder(rec.Body).Decode(&snap); err != nil {
		t.Fatal(err)
	}
	if len(snap.Servers) != 2 {
		t.Fatalf("snapshot servers = %+v, want 2 entries", snap.Servers)
	}
}

// TestFingerprintServers_EmptySeen still terminates the pending state and
// publishes the servers event for traversals that recorded no servers.
func TestFingerprintServers_EmptySeen(t *testing.T) {
	h := &Handler{st: newStore()}
	job := &TraversalJob{ID: "e", Status: statusError}
	h.st.set(job)
	sub, _, _ := job.subscribeSnapshot()
	defer job.unsubscribe(sub)

	h.fingerprintServers(job, nil)

	job.mu.RLock()
	defer job.mu.RUnlock()
	if !job.serversDone {
		t.Fatal("serversDone not set")
	}
	if len(job.Servers) != 0 {
		t.Fatalf("servers = %+v, want empty", job.Servers)
	}
	select {
	case ev := <-sub:
		if ev.Stage != "servers" {
			t.Fatalf("published stage = %q, want servers", ev.Stage)
		}
	default:
		t.Fatal("no servers event published")
	}
}