pangolin/server/internal/scheduler/probe/ingest_test.go

package probe_test

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"

	"github.com/alicebob/miniredis/v2"
	"github.com/redis/go-redis/v9"
	"github.com/wangjia/pangolin/server/internal/scheduler/probe"
)

// --------------------------------------------------------------------------
// Test helpers
// --------------------------------------------------------------------------

const (
	testProbeID = "probe-test-01"
	testSecret  = "s3cr3t-for-testing"
)

// newTestSetup creates a fresh miniredis instance, a Store, and an
// IngestHandler pre-configured with testProbeID/testSecret.
// The miniredis server is automatically closed when the test ends.
func newTestSetup(t *testing.T) (*probe.IngestHandler, *probe.Store, *miniredis.Miniredis) {
	t.Helper()
	mr := miniredis.RunT(t)
	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	st := probe.NewStore(rdb)
	reg := probe.NewMapRegistry(map[string]string{testProbeID: testSecret})
	h := probe.NewIngestHandler(reg, st)
	return h, st, mr
}

// sign computes the HMAC-SHA256 hex digest expected by IngestHandler.
// Message: ts + "\n" + rawBody
func sign(secret, tsStr string, body []byte) string {
	mac := hmac.New(sha256.New, []byte(secret))
	mac.Write([]byte(tsStr))
	mac.Write([]byte("\n"))
	mac.Write(body)
	return hex.EncodeToString(mac.Sum(nil))
}

// sampleRequest builds a minimal valid ReportRequest.
func sampleRequest() probe.ReportRequest {
	return probe.ReportRequest{
		ProbeID: testProbeID,
		Vantage: probe.VantagePoint{
			ISP:     "China Telecom",
			Region:  "Guangdong",
			Country: "CN",
		},
		Reports: []probe.NodeReport{
			{
				NodeID: "node-hk-01",
				Ts:     time.Now().Unix(),
				L1:     probe.L1Result{OK: true, RttMs: 42},
				L2: &probe.L2Result{
					OK:   true,
					ALPN: "h3",
				},
				L3: &probe.L3Result{OK: true},
			},
		},
	}
}

// buildHTTPRequest encodes req as JSON and attaches the three required auth
// headers with the current timestamp and a valid HMAC.
func buildHTTPRequest(t *testing.T, req probe.ReportRequest, overrideTS ...string) *http.Request {
	t.Helper()
	body, err := json.Marshal(req)
	if err != nil {
		t.Fatalf("marshal request: %v", err)
	}
	tsStr := strconv.FormatInt(time.Now().Unix(), 10)
	if len(overrideTS) > 0 {
		tsStr = overrideTS[0]
	}
	r := httptest.NewRequest(http.MethodPost, "/probe/report", bytes.NewReader(body))
	r.Header.Set("Content-Type", "application/json")
	r.Header.Set("X-Probe-Id", testProbeID)
	r.Header.Set("X-Probe-Ts", tsStr)
	r.Header.Set("X-Probe-Sign", sign(testSecret, tsStr, body))
	return r
}

// --------------------------------------------------------------------------
// Auth failure tests (no Redis touch needed, but we still pass a valid setup)
// --------------------------------------------------------------------------

// TestIngestBadSignature verifies that a wrong HMAC returns 401.
func TestIngestBadSignature(t *testing.T) {
	h, _, _ := newTestSetup(t)

	req := sampleRequest()
	body, _ := json.Marshal(req)
	tsStr := strconv.FormatInt(time.Now().Unix(), 10)

	r := httptest.NewRequest(http.MethodPost, "/probe/report", bytes.NewReader(body))
	r.Header.Set("X-Probe-Id", testProbeID)
	r.Header.Set("X-Probe-Ts", tsStr)
	// Deliberately wrong secret.
	r.Header.Set("X-Probe-Sign", sign("wrong-secret", tsStr, body))

	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)

	if w.Code != http.StatusUnauthorized {
		t.Errorf("bad signature: got %d, want 401", w.Code)
	}
}

