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

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)
}