bastion 0.53.2

client-go/pii_test.go

@@ -0,0 +1,105 @@
+package main
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+// piiProbe is opt-in: no PRIVATAE_PII_ADDR means no probe, no capability, no
+// heartbeat fields — a box without the service is unaffected by this lane.
+func TestPiiProbeRequiresEnv(t *testing.T) {
+	t.Setenv("PRIVATAE_PII_ADDR", "")
+	if adv, ok := piiProbe(); adv || ok {
+		t.Errorf("piiProbe(no env) = (%v, %v), want (false, false)", adv, ok)
+	}
+}
+
+func TestPiiProbeHealthy(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/health" {
+			t.Errorf("probe hit %s, want /health", r.URL.Path)
+		}
+		w.Write([]byte(`{"ok":true}`))
+	}))
+	defer srv.Close()
+	t.Setenv("PRIVATAE_PII_ADDR", strings.TrimPrefix(srv.URL, "http://"))
+	if adv, ok := piiProbe(); !adv || !ok {
+		t.Errorf("piiProbe(healthy) = (%v, %v), want (true, true)", adv, ok)
+	}
+}
+
+// Env set but upstream dead: still advertised (the forwarder runs), with
+// pii_ok=false as the distinct dead-upstream signal for the host.
+func TestPiiProbeDeadUpstreamStillAdvertised(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusBadGateway)
+	}))
+	addr := strings.TrimPrefix(srv.URL, "http://")
+	srv.Close() // connection refused from here on
+	t.Setenv("PRIVATAE_PII_ADDR", addr)
+	if adv, ok := piiProbe(); !adv || ok {
+		t.Errorf("piiProbe(dead upstream) = (%v, %v), want (true, false)", adv, ok)
+	}
+}
+
+func TestHeartbeatAdvertisesPii(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte(`{"ok":true}`))
+	}))
+	defer srv.Close()
+	t.Setenv("PRIVATAE_PII_ADDR", strings.TrimPrefix(srv.URL, "http://"))
+	t.Setenv("PRIVATAE_LLM_ADDR", "")
+
+	raw := buildHeartbeatBody("", "h", "linux", "test", "pk", "fp", "100.64.0.9", "ep", false, nil, nil, nil)
+	var body struct {
+		Capabilities []string `json:"capabilities"`
+		PiiPort      *int     `json:"pii_port"`
+		PiiOK        *bool    `json:"pii_ok"`
+	}
+	if err := json.Unmarshal([]byte(raw), &body); err != nil {
+		t.Fatalf("heartbeat body: %v", err)
+	}
+	found := false
+	for _, c := range body.Capabilities {
+		if c == "pii" {
+			found = true
+		}
+	}
+	if !found {
+		t.Errorf("capabilities = %v, want to include \"pii\"", body.Capabilities)
+	}
+	if body.PiiPort == nil || *body.PiiPort != piiForwardPort {
+		t.Errorf("pii_port = %v, want %d (the port servePIIForward listens on)", body.PiiPort, piiForwardPort)
+	}
+	if body.PiiOK == nil || !*body.PiiOK {
+		t.Errorf("pii_ok = %v, want true", body.PiiOK)
+	}
+}
+
+// Absence semantics: unset env produces NO pii fields at all — the host
+// reads a missing pii_port as "service not present", so an always-present
+// field (like llm_port's style) would break the honest disabled state.
+func TestHeartbeatOmitsPiiWhenUnset(t *testing.T) {
+	t.Setenv("PRIVATAE_PII_ADDR", "")
+	t.Setenv("PRIVATAE_LLM_ADDR", "")
+	raw := buildHeartbeatBody("", "h", "linux", "test", "pk", "fp", "100.64.0.9", "ep", false, nil, nil, nil)
+	var body map[string]interface{}
+	if err := json.Unmarshal([]byte(raw), &body); err != nil {
+		t.Fatalf("heartbeat body: %v", err)
+	}
+	if _, present := body["pii_port"]; present {
+		t.Error("pii_port present with PRIVATAE_PII_ADDR unset; want omitted")
+	}
+	if _, present := body["pii_ok"]; present {
+		t.Error("pii_ok present with PRIVATAE_PII_ADDR unset; want omitted")
+	}
+	caps, _ := body["capabilities"].([]interface{})
+	for _, c := range caps {
+		if c == "pii" {
+			t.Error("capability \"pii\" present with PRIVATAE_PII_ADDR unset")
+		}
+	}
+}