// TestIngestMissingHeaders verifies that missing any of the three auth headers
// results in 401.
func TestIngestMissingHeaders(t *testing.T) {
	h, _, _ := newTestSetup(t)
	req := sampleRequest()
	body, _ := json.Marshal(req)
	tsStr := strconv.FormatInt(time.Now().Unix(), 10)
	validSign := sign(testSecret, tsStr, body)

	cases := []struct {
		name   string
		id     string
		ts     string
		sigHdr string
	}{
		{"no X-Probe-Id", "", tsStr, validSign},
		{"no X-Probe-Ts", testProbeID, "", validSign},
		{"no X-Probe-Sign", testProbeID, tsStr, ""},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			r := httptest.NewRequest(http.MethodPost, "/probe/report", bytes.NewReader(body))
			if tc.id != "" {
				r.Header.Set("X-Probe-Id", tc.id)
			}
			if tc.ts != "" {
				r.Header.Set("X-Probe-Ts", tc.ts)
			}
			if tc.sigHdr != "" {
				r.Header.Set("X-Probe-Sign", tc.sigHdr)
			}
			w := httptest.NewRecorder()
			h.ServeHTTP(w, r)
			if w.Code != http.StatusUnauthorized {
				t.Errorf("%s: got %d, want 401", tc.name, w.Code)
			}
		})
	}
}

// TestIngestUnknownProbeID verifies that an unregistered probe ID returns 401.
func TestIngestUnknownProbeID(t *testing.T) {
	h, _, _ := newTestSetup(t)
	req := sampleRequest()
	body, _ := json.Marshal(req)
	tsStr := strconv.FormatInt(time.Now().Unix(), 10)

	r := httptest.NewRequest(http.MethodPost, "/probe/report", bytes.NewReader(body))
	r.Header.Set("X-Probe-Id", "not-registered")
	r.Header.Set("X-Probe-Ts", tsStr)
	r.Header.Set("X-Probe-Sign", sign(testSecret, tsStr, body))

	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)

	if w.Code != http.StatusUnauthorized {
		t.Errorf("unknown probe: got %d, want 401", w.Code)
	}
}

// TestIngestTimestampExpired verifies that a stale timestamp (> 300s old)
// returns 401.
func TestIngestTimestampExpired(t *testing.T) {
	h, _, _ := newTestSetup(t)
	req := sampleRequest()
	body, _ := json.Marshal(req)

	// 10 minutes in the past – well outside the ±300 s window.
	staleTS := strconv.FormatInt(time.Now().Add(-10*time.Minute).Unix(), 10)

	r := httptest.NewRequest(http.MethodPost, "/probe/report", bytes.NewReader(body))
	r.Header.Set("X-Probe-Id", testProbeID)
	r.Header.Set("X-Probe-Ts", staleTS)
	r.Header.Set("X-Probe-Sign", sign(testSecret, staleTS, body))

	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)

	if w.Code != http.StatusUnauthorized {
		t.Errorf("stale ts: got %d, want 401", w.Code)
	}
}

// TestIngestTimestampFuture verifies that a far-future timestamp is also
// rejected (outside the ±300 s window).
func TestIngestTimestampFuture(t *testing.T) {
	h, _, _ := newTestSetup(t)
	req := sampleRequest()
	body, _ := json.Marshal(req)

	futureTS := strconv.FormatInt(time.Now().Add(10*time.Minute).Unix(), 10)

	r := httptest.NewRequest(http.MethodPost, "/probe/report", bytes.NewReader(body))
	r.Header.Set("X-Probe-Id", testProbeID)
	r.Header.Set("X-Probe-Ts", futureTS)
	r.Header.Set("X-Probe-Sign", sign(testSecret, futureTS, body))

	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)

	if w.Code != http.StatusUnauthorized {
		t.Errorf("future ts: got %d, want 401", w.Code)
	}
}

