diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 3064ffa..838247f 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -161,6 +161,12 @@ services: # The daemon speaks newline-delimited JSON over a private Unix socket; the # gateway bridges that to plain HTTP on the compose network so odek can use # its existing HTTP-based piguard client without sharing a socket volume. + # + # Apple Silicon note: the upstream daemon Dockerfile pins its base images by + # amd64 digest, which yields x86-64 binaries that crash under Rosetta on + # macOS. If `docker compose up` fails with "rosetta error", build the daemon + # locally from the guard repo with the digest pins removed and tag it + # piguard:local (compose reuses the existing image). piguard: profiles: ["restricted", "godmode", "telegram-restricted", "telegram-godmode"] # GHCR packages for this repo are private, so we build the daemon image locally @@ -188,10 +194,12 @@ services: piguard-gateway: profiles: ["restricted", "godmode", "telegram-restricted", "telegram-godmode"] - # Built locally from the public guard repo source (GHCR packages are private). + # Vendored gateway (./piguard-gateway): forwards odek's JSON protocol + # ({"text"|"long"|"texts"}) verbatim to the daemon socket. The upstream + # examples/http-gateway is a raw-text demo that double-encodes odek's JSON + # bodies, making the model score the wrapper instead of the content. build: - context: https://github.com/BackendStack21/go-prompt-injection-guard.git#v1.0.0 - dockerfile: examples/http-gateway/Dockerfile + context: ./piguard-gateway image: piguard-gateway:local command: ["--addr=:8080", "--socket=/run/piguard/piguard.sock"] volumes: diff --git a/docker/piguard-gateway/Dockerfile b/docker/piguard-gateway/Dockerfile new file mode 100644 index 0000000..ffd61aa --- /dev/null +++ b/docker/piguard-gateway/Dockerfile @@ -0,0 +1,24 @@ +# syntax=docker/dockerfile:1 + +# Vendored HTTP gateway for the PIGuard daemon. Replaces the upstream +# examples/http-gateway (a raw-text demo) with a bridge that forwards odek's +# JSON protocol verbatim to the daemon socket. Base images are intentionally +# NOT pinned by digest so builds resolve to the host architecture (the +# upstream amd64-pinned digests produce x86-64 binaries inside arm64-labeled +# images on Apple Silicon). + +# ---- build stage ---- +FROM golang:1.26-bookworm AS build +WORKDIR /src +COPY go.mod main.go ./ +RUN CGO_ENABLED=0 go build -trimpath -ldflags "-s -w" -o /out/piguard-gateway . + +# ---- runtime stage ---- +# Runs as root (distroless default) on purpose: the daemon creates its socket +# owned by its own non-root user, and the gateway must be able to connect. +# The upstream example gateway does the same. The container is unprivileged +# otherwise and only reachable on the private compose network. +FROM gcr.io/distroless/static-debian12 +COPY --from=build /out/piguard-gateway /usr/local/bin/piguard-gateway +ENTRYPOINT ["piguard-gateway"] +CMD ["--addr=:8080", "--socket=/run/piguard/piguard.sock"] diff --git a/docker/piguard-gateway/go.mod b/docker/piguard-gateway/go.mod new file mode 100644 index 0000000..95432f6 --- /dev/null +++ b/docker/piguard-gateway/go.mod @@ -0,0 +1,3 @@ +module piguard-gateway + +go 1.24 diff --git a/docker/piguard-gateway/main.go b/docker/piguard-gateway/main.go new file mode 100644 index 0000000..82a1e98 --- /dev/null +++ b/docker/piguard-gateway/main.go @@ -0,0 +1,126 @@ +// Command piguard-gateway is a small HTTP front end for the PIGuard daemon. +// It forwards requests to the daemon's Unix socket and returns the JSON +// reply, so odek's HTTP-based guard client can reach the socket-only daemon +// over the compose network. +// +// The daemon speaks newline-delimited JSON ({"text":...}, {"long":...}, +// {"texts":[...]}), which is exactly what odek's piguard client sends, so +// JSON bodies are forwarded verbatim. A non-JSON body is treated as raw text +// and wrapped as {"text": } for convenience (curl -d 'some text'). +// +// Endpoints: +// +// GET /healthz liveness — checks the daemon socket is reachable +// POST /detect {"text": "..."} (or raw text) +// POST /long {"long": "..."} (or raw text) +// POST /raw {"texts": [...]} (batch; or any daemon JSON line) +// +// Note: the daemon dispatches on the JSON keys, not the HTTP path; the paths +// exist only to mirror odek's client configuration. +package main + +import ( + "bufio" + "bytes" + "encoding/json" + "flag" + "io" + "log" + "net" + "net/http" + "time" +) + +// maxBody caps a request body, mirroring the daemon's 1 MiB request limit +// (with a little headroom for JSON framing). +const maxBody = 1 << 20 + +type gateway struct { + socket string +} + +// forward sends one newline-delimited request to the daemon and returns its +// single-line reply. +func (g *gateway) forward(line []byte) ([]byte, error) { + conn, err := net.DialTimeout("unix", g.socket, 2*time.Second) + if err != nil { + return nil, err + } + defer conn.Close() + _ = conn.SetDeadline(time.Now().Add(60 * time.Second)) + + if _, err := conn.Write(append(bytes.TrimRight(line, "\n"), '\n')); err != nil { + return nil, err + } + resp, err := bufio.NewReader(conn).ReadBytes('\n') + if err != nil && len(resp) == 0 { + return nil, err + } + return resp, nil +} + +// daemonLine converts an HTTP request body into one daemon protocol line. +// Bodies that already are daemon protocol JSON (carrying a text, long, or +// texts key) pass through verbatim; anything else is wrapped as raw text. +func daemonLine(body []byte) ([]byte, error) { + trimmed := bytes.TrimSpace(body) + var probe map[string]json.RawMessage + if err := json.Unmarshal(trimmed, &probe); err == nil { + for _, key := range []string{"text", "long", "texts"} { + if _, ok := probe[key]; ok { + return trimmed, nil + } + } + } + return json.Marshal(map[string]string{"text": string(body)}) +} + +func (g *gateway) handleDetect(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxBody)) + if err != nil { + http.Error(w, "request too large", http.StatusRequestEntityTooLarge) + return + } + line, err := daemonLine(body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + resp, err := g.forward(line) + if err != nil { + http.Error(w, err.Error(), http.StatusBadGateway) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(resp) +} + +func (g *gateway) handleHealth(w http.ResponseWriter, r *http.Request) { + conn, err := net.DialTimeout("unix", g.socket, 2*time.Second) + if err != nil { + http.Error(w, "daemon unreachable: "+err.Error(), http.StatusServiceUnavailable) + return + } + conn.Close() + _, _ = w.Write([]byte("ok\n")) +} + +func main() { + addr := flag.String("addr", ":8080", "HTTP listen address") + socket := flag.String("socket", "/run/piguard/piguard.sock", "PIGuard daemon Unix socket") + flag.Parse() + + g := &gateway{socket: *socket} + mux := http.NewServeMux() + mux.HandleFunc("/healthz", g.handleHealth) + mux.HandleFunc("/detect", g.handleDetect) + mux.HandleFunc("/long", g.handleDetect) + mux.HandleFunc("/raw", g.handleDetect) + + log.Printf("piguard-gateway: listening on %s, daemon socket %s", *addr, *socket) + log.Fatal(http.ListenAndServe(*addr, mux)) +} diff --git a/docker/piguard/.gitignore b/docker/piguard/.gitignore index 187c327..8e4ab80 100644 --- a/docker/piguard/.gitignore +++ b/docker/piguard/.gitignore @@ -3,6 +3,7 @@ *.data *.json tokenizer.* +modeling_*.py !/.gitkeep !/.gitignore diff --git a/internal/guard/piguard.go b/internal/guard/piguard.go index 7e46115..dbfd566 100644 --- a/internal/guard/piguard.go +++ b/internal/guard/piguard.go @@ -211,8 +211,14 @@ func (p *piguardClient) post(ctx context.Context, urlStr string, body []byte) (* // resultFromResponse converts a PIGuard response into a Result, applying the // configured threshold. +// +// The sidecar's score is the confidence of the predicted label — whichever +// label that is — not the injection probability (a confident BENIGN result +// also scores ~1.0). The threshold therefore only applies to INJECTION +// labels; comparing it against the score of a BENIGN result would reject +// virtually everything, since the model is confident on most inputs. func resultFromResponse(r detectResponse, start time.Time, threshold float64) Result { - injected := r.Label == "INJECTION" || r.Score >= threshold + injected := r.Label == "INJECTION" && r.Score >= threshold return Result{ Label: r.Label, Score: r.Score, diff --git a/internal/guard/piguard_test.go b/internal/guard/piguard_test.go index a0fcdc4..ebb267b 100644 --- a/internal/guard/piguard_test.go +++ b/internal/guard/piguard_test.go @@ -7,6 +7,7 @@ import ( "net/http/httptest" "strings" "testing" + "time" ) // fakePiguardServer returns an httptest.Server that mimics the PIGuard HTTP gateway. @@ -29,7 +30,7 @@ func fakePiguardServer(t *testing.T) *httptest.Server { http.Error(w, err.Error(), http.StatusBadRequest) return } - resp := detectResponse{Label: "BENIGN", Score: 0.1} + resp := detectResponse{Label: "BENIGN", Score: 0.9997} // realistic high-confidence BENIGN if strings.Contains(strings.ToLower(req.Text), "ignore") { resp.Label = "INJECTION" resp.Score = 0.999 @@ -42,7 +43,7 @@ func fakePiguardServer(t *testing.T) *httptest.Server { http.Error(w, err.Error(), http.StatusBadRequest) return } - resp := detectResponse{Label: "BENIGN", Score: 0.1} + resp := detectResponse{Label: "BENIGN", Score: 0.9997} // realistic high-confidence BENIGN if strings.Contains(strings.ToLower(req.Long), "ignore") { resp.Label = "INJECTION" resp.Score = 0.999 @@ -57,7 +58,7 @@ func fakePiguardServer(t *testing.T) *httptest.Server { } results := make([]detectResponse, len(req.Texts)) for i, text := range req.Texts { - results[i] = detectResponse{Label: "BENIGN", Score: 0.1} + results[i] = detectResponse{Label: "BENIGN", Score: 0.9997} if strings.Contains(strings.ToLower(text), "ignore") { results[i] = detectResponse{Label: "INJECTION", Score: 0.999} } @@ -224,3 +225,35 @@ func TestPiguardClient_NonOKStatus(t *testing.T) { t.Fatal("expected error for non-OK status") } } + +// TestResultFromResponse_ThresholdSemantics pins the score semantics of the +// PIGuard sidecar: the score is the confidence of the predicted label, not +// the injection probability. A high-confidence BENIGN result (score ~1.0) +// must NOT be treated as an injection, and the threshold only gates INJECTION +// labels. Regression test for the bug where every confident BENIGN memory +// fact was rejected with score >= threshold (default 0.9). +func TestResultFromResponse_ThresholdSemantics(t *testing.T) { + cases := []struct { + name string + label string + score float64 + threshold float64 + want bool + }{ + {"confident benign passes", "BENIGN", 0.9999, 0.9, false}, + {"weak benign passes", "BENIGN", 0.55, 0.9, false}, + {"confident injection rejected", "INJECTION", 0.98, 0.9, true}, + {"injection at threshold rejected", "INJECTION", 0.9, 0.9, true}, + {"weak injection below threshold passes", "INJECTION", 0.62, 0.9, false}, + {"unknown label passes regardless of score", "SOMETHING", 1.0, 0.9, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := resultFromResponse(detectResponse{Label: tc.label, Score: tc.score}, time.Now(), tc.threshold) + if got.Injected != tc.want { + t.Errorf("resultFromResponse(%s, %.4f, thr %.2f).Injected = %v, want %v", + tc.label, tc.score, tc.threshold, got.Injected, tc.want) + } + }) + } +}