llmgw/internal/guard/service_test.go

package guard

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// classifier is a stand-in for Lakera or Azure Content Safety: it finds
// "evil" unsafe, and refuses a call without its key.
func classifier(t *testing.T) (*httptest.Server, string) {
	key := filepath.Join(t.TempDir(), "key")
	if err := os.WriteFile(key, []byte("Bearer k1\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer k1" {
			http.Error(w, "no", http.StatusUnauthorized)
			return
		}
		var in struct{ Input string }
		_ = json.NewDecoder(r.Body).Decode(&in)
		flagged := strings.Contains(in.Input, "evil")
		_ = json.NewEncoder(w).Encode(map[string]any{"results": []any{map[string]any{"flagged": flagged, "category": "injection"}}})
	}))
	t.Cleanup(srv.Close)
	return srv, key
}

func svcPolicy(url, key, on, action, fail string) Config {
	return lib(map[string]Policy{"svc": {Fail: fail, Services: []Service{{
		Name: "cls", URL: url, HeaderFiles: map[string]string{"Authorization": key},
		Blocked: "results.0.flagged", Reason: "results.0.category", On: on, Action: action,
	}}}}, "svc")
}

func TestAnOutsideServiceBlocksOrFlagsWhatItFindsUnsafe(t *testing.T) {
	srv, key := classifier(t)
	ctx := context.Background()
	g := gate(t, svcPolicy(srv.URL, key, "", "", ""))
	v, err := g.Input(ctx, Subject{}, chat("user", "something evil"))
	if err != nil || v.Action != Block || v.Rule != "service:cls:injection" {
		t.Errorf("unsafe: %+v %v", v, err)
	}
	if v, _ := g.Input(ctx, Subject{}, chat("user", "hello")); v.Action != Allow {
		t.Errorf("safe: %+v", v)
	}
	g = gate(t, svcPolicy(srv.URL, key, "input", "flag", ""))
	if v, _ := g.Input(ctx, Subject{}, chat("user", "evil")); v.Action != Flag {
		t.Errorf("flag: %+v", v)
	}
}

func TestAServiceThatCannotBeAskedIsWhatFailSays(t *testing.T) {
	srv, _ := classifier(t)
	ctx := context.Background()
	// The wrong key: the service answers 401.
	wrong := filepath.Join(t.TempDir(), "wrong")
	_ = os.WriteFile(wrong, []byte("Bearer nope"), 0o600)
	if _, err := gate(t, svcPolicy(srv.URL, wrong, "", "", "closed")).Input(ctx, Subject{}, chat("user", "hi")); err == nil {
		t.Error("fail: closed served a request no service could check")
	}
	if v, err := gate(t, svcPolicy(srv.URL, wrong, "", "", "open")).Input(ctx, Subject{}, chat("user", "hi")); err != nil || v.Action != Allow {
		t.Errorf("fail: open: %+v %v", v, err)
	}
}

func TestAnOutsideServiceReadsTheWholeAnswer(t *testing.T) {
	srv, key := classifier(t)
	g := gate(t, svcPolicy(srv.URL, key, "output", "", ""))
	f := g.Output(context.Background(), Subject{})
	if f == nil {
		t.Fatal("no output filter for an output service")
	}
	out := f.Buffered([]byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":"an evil plan"},"finish_reason":"stop"}]}`))
	if !strings.Contains(string(out), "content_filter") || strings.Contains(string(out), "evil") || f.Verdict().Action != Block {
		t.Errorf("%s %+v", out, f.Verdict())
	}
}

func TestAServiceIsCheckedWhenThePolicyIsRead(t *testing.T) {
	for _, s := range []Service{
		{Name: "a", URL: "ftp://x", Blocked: "x"},
		{Name: "b", URL: "https://x"},
		{Name: "c", URL: "https://x", Blocked: "x", On: "sideways"},
		{Name: "d", URL: "https://x", Blocked: "x", Body: "{{"},
	} {
		if err := (Policy{Services: []Service{s}}).validate(); err == nil {
			t.Errorf("%s passed", s.Name)
		}
	}
	if v, ok := at(map[string]any{"a": []any{map[string]any{"b": 2.0}}}, "a.0.b"); !ok || !truthy(v) {
		t.Errorf("at: %v %v", v, ok)
	}
}

// recorder is a service that finds nothing unsafe and keeps what it read.
func recorder(t *testing.T) (*httptest.Server, *[]string) {
	var seen []string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var in struct{ Input string }
		_ = json.NewDecoder(r.Body).Decode(&in)
		seen = append(seen, in.Input)
		_, _ = w.Write([]byte(`{"flagged": false}`))
	}))
	t.Cleanup(srv.Close)
	return srv, &seen
}