// TestIngestHMACConstantTime verifies that an HMAC of correct length but with
// one byte flipped is still rejected.  This guards against a non-constant-time
// comparison that might accept a zero-byte suffix match.
func TestIngestHMACConstantTime(t *testing.T) {
	h, _, _ := newTestSetup(t)
	req := sampleRequest()
	body, _ := json.Marshal(req)
	tsStr := strconv.FormatInt(time.Now().Unix(), 10)

	// Produce a valid MAC then flip the last byte.
	mac := hmac.New(sha256.New, []byte(testSecret))
	mac.Write([]byte(tsStr))
	mac.Write([]byte("\n"))
	mac.Write(body)
	correctMAC := mac.Sum(nil)
	correctMAC[len(correctMAC)-1] ^= 0xFF
	tamperedSign := hex.EncodeToString(correctMAC)

	r := httptest.NewRequest(http.MethodPost, "/probe/report", bytes.NewReader(body))
	r.Header.Set("X-Probe-Id", testProbeID)
	r.Header.Set("X-Probe-Ts", tsStr)
	r.Header.Set("X-Probe-Sign", tamperedSign)

	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)

	if w.Code != http.StatusUnauthorized {
		t.Errorf("tampered HMAC: got %d, want 401", w.Code)
	}
}

// --------------------------------------------------------------------------
// Integration tests – valid requests + Redis state
// --------------------------------------------------------------------------

// TestIngestValidReport is the full happy-path integration test.
// After a valid POST /probe/report the handler must:
//  1. Return 202 Accepted.
//  2. Write probe:{nodeID}:{vantage} in Redis with TTL ≈ 30 min.
//  3. Write probe:hb:{probeID} in Redis with TTL ≈ 15 min.
func TestIngestValidReport(t *testing.T) {
	h, _, mr := newTestSetup(t)

	req := sampleRequest()
	httpReq := buildHTTPRequest(t, req)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, httpReq)

	if w.Code != http.StatusAccepted {
		t.Fatalf("valid report: got %d, want 202; body=%s", w.Code, w.Body.String())
	}

	// --- Verify Redis snapshot key ---
	// Expected key: probe:node-hk-01:CN:Guangdong:China_Telecom
	snapKey := "probe:node-hk-01:CN:Guangdong:China_Telecom"
	val, err := mr.Get(snapKey)
	if err != nil {
		t.Fatalf("snapshot key %q not found in Redis: %v", snapKey, err)
	}

	var snap probe.ProbeSnapshot
	if err := json.Unmarshal([]byte(val), &snap); err != nil {
		t.Fatalf("unmarshal snapshot: %v", err)
	}
	if snap.ProbeID != testProbeID {
		t.Errorf("snapshot.ProbeID = %q, want %q", snap.ProbeID, testProbeID)
	}
	if snap.Report.NodeID != "node-hk-01" {
		t.Errorf("snapshot.Report.NodeID = %q, want node-hk-01", snap.Report.NodeID)
	}
	if !snap.Report.L1.OK {
		t.Error("snapshot L1.OK should be true")
	}
	if snap.ReceivedAt == 0 {
		t.Error("snapshot.ReceivedAt should be non-zero")
	}

	// TTL should be close to 30 minutes (1800s).  Allow ±5s for test latency.
	snapTTL := mr.TTL(snapKey)
	if snapTTL < 1795*time.Second || snapTTL > 1800*time.Second {
		t.Errorf("snapshot TTL = %v, want ≈30 min", snapTTL)
	}

	// --- Verify heartbeat key ---
	hbKey := "probe:hb:" + testProbeID
	hbVal, err := mr.Get(hbKey)
	if err != nil {
		t.Fatalf("heartbeat key %q not found in Redis: %v", hbKey, err)
	}
	if hbVal == "" {
		t.Error("heartbeat value should not be empty")
	}

	// TTL should be close to 15 minutes (900s).
	hbTTL := mr.TTL(hbKey)
	if hbTTL < 895*time.Second || hbTTL > 900*time.Second {
		t.Errorf("heartbeat TTL = %v, want ≈15 min", hbTTL)
	}
}

