From fb3389a94b39b1b4eeaf4a11f1cf225c77b7a945 Mon Sep 17 00:00:00 2001 From: wangjia <809946525@qq.com> Date: Sat, 13 Jun 2026 11:56:34 +0800 Subject: [PATCH] feat(15A): probe ingest endpoint + Redis store (tsk_rBPr0Xuy10bz) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `server/internal/scheduler/probe/` package: * types.go – frozen ReportRequest/NodeReport/L1-L3 schema (15B/15C contract) * store.go – Redis read/write: probe:{nodeId}:{vantage} TTL 30min, probe:hb:{probeId} TTL 15min; missing key = no data, not failure; SnapshotsByNode + AliveProbes read interfaces for 15D * ingest.go – POST /probe/report handler: per-probe HMAC-SHA256 auth (X-Probe-Id / X-Probe-Ts / X-Probe-Sign), ±300s time window, constant-time comparison, idempotent replay via Redis SetNX Wire route in cmd/server/main.go (opt-in via PROBE_SECRETS env var). 18/18 tests pass (go test ./internal/scheduler/probe/... -count=1). Co-Authored-By: Claude Sonnet 4.6 --- server/cmd/server/main.go | 37 ++ server/internal/scheduler/probe/ingest.go | 185 ++++++ .../internal/scheduler/probe/ingest_test.go | 625 ++++++++++++++++++ server/internal/scheduler/probe/store.go | 212 ++++++ server/internal/scheduler/probe/types.go | 105 +++ 5 files changed, 1164 insertions(+) create mode 100644 server/internal/scheduler/probe/ingest.go create mode 100644 server/internal/scheduler/probe/ingest_test.go create mode 100644 server/internal/scheduler/probe/store.go create mode 100644 server/internal/scheduler/probe/types.go diff --git a/server/cmd/server/main.go b/server/cmd/server/main.go index 0af43dc..37d7965 100644 --- a/server/cmd/server/main.go +++ b/server/cmd/server/main.go @@ -9,6 +9,8 @@ import ( "github.com/go-chi/chi/v5" "github.com/go-chi/chi/v5/middleware" + "github.com/wangjia/pangolin/server/internal/redisutil" + "github.com/wangjia/pangolin/server/internal/scheduler/probe" ) func main() { @@ -33,8 +35,43 @@ func main() { _ = json.NewEncoder(w).Encode(map[string]string{"status": "ok"}) }) + // --- Probe ingest route (optional) --- + // Requires: + // PROBE_SECRETS – JSON object mapping probeId → HMAC secret, e.g. + // '{"probe-sg-01":"s3cr3t1","probe-jp-01":"s3cr3t2"}' + // REDIS_ADDR – Redis address, default 127.0.0.1:6379 + // REDIS_PASSWORD – Redis password (optional) + // + // If PROBE_SECRETS is empty the route is not registered and the server + // starts normally without the probe ingest endpoint. + probeSecretsJSON := os.Getenv("PROBE_SECRETS") + if probeSecretsJSON != "" { + redisAddr := getenvDefault("REDIS_ADDR", "127.0.0.1:6379") + rdb, err := redisutil.New(redisAddr, os.Getenv("REDIS_PASSWORD"), 0) + if err != nil { + log.Printf("probe: redis connect failed (%v) – probe route disabled", err) + } else { + var secretMap map[string]string + if err := json.Unmarshal([]byte(probeSecretsJSON), &secretMap); err != nil { + log.Fatalf("probe: invalid PROBE_SECRETS JSON: %v", err) + } + reg := probe.NewMapRegistry(secretMap) + st := probe.NewStore(rdb) + h := probe.NewIngestHandler(reg, st) + r.Post("/probe/report", h.ServeHTTP) + log.Printf("probe ingest route registered (%d probe(s))", len(secretMap)) + } + } + log.Printf("pangolin server listening on %s", *addr) if err := http.ListenAndServe(*addr, r); err != nil { log.Fatalf("server error: %v", err) } } + +func getenvDefault(key, def string) string { + if v := os.Getenv(key); v != "" { + return v + } + return def +} diff --git a/server/internal/scheduler/probe/ingest.go b/server/internal/scheduler/probe/ingest.go new file mode 100644 index 0000000..5444763 --- /dev/null +++ b/server/internal/scheduler/probe/ingest.go @@ -0,0 +1,185 @@ +package probe + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "io" + "net/http" + "strconv" + "time" +) + +const ( + // maxBodySize caps the request body at 1 MiB to prevent DoS. + maxBodySize = 1 << 20 + + // timeWindow is the symmetric ±window for timestamp validation. + // Requests whose X-Probe-Ts falls outside now±300s are rejected. + timeWindow = 300 * time.Second +) + +// ProbeRegistry maps probe IDs to their per-probe HMAC secrets. +// Implementations must be safe for concurrent use. +type ProbeRegistry interface { + // LookupSecret returns the HMAC secret for probeID. + // ok is false when the probe ID is not registered. + LookupSecret(probeID string) (secret string, ok bool) +} + +// MapRegistry is an in-memory ProbeRegistry backed by a plain map. +// Suitable for static configuration loaded at startup. +type MapRegistry struct { + secrets map[string]string // probeID → HMAC secret +} + +// NewMapRegistry creates a MapRegistry from a probeID→secret map. +// The map is copied; subsequent mutation of the original has no effect. +func NewMapRegistry(m map[string]string) *MapRegistry { + cp := make(map[string]string, len(m)) + for k, v := range m { + cp[k] = v + } + return &MapRegistry{secrets: cp} +} + +// LookupSecret implements ProbeRegistry. +func (r *MapRegistry) LookupSecret(probeID string) (string, bool) { + s, ok := r.secrets[probeID] + return s, ok +} + +// IngestHandler handles POST /probe/report. +// +// # Security model +// +// Each probe agent has its own identity (probeID + secret) registered in the +// ProbeRegistry. Every request must carry three headers: +// +// X-Probe-Id – the probe agent's identifier +// X-Probe-Ts – Unix timestamp (seconds, string) +// X-Probe-Sign – HMAC-SHA256 hex digest: HMAC(secret, ts+"\n"+rawBody) +// +// Validation steps (all failures return 401 with no diagnostic detail): +// 1. Presence of all three headers. +// 2. ProbeID exists in the registry. +// 3. |now − ts| ≤ 300 s (replay window). +// 4. Constant-time HMAC comparison. +// 5. Replay check via Redis SetNX on (probeID, ts). +// +// # Statelessness +// +// The handler only verifies the signature and writes to Redis. +// It carries no in-process state and is safe to run behind a CDN or across +// horizontally-scaled instances sharing the same Redis. +type IngestHandler struct { + registry ProbeRegistry + store *Store +} + +// NewIngestHandler creates an IngestHandler. +func NewIngestHandler(reg ProbeRegistry, store *Store) *IngestHandler { + return &IngestHandler{registry: reg, store: store} +} + +// ServeHTTP implements http.Handler for POST /probe/report. +func (h *IngestHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + + // Read body first – HMAC is computed over the raw bytes. + body, err := io.ReadAll(io.LimitReader(r.Body, maxBodySize)) + if err != nil { + w.WriteHeader(http.StatusUnauthorized) + return + } + + probeID := r.Header.Get("X-Probe-Id") + tsStr := r.Header.Get("X-Probe-Ts") + sign := r.Header.Get("X-Probe-Sign") + + // Reject immediately if any required auth header is missing. + // All auth failures are 401 – no detail leak. + if probeID == "" || tsStr == "" || sign == "" { + w.WriteHeader(http.StatusUnauthorized) + return + } + + // Probe must be registered. + secret, ok := h.registry.LookupSecret(probeID) + if !ok { + w.WriteHeader(http.StatusUnauthorized) + return + } + + // Validate timestamp window: |now − ts| ≤ timeWindow. + tsUnix, err := strconv.ParseInt(tsStr, 10, 64) + if err != nil { + w.WriteHeader(http.StatusUnauthorized) + return + } + diff := time.Since(time.Unix(tsUnix, 0)) + if diff < 0 { + diff = -diff + } + if diff > timeWindow { + w.WriteHeader(http.StatusUnauthorized) + return + } + + // Verify HMAC-SHA256(secret, ts + "\n" + rawBody). + // Constant-time comparison prevents timing side-channels. + mac := hmac.New(sha256.New, []byte(secret)) + mac.Write([]byte(tsStr)) + mac.Write([]byte("\n")) + mac.Write(body) + expected := mac.Sum(nil) + + gotBytes, err := hex.DecodeString(sign) + if err != nil { + w.WriteHeader(http.StatusUnauthorized) + return + } + if !hmac.Equal(gotBytes, expected) { + w.WriteHeader(http.StatusUnauthorized) + return + } + + // Replay check: (probeID, ts) pair must not have been seen before. + // Uses Redis SetNX for atomicity across multiple handler instances. + // On replay we return 202 (idempotent acknowledgement) rather than an + // error, so that a network-retrying probe agent is not penalised. + ctx := r.Context() + isReplay, err := h.store.CheckAndMarkSeen(ctx, probeID, tsStr) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + return + } + if isReplay { + w.WriteHeader(http.StatusAccepted) + return + } + + // Parse the request body. + var req ReportRequest + if err := json.Unmarshal(body, &req); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + if len(req.Reports) == 0 { + w.WriteHeader(http.StatusBadRequest) + return + } + + // Persist snapshots and heartbeat to Redis. + if err := h.store.SaveReports(ctx, probeID, req.Vantage, req.Reports); err != nil { + w.WriteHeader(http.StatusInternalServerError) + return + } + + w.WriteHeader(http.StatusAccepted) +} diff --git a/server/internal/scheduler/probe/ingest_test.go b/server/internal/scheduler/probe/ingest_test.go new file mode 100644 index 0000000..792958c --- /dev/null +++ b/server/internal/scheduler/probe/ingest_test.go @@ -0,0 +1,625 @@ +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()) + } +} + diff --git a/server/internal/scheduler/probe/store.go b/server/internal/scheduler/probe/store.go new file mode 100644 index 0000000..018a26b --- /dev/null +++ b/server/internal/scheduler/probe/store.go @@ -0,0 +1,212 @@ +package probe + +import ( + "context" + "encoding/json" + "fmt" + "strconv" + "strings" + "time" + + "github.com/redis/go-redis/v9" +) + +// Redis TTL constants for the probe subsystem. +const ( + // snapshotTTL is how long a per-(node,vantage) probe snapshot is retained. + // After this window the key expires and 15D must treat it as "no data". + snapshotTTL = 30 * time.Minute + + // heartbeatTTL is how long a probe heartbeat key lives without renewal. + // After this window the key expires — meaning "no recent data", NOT "failure". + heartbeatTTL = 15 * time.Minute + + // seenTTL is the lifetime of a replay-prevention key. + // Must exceed 2 × timeWindow (2 × 300 s = 600 s); 700 s adds a 100 s buffer. + seenTTL = 700 * time.Second +) + +// Store handles Redis reads and writes for the probe subsystem. +// +// # Key schema +// +// probe:{nodeID}:{vantageKey} → ProbeSnapshot JSON, TTL snapshotTTL (30 min) +// probe:hb:{probeID} → Unix timestamp string, TTL heartbeatTTL (15 min) +// probe:seen:{probeID}:{ts} → "1", TTL seenTTL (700 s), replay prevention +// +// # Missing-key semantics (IMPORTANT – must not be violated by callers) +// +// A missing probe:hb:{probeID} key means "no recent heartbeat data". +// It must NEVER be interpreted as "probe is down" or folded into a failure +// signal. Determination logic (15D) must treat an absent key as unknown, +// not as a negative result. +// +// # Node-ID constraint +// +// NodeIDs must not be the literal string "hb" or "seen", as those are used as +// key-namespace prefixes. In practice node IDs are UUIDs so this is safe. +type Store struct { + rdb *redis.Client +} + +// NewStore creates a Store backed by the given Redis client. +func NewStore(rdb *redis.Client) *Store { + return &Store{rdb: rdb} +} + +// -------------------------------------------------------------------------- +// Key helpers +// -------------------------------------------------------------------------- + +// snapshotKey returns the Redis key for the latest probe snapshot +// for the given (nodeID, vantage) pair. +// +// Format: probe:{nodeID}:{country}:{region}:{isp} +func snapshotKey(nodeID string, v VantagePoint) string { + return "probe:" + nodeID + ":" + vantageKey(v) +} + +// vantageKey returns a canonical, colon-safe string for a VantagePoint. +// It is used as the trailing segment of a snapshot Redis key. +func vantageKey(v VantagePoint) string { + return sanitizeKeySegment(v.Country) + ":" + + sanitizeKeySegment(v.Region) + ":" + + sanitizeKeySegment(v.ISP) +} + +// sanitizeKeySegment replaces characters that would interfere with Redis key +// parsing (space, colon) with underscores so they are safe in compound keys. +func sanitizeKeySegment(s string) string { + s = strings.ReplaceAll(s, " ", "_") + s = strings.ReplaceAll(s, ":", "_") + return s +} + +// heartbeatKey returns the Redis key for a probe agent heartbeat. +func heartbeatKey(probeID string) string { + return "probe:hb:" + probeID +} + +// seenKey returns the Redis key used for replay-prevention on a +// (probeID, ts) tuple. +func seenKey(probeID, ts string) string { + return "probe:seen:" + probeID + ":" + ts +} + +// -------------------------------------------------------------------------- +// Write path +// -------------------------------------------------------------------------- + +// SaveReports persists one ProbeSnapshot per (nodeID, vantage) pair from the +// given report batch and updates the probe heartbeat. +// All writes are issued in a single pipeline for efficiency. +func (s *Store) SaveReports(ctx context.Context, probeID string, vantage VantagePoint, reports []NodeReport) error { + now := time.Now().Unix() + pipe := s.rdb.Pipeline() + + for _, rep := range reports { + snap := ProbeSnapshot{ + ProbeID: probeID, + Vantage: vantage, + Report: rep, + ReceivedAt: now, + } + data, err := json.Marshal(snap) + if err != nil { + return fmt.Errorf("probe: marshal snapshot for node %s: %w", rep.NodeID, err) + } + pipe.Set(ctx, snapshotKey(rep.NodeID, vantage), data, snapshotTTL) + } + + // Update heartbeat: value is the ingest timestamp so readers can compute + // staleness without needing a separate TTL query. + pipe.Set(ctx, heartbeatKey(probeID), strconv.FormatInt(now, 10), heartbeatTTL) + + if _, err := pipe.Exec(ctx); err != nil { + return fmt.Errorf("probe: redis pipeline exec: %w", err) + } + return nil +} + +// CheckAndMarkSeen atomically checks whether the (probeID, ts) pair has been +// seen before, and marks it as seen if not. +// +// Returns (true, nil) – this is a replay; the caller should handle idempotently. +// Returns (false, nil) – first time seen; the caller should process normally. +func (s *Store) CheckAndMarkSeen(ctx context.Context, probeID, ts string) (bool, error) { + key := seenKey(probeID, ts) + // SET … NX EX: atomically set iff the key does not exist. + set, err := s.rdb.SetNX(ctx, key, "1", seenTTL).Result() + if err != nil { + return false, fmt.Errorf("probe: seen check: %w", err) + } + // SetNX returns true if the key was newly created (not a replay). + return !set, nil +} + +// -------------------------------------------------------------------------- +// Read interfaces consumed by 15D (determination logic) +// -------------------------------------------------------------------------- + +// SnapshotsByNode returns all cached probe snapshots for the given nodeID, +// indexed by vantage key string (Country:Region:ISP). +// +// A missing key — and therefore an empty map — means "no recent probe data +// from any vantage point". Callers (15D) must treat this as "unknown", not +// as a failure signal. +// +// Uses SCAN + GET in two passes. For the expected scale (dozens of vantages +// per node) this is efficient enough; 15D may add caching on top if needed. +func (s *Store) SnapshotsByNode(ctx context.Context, nodeID string) (map[string]ProbeSnapshot, error) { + pattern := "probe:" + nodeID + ":*" + + var keys []string + iter := s.rdb.Scan(ctx, 0, pattern, 0).Iterator() + for iter.Next(ctx) { + keys = append(keys, iter.Val()) + } + if err := iter.Err(); err != nil { + return nil, fmt.Errorf("probe: scan snapshots for node %s: %w", nodeID, err) + } + + result := make(map[string]ProbeSnapshot, len(keys)) + prefix := "probe:" + nodeID + ":" + + for _, k := range keys { + val, err := s.rdb.Get(ctx, k).Result() + if err == redis.Nil { + // Key expired between SCAN and GET – not an error. + continue + } + if err != nil { + return nil, fmt.Errorf("probe: get snapshot %s: %w", k, err) + } + var snap ProbeSnapshot + if err := json.Unmarshal([]byte(val), &snap); err != nil { + // Corrupted data – log-worthy but non-fatal; skip this entry. + continue + } + vk := strings.TrimPrefix(k, prefix) + result[vk] = snap + } + return result, nil +} + +// AliveProbes returns the IDs of probe agents that have sent a heartbeat +// within the last heartbeatTTL window (15 min). +// +// An empty result means "no probes have reported recently" – not that all +// probes are down. Missing heartbeat keys must never be folded into a +// failure signal by callers. +func (s *Store) AliveProbes(ctx context.Context) ([]string, error) { + var ids []string + iter := s.rdb.Scan(ctx, 0, "probe:hb:*", 0).Iterator() + for iter.Next(ctx) { + k := iter.Val() + ids = append(ids, strings.TrimPrefix(k, "probe:hb:")) + } + if err := iter.Err(); err != nil { + return nil, fmt.Errorf("probe: scan heartbeats: %w", err) + } + return ids, nil +} diff --git a/server/internal/scheduler/probe/types.go b/server/internal/scheduler/probe/types.go new file mode 100644 index 0000000..8e48028 --- /dev/null +++ b/server/internal/scheduler/probe/types.go @@ -0,0 +1,105 @@ +// Package probe implements the probe-ingest HTTP endpoint and the Redis +// storage layer for Pangolin's node-health probing subsystem. +// +// Schema contract (frozen after task 15A): +// +// POST /probe/report – ingest.go – consumed by 15B (probe client) +// and 15C (third-party dial-test) +// +// 15D (determination logic) reads from Redis via Store, but does not live here. +// This package intentionally has no determination logic, no gRPC, and no alerts. +package probe + +// ReportRequest is the JSON body for POST /probe/report. +// This schema is frozen; 15B and 15C must conform. +type ReportRequest struct { + // ProbeID identifies the probe agent submitting this report. + // Must match the X-Probe-Id request header. + ProbeID string `json:"probeId"` + + // Vantage describes the network vantage point of this probe agent. + Vantage VantagePoint `json:"vantage"` + + // Reports contains one result entry per tested node. + // A single request may carry results for multiple nodes. + Reports []NodeReport `json:"reports"` +} + +// VantagePoint describes the network vantage point of a probe agent. +type VantagePoint struct { + // ISP is the Internet Service Provider name, e.g. "China Telecom". + ISP string `json:"isp"` + + // Region is the geographic region or province, e.g. "Guangdong". + Region string `json:"region"` + + // Country is the ISO 3166-1 alpha-2 country code, e.g. "CN". + Country string `json:"country"` +} + +// NodeReport holds the probe results for a single target node. +type NodeReport struct { + // NodeID is the target node identifier (matches nodes.id in the DB). + NodeID string `json:"nodeId"` + + // Ts is the Unix timestamp (seconds) when the probe was conducted. + Ts int64 `json:"ts"` + + // L1 is the Layer-1 (TCP connectivity + RTT) result. + L1 L1Result `json:"l1"` + + // L2 is the Layer-2 (TLS handshake + ALPN negotiation) result. + // May be absent if L1 failed. + L2 *L2Result `json:"l2,omitempty"` + + // L3 is the Layer-3 (application-protocol) result. + // Optional: third-party dial-tests (15C) only populate L1/L2. + L3 *L3Result `json:"l3,omitempty"` +} + +// L1Result captures TCP-level reachability. +type L1Result struct { + // OK indicates whether a TCP connection was successfully established. + OK bool `json:"ok"` + + // RttMs is the round-trip time in milliseconds; 0 if connection failed. + RttMs int `json:"rttMs"` +} + +// L2Result captures TLS-level reachability. +type L2Result struct { + // OK indicates whether the TLS handshake succeeded. + OK bool `json:"ok"` + + // ALPN is the negotiated ALPN protocol identifier, e.g. "h3", "h2", or "". + ALPN string `json:"alpn"` + + // Err is a short error description when OK is false; empty on success. + Err string `json:"err,omitempty"` +} + +// L3Result captures application-protocol-level reachability. +// Only first-party probe agents (15B) populate this; third-party (15C) omit it. +type L3Result struct { + // OK indicates whether the application-level exchange succeeded. + OK bool `json:"ok"` + + // Err is a short error description when OK is false; empty on success. + Err string `json:"err,omitempty"` +} + +// ProbeSnapshot is the stored form of a single probe result for one +// (node, vantage) pair, persisted as JSON in Redis. +type ProbeSnapshot struct { + // ProbeID is the probe agent that produced this snapshot. + ProbeID string `json:"probeId"` + + // Vantage is the vantage point from which the probe was conducted. + Vantage VantagePoint `json:"vantage"` + + // Report is the original NodeReport from the probe agent. + Report NodeReport `json:"report"` + + // ReceivedAt is the Unix timestamp (seconds) when this report was ingested. + ReceivedAt int64 `json:"receivedAt"` +}