func TestAnOutsideServiceReadsWhatTheModelReadsNotWhatWasRedacted(t *testing.T) {
	srv, seen := recorder(t)
	g := gate(t, lib(map[string]Policy{"p": {
		Input:    Rules{PII: map[string]string{"email": "redact"}},
		Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged"}},
	}}, "p"))
	if _, err := g.Input(context.Background(), Subject{}, chat("user", "mail ayse@example.com")); err != nil {
		t.Fatal(err)
	}
	if len(*seen) != 1 || strings.Contains((*seen)[0], "ayse@") || !strings.Contains((*seen)[0], "[EMAIL_1]") {
		t.Errorf("the service read %q", *seen)
	}
}

func TestAStreamedAnswerIsReadByAServiceOnce(t *testing.T) {
	srv, seen := recorder(t)
	g := gate(t, lib(map[string]Policy{"p": {
		Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged", On: "output"}},
	}}, "p"))
	got, _, _ := stream(g.Output(context.Background(), Subject{Stream: true}), "Hello ", "there.")
	if got != "Hello there." {
		t.Errorf("client got %q", got)
	}
	if len(*seen) != 1 || (*seen)[0] != "Hello there." {
		t.Errorf("the service was asked %d times: %q", len(*seen), *seen)
	}
}

func TestAServiceNamedByTwoPoliciesIsAskedTheStricterWay(t *testing.T) {
	a := Policy{Services: []Service{{Name: "cls", URL: "https://c", Blocked: "x", On: "input", Action: "flag"}}}
	b := Policy{Services: []Service{{Name: "cls", URL: "https://c", Blocked: "x", On: "output", Action: "block"}}}
	got := Merge([]string{"a", "b"}, []Policy{a, b}).Services
	if len(got) != 1 || got[0].action() != Block || !got[0].onInput() || !got[0].onOutput() {
		t.Errorf("%+v", got)
	}
}

func TestAnOutputServiceReadsTheAnswerWithoutTheRestoredValues(t *testing.T) {
	srv, seen := recorder(t)
	g := gate(t, lib(map[string]Policy{"p": {
		Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged", On: "output"}},
	}}, "p"))
	f := g.Output(context.Background(), Subject{Restored: map[string]string{"[EMAIL_1]": "ayse@example.com"}})
	f.Buffered([]byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":"Sent to ayse@example.com."},"finish_reason":"stop"}]}`))
	if len(*seen) != 1 || strings.Contains((*seen)[0], "ayse@") || !strings.Contains((*seen)[0], "[EMAIL_1]") {
		t.Errorf("the service read %q", *seen)
	}
}

func TestOfTwoChoicesTheServiceReadsTheFirstWhole(t *testing.T) {
	// The second finishes first; the first is still arriving then.
	srv, seen := recorder(t)
	g := gate(t, lib(map[string]Policy{"p": {
		Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged", On: "output"}},
	}}, "p"))
	f := g.Output(context.Background(), Subject{Stream: true})
	ev := func(idx int, text string, finish bool) string {
		c := map[string]any{"index": idx, "delta": map[string]any{"content": text}, "finish_reason": nil}
		if finish {
			c["finish_reason"] = "stop"
		}
		b, _ := json.Marshal(map[string]any{"choices": []any{c}})
		return string(b)
	}
	for _, e := range []string{ev(0, "The first ", false), ev(1, "Second.", true), ev(0, "answer.", true)} {
		f.OnEvent(e)
	}
	f.OnEnd()
	if len(*seen) != 1 || (*seen)[0] != "The first answer." {
		t.Errorf("the service read %q", *seen)
	}
}