// TestIngestReplayIdempotent verifies that submitting the same (probeID, ts)
// a second time is treated as an idempotent replay:
//   - First request: 202, Redis written.
//   - Second request: 202, no error (no double-write).
func TestIngestReplayIdempotent(t *testing.T) {
	h, _, mr := newTestSetup(t)

	req := sampleRequest()
	body, _ := json.Marshal(req)
	tsStr := strconv.FormatInt(time.Now().Unix(), 10)

	buildReq := func() *http.Request {
		r := httptest.NewRequest(http.MethodPost, "/probe/report", bytes.NewReader(body))
		r.Header.Set("Content-Type", "application/json")
		r.Header.Set("X-Probe-Id", testProbeID)
		r.Header.Set("X-Probe-Ts", tsStr) // same ts both times
		r.Header.Set("X-Probe-Sign", sign(testSecret, tsStr, body))
		return r
	}

	// First request.
	w1 := httptest.NewRecorder()
	h.ServeHTTP(w1, buildReq())
	if w1.Code != http.StatusAccepted {
		t.Fatalf("first request: got %d, want 202", w1.Code)
	}

	// Snapshot key must exist after first request.
	snapKey := "probe:node-hk-01:CN:Guangdong:China_Telecom"
	if _, err := mr.Get(snapKey); err != nil {
		t.Fatalf("snapshot key missing after first request: %v", err)
	}

	// Second request (same ts – replay).
	w2 := httptest.NewRecorder()
	h.ServeHTTP(w2, buildReq())
	if w2.Code != http.StatusAccepted {
		t.Fatalf("replay request: got %d, want 202 (idempotent)", w2.Code)
	}

	// Replay prevention key should exist.
	seenKey := "probe:seen:" + testProbeID + ":" + tsStr
	if _, err := mr.Get(seenKey); err != nil {
		t.Errorf("seen key %q not found: %v", seenKey, err)
	}
}

// TestIngestMultipleReports verifies that a batch with multiple node reports
// writes one snapshot key per node.
func TestIngestMultipleReports(t *testing.T) {
	h, _, mr := newTestSetup(t)

	req := probe.ReportRequest{
		ProbeID: testProbeID,
		Vantage: probe.VantagePoint{Country: "CN", Region: "Shanghai", ISP: "ChinaNet"},
		Reports: []probe.NodeReport{
			{NodeID: "node-sg-01", Ts: time.Now().Unix(), L1: probe.L1Result{OK: true, RttMs: 80}},
			{NodeID: "node-jp-01", Ts: time.Now().Unix(), L1: probe.L1Result{OK: false}},
		},
	}
	httpReq := buildHTTPRequest(t, req)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, httpReq)

	if w.Code != http.StatusAccepted {
		t.Fatalf("multi-report: got %d, want 202", w.Code)
	}

	for _, nodeID := range []string{"node-sg-01", "node-jp-01"} {
		key := "probe:" + nodeID + ":CN:Shanghai:ChinaNet"
		if _, err := mr.Get(key); err != nil {
			t.Errorf("snapshot for %s not found: %v", nodeID, err)
		}
	}
}

// TestIngestL3Optional verifies that a report without L3 (third-party probe)
// is accepted normally.
func TestIngestL3Optional(t *testing.T) {
	h, _, _ := newTestSetup(t)

	req := probe.ReportRequest{
		ProbeID: testProbeID,
		Vantage: probe.VantagePoint{Country: "JP", Region: "Tokyo", ISP: "NTT"},
		Reports: []probe.NodeReport{
			{
				NodeID: "node-jp-02",
				Ts:     time.Now().Unix(),
				L1:     probe.L1Result{OK: true, RttMs: 15},
				L2:     &probe.L2Result{OK: true, ALPN: "h3"},
				// L3 intentionally absent.
			},
		},
	}
	httpReq := buildHTTPRequest(t, req)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, httpReq)

	if w.Code != http.StatusAccepted {
		t.Errorf("L3-absent report: got %d, want 202", w.Code)
	}
}

