merge: maestro/tsk_rBPr0Xuy10bz [tsk_rBPr0Xuy10bz] 探针汇聚入口 + Redis 探针存储 (tsk_eu2qVvni1HJR)

Manually apply changes from maestro/tsk_rBPr0Xuy10bz that could not be
auto-merged due to uncommitted changes on main at merge time.

Added:
- server/internal/scheduler/probe/types.go  – frozen schema types (ReportRequest,
  VantagePoint, NodeReport, L1/L2/L3Result, ProbeSnapshot)
- server/internal/scheduler/probe/store.go  – Redis Store: SaveReports, CheckAndMarkSeen,
  SnapshotsByNode, AliveProbes with TTL constants and key-schema docs
- server/internal/scheduler/probe/ingest.go – IngestHandler (POST /probe/report),
  HMAC-SHA256 auth + timestamp window + replay prevention via SetNX
- server/internal/scheduler/probe/ingest_test.go – 17 tests (auth failures, integration,
  Store unit tests); all pass with miniredis

Updated:
- server/cmd/server/main.go – register /probe/report route when PROBE_SECRETS env is set

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
wangjia
2026-06-13 14:35:24 +08:00
parent 787151245e
commit f76e1797c2
5 changed files with 1163 additions and 0 deletions
+37
View File
@@ -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
}
+185
View File
@@ -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)
}
@@ -0,0 +1,624 @@
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())
}
}
+212
View File
@@ -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
}
+105
View File
@@ -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"`
}