Skip to content

Commit 8275f7f

Browse files
authored
fix(guard): correct PIGuard score semantics and vendor JSON-faithful HTTP gateway (#85)
Two independent bugs caused the PIGuard sidecar to reject virtually every memory fact with scores pegged at 0.98-0.99, including innocuous content: 1. resultFromResponse treated the sidecar score as an injection probability (label == INJECTION || score >= threshold). The score is actually the confidence of the predicted label - a confident BENIGN result scores ~1.0 as well - so the threshold clause rejected almost everything at the default 0.9. The threshold now only gates INJECTION labels. 2. The docker-compose gateway (upstream examples/http-gateway, a raw-text demo) wrapped odek's JSON {"text":...} bodies as literal text, so the model scored the JSON wrapper instead of the content, flipping benign facts to INJECTION at 0.98-0.99. Replaced with a vendored gateway (docker/piguard-gateway) that forwards daemon-protocol JSON verbatim (with raw-text fallback for curl) and builds without amd64-pinned base digests (Apple Silicon safe). Verified end-to-end against a live stack: 8/8 benign facts classified BENIGN, 4/4 real injections classified INJECTION (still rejected), plus regression tests for the threshold semantics.
1 parent 5a28509 commit 8275f7f

7 files changed

Lines changed: 208 additions & 7 deletions

File tree

docker/docker-compose.yml

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,12 @@ services:
161161
# The daemon speaks newline-delimited JSON over a private Unix socket; the
162162
# gateway bridges that to plain HTTP on the compose network so odek can use
163163
# its existing HTTP-based piguard client without sharing a socket volume.
164+
#
165+
# Apple Silicon note: the upstream daemon Dockerfile pins its base images by
166+
# amd64 digest, which yields x86-64 binaries that crash under Rosetta on
167+
# macOS. If `docker compose up` fails with "rosetta error", build the daemon
168+
# locally from the guard repo with the digest pins removed and tag it
169+
# piguard:local (compose reuses the existing image).
164170
piguard:
165171
profiles: ["restricted", "godmode", "telegram-restricted", "telegram-godmode"]
166172
# GHCR packages for this repo are private, so we build the daemon image locally
@@ -188,10 +194,12 @@ services:
188194

189195
piguard-gateway:
190196
profiles: ["restricted", "godmode", "telegram-restricted", "telegram-godmode"]
191-
# Built locally from the public guard repo source (GHCR packages are private).
197+
# Vendored gateway (./piguard-gateway): forwards odek's JSON protocol
198+
# ({"text"|"long"|"texts"}) verbatim to the daemon socket. The upstream
199+
# examples/http-gateway is a raw-text demo that double-encodes odek's JSON
200+
# bodies, making the model score the wrapper instead of the content.
192201
build:
193-
context: https://github.com/BackendStack21/go-prompt-injection-guard.git#v1.0.0
194-
dockerfile: examples/http-gateway/Dockerfile
202+
context: ./piguard-gateway
195203
image: piguard-gateway:local
196204
command: ["--addr=:8080", "--socket=/run/piguard/piguard.sock"]
197205
volumes:

docker/piguard-gateway/Dockerfile

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# syntax=docker/dockerfile:1
2+
3+
# Vendored HTTP gateway for the PIGuard daemon. Replaces the upstream
4+
# examples/http-gateway (a raw-text demo) with a bridge that forwards odek's
5+
# JSON protocol verbatim to the daemon socket. Base images are intentionally
6+
# NOT pinned by digest so builds resolve to the host architecture (the
7+
# upstream amd64-pinned digests produce x86-64 binaries inside arm64-labeled
8+
# images on Apple Silicon).
9+
10+
# ---- build stage ----
11+
FROM golang:1.26-bookworm AS build
12+
WORKDIR /src
13+
COPY go.mod main.go ./
14+
RUN CGO_ENABLED=0 go build -trimpath -ldflags "-s -w" -o /out/piguard-gateway .
15+
16+
# ---- runtime stage ----
17+
# Runs as root (distroless default) on purpose: the daemon creates its socket
18+
# owned by its own non-root user, and the gateway must be able to connect.
19+
# The upstream example gateway does the same. The container is unprivileged
20+
# otherwise and only reachable on the private compose network.
21+
FROM gcr.io/distroless/static-debian12
22+
COPY --from=build /out/piguard-gateway /usr/local/bin/piguard-gateway
23+
ENTRYPOINT ["piguard-gateway"]
24+
CMD ["--addr=:8080", "--socket=/run/piguard/piguard.sock"]

docker/piguard-gateway/go.mod

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
module piguard-gateway
2+
3+
go 1.24

docker/piguard-gateway/main.go

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
// Command piguard-gateway is a small HTTP front end for the PIGuard daemon.
2+
// It forwards requests to the daemon's Unix socket and returns the JSON
3+
// reply, so odek's HTTP-based guard client can reach the socket-only daemon
4+
// over the compose network.
5+
//
6+
// The daemon speaks newline-delimited JSON ({"text":...}, {"long":...},
7+
// {"texts":[...]}), which is exactly what odek's piguard client sends, so
8+
// JSON bodies are forwarded verbatim. A non-JSON body is treated as raw text
9+
// and wrapped as {"text": <body>} for convenience (curl -d 'some text').
10+
//
11+
// Endpoints:
12+
//
13+
// GET /healthz liveness — checks the daemon socket is reachable
14+
// POST /detect {"text": "..."} (or raw text)
15+
// POST /long {"long": "..."} (or raw text)
16+
// POST /raw {"texts": [...]} (batch; or any daemon JSON line)
17+
//
18+
// Note: the daemon dispatches on the JSON keys, not the HTTP path; the paths
19+
// exist only to mirror odek's client configuration.
20+
package main
21+
22+
import (
23+
"bufio"
24+
"bytes"
25+
"encoding/json"
26+
"flag"
27+
"io"
28+
"log"
29+
"net"
30+
"net/http"
31+
"time"
32+
)
33+
34+
// maxBody caps a request body, mirroring the daemon's 1 MiB request limit
35+
// (with a little headroom for JSON framing).
36+
const maxBody = 1 << 20
37+
38+
type gateway struct {
39+
socket string
40+
}
41+
42+
// forward sends one newline-delimited request to the daemon and returns its
43+
// single-line reply.
44+
func (g *gateway) forward(line []byte) ([]byte, error) {
45+
conn, err := net.DialTimeout("unix", g.socket, 2*time.Second)
46+
if err != nil {
47+
return nil, err
48+
}
49+
defer conn.Close()
50+
_ = conn.SetDeadline(time.Now().Add(60 * time.Second))
51+
52+
if _, err := conn.Write(append(bytes.TrimRight(line, "\n"), '\n')); err != nil {
53+
return nil, err
54+
}
55+
resp, err := bufio.NewReader(conn).ReadBytes('\n')
56+
if err != nil && len(resp) == 0 {
57+
return nil, err
58+
}
59+
return resp, nil
60+
}
61+
62+
// daemonLine converts an HTTP request body into one daemon protocol line.
63+
// Bodies that already are daemon protocol JSON (carrying a text, long, or
64+
// texts key) pass through verbatim; anything else is wrapped as raw text.
65+
func daemonLine(body []byte) ([]byte, error) {
66+
trimmed := bytes.TrimSpace(body)
67+
var probe map[string]json.RawMessage
68+
if err := json.Unmarshal(trimmed, &probe); err == nil {
69+
for _, key := range []string{"text", "long", "texts"} {
70+
if _, ok := probe[key]; ok {
71+
return trimmed, nil
72+
}
73+
}
74+
}
75+
return json.Marshal(map[string]string{"text": string(body)})
76+
}
77+
78+
func (g *gateway) handleDetect(w http.ResponseWriter, r *http.Request) {
79+
if r.Method != http.MethodPost {
80+
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
81+
return
82+
}
83+
body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxBody))
84+
if err != nil {
85+
http.Error(w, "request too large", http.StatusRequestEntityTooLarge)
86+
return
87+
}
88+
line, err := daemonLine(body)
89+
if err != nil {
90+
http.Error(w, err.Error(), http.StatusBadRequest)
91+
return
92+
}
93+
resp, err := g.forward(line)
94+
if err != nil {
95+
http.Error(w, err.Error(), http.StatusBadGateway)
96+
return
97+
}
98+
w.Header().Set("Content-Type", "application/json")
99+
_, _ = w.Write(resp)
100+
}
101+
102+
func (g *gateway) handleHealth(w http.ResponseWriter, r *http.Request) {
103+
conn, err := net.DialTimeout("unix", g.socket, 2*time.Second)
104+
if err != nil {
105+
http.Error(w, "daemon unreachable: "+err.Error(), http.StatusServiceUnavailable)
106+
return
107+
}
108+
conn.Close()
109+
_, _ = w.Write([]byte("ok\n"))
110+
}
111+
112+
func main() {
113+
addr := flag.String("addr", ":8080", "HTTP listen address")
114+
socket := flag.String("socket", "/run/piguard/piguard.sock", "PIGuard daemon Unix socket")
115+
flag.Parse()
116+
117+
g := &gateway{socket: *socket}
118+
mux := http.NewServeMux()
119+
mux.HandleFunc("/healthz", g.handleHealth)
120+
mux.HandleFunc("/detect", g.handleDetect)
121+
mux.HandleFunc("/long", g.handleDetect)
122+
mux.HandleFunc("/raw", g.handleDetect)
123+
124+
log.Printf("piguard-gateway: listening on %s, daemon socket %s", *addr, *socket)
125+
log.Fatal(http.ListenAndServe(*addr, mux))
126+
}