// --------------------------------------------------------------------------
// Store unit tests
// --------------------------------------------------------------------------

// TestStoreSnapshotsByNode verifies that SnapshotsByNode returns the correct
// map and that a missing key (no data) is represented as an empty map, not
// as an error or failure.
func TestStoreSnapshotsByNode(t *testing.T) {
	mr := miniredis.RunT(t)
	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	st := probe.NewStore(rdb)
	ctx := context.Background()

	// No data yet – must return empty map, not error.
	snaps, err := st.SnapshotsByNode(ctx, "node-xx-01")
	if err != nil {
		t.Fatalf("SnapshotsByNode with no data: unexpected error: %v", err)
	}
	if len(snaps) != 0 {
		t.Errorf("expected empty map, got %d entries", len(snaps))
	}

	// Write one snapshot via SaveReports.
	reports := []probe.NodeReport{
		{NodeID: "node-xx-01", Ts: time.Now().Unix(), L1: probe.L1Result{OK: true, RttMs: 10}},
	}
	if err := st.SaveReports(ctx, "probe-a", probe.VantagePoint{Country: "SG", Region: "Central", ISP: "Singtel"}, reports); err != nil {
		t.Fatalf("SaveReports: %v", err)
	}

	snaps, err = st.SnapshotsByNode(ctx, "node-xx-01")
	if err != nil {
		t.Fatalf("SnapshotsByNode after write: %v", err)
	}
	if len(snaps) != 1 {
		t.Fatalf("expected 1 snapshot, got %d", len(snaps))
	}
	// The vantage key in the map should be Country:Region:ISP.
	if _, ok := snaps["SG:Central:Singtel"]; !ok {
		t.Errorf("expected key SG:Central:Singtel in snapshots; got keys: %v", mapKeys(snaps))
	}
}

// TestStoreAliveProbes verifies that AliveProbes returns probes that have an
// active heartbeat and an empty slice (not an error) when none have reported.
func TestStoreAliveProbes(t *testing.T) {
	mr := miniredis.RunT(t)
	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	st := probe.NewStore(rdb)
	ctx := context.Background()

	// No heartbeats yet – empty list, not error.
	ids, err := st.AliveProbes(ctx)
	if err != nil {
		t.Fatalf("AliveProbes with no data: unexpected error: %v", err)
	}
	if len(ids) != 0 {
		t.Errorf("expected empty list, got %v", ids)
	}

	// Write a heartbeat for two probes.
	for _, pid := range []string{"probe-a", "probe-b"} {
		reports := []probe.NodeReport{
			{NodeID: "node-xx-01", Ts: time.Now().Unix(), L1: probe.L1Result{OK: true}},
		}
		if err := st.SaveReports(ctx, pid, probe.VantagePoint{Country: "HK"}, reports); err != nil {
			t.Fatalf("SaveReports(%s): %v", pid, err)
		}
	}

	ids, err = st.AliveProbes(ctx)
	if err != nil {
		t.Fatalf("AliveProbes after writes: %v", err)
	}
	if len(ids) != 2 {
		t.Errorf("expected 2 alive probes, got %d: %v", len(ids), ids)
	}
}

// TestStoreMissingHeartbeatIsNotFailure documents the "missing key = no data,
// not failure" contract for probe:hb:{probeID}.
func TestStoreMissingHeartbeatIsNotFailure(t *testing.T) {
	mr := miniredis.RunT(t)
	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	st := probe.NewStore(rdb)
	ctx := context.Background()

	// AliveProbes on an empty Redis must NOT return an error or non-nil slice.
	ids, err := st.AliveProbes(ctx)
	if err != nil {
		t.Fatalf("AliveProbes must not error on empty Redis: %v", err)
	}
	// An empty/nil slice is acceptable; a non-nil error is not.
	// Callers (15D) must treat this as "unknown", not "all probes down".
	_ = ids // zero-length is fine
}

