sift/internal/fleetctx/fleetctx_test.go

package fleetctx

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestRestrictedRefusedEvenIfAllowlisted(t *testing.T) {
	c := &Client{Base: "http://127.0.0.1:1", Allow: []string{"guardian-star", "hermes-fleet"}, Permit: []string{"straylight-gs"}}
	items, warn := c.Search(context.Background(), "guardian-star", "talos", "anything")
	if len(items) != 0 {
		t.Fatalf("items %v", items)
	}
	if warn == "" {
		t.Fatal("want warning")
	}
}

func TestRestrictedRefusedWhenFleetCallerRemapped(t *testing.T) {
	c := &Client{Base: "http://127.0.0.1:1", Allow: []string{"guardian-star"}, Permit: []string{"straylight-gs"}}
	items, warn := c.Search(context.Background(), "guardian-star", "straylight", "q")
	if len(items) != 0 || warn == "" {
		t.Fatalf("items %v warn %q", items, warn)
	}
}

func TestRestrictedCallerMappedToFleetMakesNoRequest(t *testing.T) {
	var hits atomic.Int32
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.WriteHeader(200)
		_, _ = w.Write([]byte("[]"))
	}))
	defer ts.Close()
	c := &Client{
		Base:   ts.URL,
		Allow:  []string{"guardian-star", "hermes-fleet"},
		Permit: []string{"straylight-gs"},
		HTTP:   ts.Client(),
	}
	items, warn := c.Search(context.Background(), "hermes-fleet", "straylight-gs", "secret query")
	if hits.Load() != 0 {
		t.Fatalf("honcho hits %d", hits.Load())
	}
	if len(items) != 0 || warn == "" {
		t.Fatalf("items %v warn %q", items, warn)
	}
}

func TestRestrictedAllowedForPermittedCaller(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v3/workspaces/guardian-star/search" {
			t.Errorf("path %s", r.URL.Path)
		}
		_ = json.NewEncoder(w).Encode([]map[string]any{{
			"content":      "gs already knows Y",
			"created_at":   "2026-09-23T00:00:00Z",
			"peer_id":      "kyle",
			"session_id":   "s1",
			"workspace_id": "guardian-star",
			"token_count":  4,
		}})
	}))
	defer ts.Close()
	c := &Client{
		Base:   ts.URL,
		Allow:  []string{"guardian-star", "hermes-fleet"},
		Permit: []string{"straylight-gs"},
		HTTP:   ts.Client(),
	}
	items, warn := c.Search(context.Background(), "guardian-star", "straylight-gs", "Y")
	if warn != "" || len(items) != 1 || items[0].Peer != "kyle" || items[0].Content != "gs already knows Y" {
		t.Fatalf("%v %q", items, warn)
	}
}

func TestRestrictedSkippedIfNotAllowlisted(t *testing.T) {
	c := &Client{Base: "http://127.0.0.1:1", Permit: []string{"straylight-gs"}}
	_, warn := c.Search(context.Background(), "guardian-star", "straylight-gs", "q")
	if warn == "" {
		t.Fatal("want skip")
	}
}

func TestAllowlist(t *testing.T) {
	c := &Client{Base: "http://127.0.0.1:1", Allow: []string{"hermes-fleet"}}
	_, warn := c.Search(context.Background(), "other-ws", "straylight", "q")
	if warn == "" {
		t.Fatal("want skip")
	}
}

func TestTimeoutWarning(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(4 * time.Second)
		w.WriteHeader(200)
		_, _ = w.Write([]byte("[]"))
	}))
	defer ts.Close()
	c := &Client{Base: ts.URL, Allow: []string{"hermes-fleet"}, HTTP: ts.Client()}
	start := time.Now()
	_, warn := c.Search(context.Background(), "hermes-fleet", "straylight", "q")
	if time.Since(start) > 3500*time.Millisecond {
		t.Fatalf("took %s", time.Since(start))
	}
	if warn == "" {
		t.Fatal("want warning on timeout")
	}
}

func TestHappyPath(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		b, _ := io.ReadAll(r.Body)
		var req map[string]any
		_ = json.Unmarshal(b, &req)
		if req["limit"] != float64(5) {
			t.Errorf("limit %v", req["limit"])
		}
		_ = json.NewEncoder(w).Encode([]map[string]any{{
			"content":      "fleet already knows X",
			"created_at":   "2026-09-22T00:00:00Z",
			"peer_id":      "kyle",
			"session_id":   "s1",
			"workspace_id": "hermes-fleet",
			"token_count":  4,
		}})
	}))
	defer ts.Close()
	c := &Client{Base: ts.URL, Allow: []string{"hermes-fleet"}, HTTP: ts.Client()}
	items, warn := c.Search(context.Background(), "hermes-fleet", "straylight", "X")
	if warn != "" || len(items) != 1 || items[0].Peer != "kyle" {
		t.Fatalf("%v %q", items, warn)
	}
}

func TestTraversalRejected(t *testing.T) {
	c := &Client{Base: "http://127.0.0.1:1", Allow: []string{"hermes-fleet"}}
	_, warn := c.Search(context.Background(), "hermes-fleet/../guardian-star", "straylight-gs", "q")
	if warn == "" {
		t.Fatal("want skip")
	}
}

func TestDropsForeignWorkspace(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode([]map[string]any{
			{"content": "secret", "peer_id": "x", "workspace_id": "guardian-star", "session_id": "s", "created_at": "t"},
			{"content": "ok", "peer_id": "y", "workspace_id": "hermes-fleet", "session_id": "s", "created_at": "t"},
		})
	}))
	defer ts.Close()
	c := &Client{Base: ts.URL, Allow: []string{"hermes-fleet"}, HTTP: ts.Client()}
	items, warn := c.Search(context.Background(), "hermes-fleet", "straylight", "q")
	if warn != "" || len(items) != 1 || items[0].Content != "ok" {
		t.Fatalf("%v %q", items, warn)
	}
}

func TestConcurrentSearch(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode([]map[string]any{{
			"content": "ok", "peer_id": "y", "workspace_id": "hermes-fleet", "session_id": "s", "created_at": "t",
		}})
	}))
	defer ts.Close()
	c := &Client{Base: ts.URL, Allow: []string{"hermes-fleet"}, HTTP: ts.Client()}
	var wg sync.WaitGroup
	for i := 0; i < 32; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, _ = c.Search(context.Background(), "hermes-fleet", "straylight", "q")
		}()
	}
	wg.Wait()
	if c.Calls.Load() != 32 {
		t.Fatalf("calls %d", c.Calls.Load())
	}
}

func TestCheckMappingsWarnsMissingKey(t *testing.T) {
	var buf bytes.Buffer
	log := slog.New(slog.NewTextHandler(&buf, nil))
	c := &Client{Permit: []string{"straylight-gs"}, Log: log}
	c.CheckMappings(nil)
	got := buf.String()
	if !strings.Contains(got, "permit_restricted caller has no key") || !strings.Contains(got, "straylight-gs") {
		t.Fatalf("missing warn: %s", got)
	}
	buf.Reset()
	c.CheckMappings([]CallerWS{{Name: "straylight-gs", Workspace: "guardian-star"}})
	if strings.Contains(buf.String(), "no key") || strings.Contains(buf.String(), "mapped off") {
		t.Fatalf("mapped ok still logged: %s", buf.String())
	}
}