docker/piguard/.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
*.data
44
*.json
55
tokenizer.*
6+
modeling_*.py
67
!/.gitkeep
78
!/.gitignore
89

internal/guard/piguard.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -211,8 +211,14 @@ func (p *piguardClient) post(ctx context.Context, urlStr string, body []byte) (*
211211

212212
// resultFromResponse converts a PIGuard response into a Result, applying the
213213
// configured threshold.
214+
//
215+
// The sidecar's score is the confidence of the predicted label — whichever
216+
// label that is — not the injection probability (a confident BENIGN result
217+
// also scores ~1.0). The threshold therefore only applies to INJECTION
218+
// labels; comparing it against the score of a BENIGN result would reject
219+
// virtually everything, since the model is confident on most inputs.
214220
func resultFromResponse(r detectResponse, start time.Time, threshold float64) Result {
215-
injected := r.Label == "INJECTION" || r.Score >= threshold
221+
injected := r.Label == "INJECTION" && r.Score >= threshold
216222
return Result{
217223
Label: r.Label,
218224
Score: r.Score,

internal/guard/piguard_test.go

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"net/http/httptest"
88
"strings"
99
"testing"
10+
"time"
1011
)
1112

1213
// fakePiguardServer returns an httptest.Server that mimics the PIGuard HTTP gateway.
@@ -29,7 +30,7 @@ func fakePiguardServer(t *testing.T) *httptest.Server {
2930
http.Error(w, err.Error(), http.StatusBadRequest)
3031
return
3132
}
32-
resp := detectResponse{Label: "BENIGN", Score: 0.1}
33+
resp := detectResponse{Label: "BENIGN", Score: 0.9997} // realistic high-confidence BENIGN
3334
if strings.Contains(strings.ToLower(req.Text), "ignore") {
3435
resp.Label = "INJECTION"
3536
resp.Score = 0.999
@@ -42,7 +43,7 @@ func fakePiguardServer(t *testing.T) *httptest.Server {
4243
http.Error(w, err.Error(), http.StatusBadRequest)
4344
return
4445
}
45-
resp := detectResponse{Label: "BENIGN", Score: 0.1}
46+
resp := detectResponse{Label: "BENIGN", Score: 0.9997} // realistic high-confidence BENIGN
4647
if strings.Contains(strings.ToLower(req.Long), "ignore") {
4748
resp.Label = "INJECTION"
4849
resp.Score = 0.999
@@ -57,7 +58,7 @@ func fakePiguardServer(t *testing.T) *httptest.Server {
5758
}
5859
results := make([]detectResponse, len(req.Texts))
5960
for i, text := range req.Texts {
60-
results[i] = detectResponse{Label: "BENIGN", Score: 0.1}
61+
results[i] = detectResponse{Label: "BENIGN", Score: 0.9997}
6162
if strings.Contains(strings.ToLower(text), "ignore") {
6263
results[i] = detectResponse{Label: "INJECTION", Score: 0.999}
6364
}
@@ -224,3 +225,35 @@ func TestPiguardClient_NonOKStatus(t *testing.T) {
224225
t.Fatal("expected error for non-OK status")
225226
}
226227
}
228+
229+
// TestResultFromResponse_ThresholdSemantics pins the score semantics of the
230+
// PIGuard sidecar: the score is the confidence of the predicted label, not
231+
// the injection probability. A high-confidence BENIGN result (score ~1.0)
232+
// must NOT be treated as an injection, and the threshold only gates INJECTION
233+
// labels. Regression test for the bug where every confident BENIGN memory
234+
// fact was rejected with score >= threshold (default 0.9).
235+
func TestResultFromResponse_ThresholdSemantics(t *testing.T) {
236+
cases := []struct {
237+
name string
238+
label string
239+
score float64
240+
threshold float64
241+
want bool
242+
}{
243+
{"confident benign passes", "BENIGN", 0.9999, 0.9, false},
244+
{"weak benign passes", "BENIGN", 0.55, 0.9, false},
245+
{"confident injection rejected", "INJECTION", 0.98, 0.9, true},
246+
{"injection at threshold rejected", "INJECTION", 0.9, 0.9, true},
247+
{"weak injection below threshold passes", "INJECTION", 0.62, 0.9, false},
248+
{"unknown label passes regardless of score", "SOMETHING", 1.0, 0.9, false},
249+
}
250+
for _, tc := range cases {
251+
t.Run(tc.name, func(t *testing.T) {
252+
got := resultFromResponse(detectResponse{Label: tc.label, Score: tc.score}, time.Now(), tc.threshold)
253+
if got.Injected != tc.want {
254+
t.Errorf("resultFromResponse(%s, %.4f, thr %.2f).Injected = %v, want %v",
255+
tc.label, tc.score, tc.threshold, got.Injected, tc.want)
256+
}
257+
})
258+
}
259+
}

0 commit comments

Comments
 (0)