// --------------------------------------------------------------------------
// Helper
// --------------------------------------------------------------------------

func mapKeys(m map[string]probe.ProbeSnapshot) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}

// TestIngestMethodNotAllowed verifies GET returns 405.
func TestIngestMethodNotAllowed(t *testing.T) {
	h, _, _ := newTestSetup(t)
	r := httptest.NewRequest(http.MethodGet, "/probe/report", nil)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)
	if w.Code != http.StatusMethodNotAllowed {
		t.Errorf("GET: got %d, want 405", w.Code)
	}
}

// TestIngestEmptyReports verifies that a request with zero reports is rejected
// with 400.
func TestIngestEmptyReports(t *testing.T) {
	h, _, _ := newTestSetup(t)

	req := probe.ReportRequest{
		ProbeID: testProbeID,
		Vantage: probe.VantagePoint{Country: "CN", Region: "BJ", ISP: "Unicom"},
		Reports: []probe.NodeReport{}, // empty
	}
	httpReq := buildHTTPRequest(t, req)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, httpReq)

	if w.Code != http.StatusBadRequest {
		t.Errorf("empty reports: got %d, want 400", w.Code)
	}
}

// TestVantageKeyEncoding verifies that vantage values with spaces and colons
// are sanitized and produce a valid Redis key.
func TestVantageKeyEncoding(t *testing.T) {
	h, _, mr := newTestSetup(t)

	req := probe.ReportRequest{
		ProbeID: testProbeID,
		Vantage: probe.VantagePoint{
			ISP:     "China Telecom:CN2", // contains space and colon
			Region:  "Inner Mongolia",
			Country: "CN",
		},
		Reports: []probe.NodeReport{
			{NodeID: "node-bj-01", Ts: time.Now().Unix(), L1: probe.L1Result{OK: true}},
		},
	}
	httpReq := buildHTTPRequest(t, req)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, httpReq)

	if w.Code != http.StatusAccepted {
		t.Fatalf("vantage encoding: got %d, want 202", w.Code)
	}

	// Space → _ and colon → _ in the ISP field.
	expectedKey := "probe:node-bj-01:CN:Inner_Mongolia:China_Telecom_CN2"
	if _, err := mr.Get(expectedKey); err != nil {
		// List all keys to help diagnose.
		keys := mr.Keys()
		t.Errorf("expected key %q not found; existing keys: %v", expectedKey, keys)
	}
}

// TestIngestSignatureUsesRawBody verifies that the HMAC is computed over the
// exact raw bytes sent on the wire, not a re-serialised form.
func TestIngestSignatureUsesRawBody(t *testing.T) {
	h, _, _ := newTestSetup(t)

	// Build raw body with extra whitespace (valid JSON, different bytes).
	rawBody := []byte(`{  "probeId": "` + testProbeID + `",  "vantage":{"isp":"ISP1","region":"R1","country":"CN"},  "reports":[{"nodeId":"n1","ts":1234567890,"l1":{"ok":true,"rttMs":1}}]  }`)
	tsStr := strconv.FormatInt(time.Now().Unix(), 10)

	r := httptest.NewRequest(http.MethodPost, "/probe/report", bytes.NewReader(rawBody))
	r.Header.Set("Content-Type", "application/json")
	r.Header.Set("X-Probe-Id", testProbeID)
	r.Header.Set("X-Probe-Ts", tsStr)
	r.Header.Set("X-Probe-Sign", sign(testSecret, tsStr, rawBody))

	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)

	if w.Code != http.StatusAccepted {
		t.Errorf("raw body sign: got %d, want 202; body=%s", w.Code, w.Body.String())
	}
}