bastion 0.53.2

client-go/inferencefilter.go

@@ -0,0 +1,471 @@
+package main
+
+// The inference forwards, with an opinion.
+//
+// serveOllamaForward / serveLLMForward (run.go) are raw TCP splices: bytes in,
+// bytes out, nothing looked at. That is the right shape for an UNCONSTRAINED
+// pairing and it stays the path for one. It cannot be the shape for a pairing
+// whose grant names models, for a reason that is worth stating plainly rather
+// than discovering later: A MODEL NAME LIVES IN THE HTTP BODY. A TCP connection
+// carries no model, so there is nothing an accept-time check could look at —
+// per-model enforcement is only possible one layer in, where the request is.
+//
+// So when a scope is in force this file terminates HTTP on the forward, decides
+// each request on its own, and proxies the rest untouched. Two things are
+// enforced, and they are different:
+//
+//	1. THE MODEL. A request naming a model outside the grant is refused 403 and
+//	   never reaches the upstream. A request on a model-bearing path that names
+//	   NO model is also refused: "which model" is the whole question here, and a
+//	   request that does not answer it cannot be checked (ollama would fall back
+//	   to its own default, which is a model nobody approved).
+//
+//	2. THE BOX'S MODEL STORE. ollama's API on the very same port also serves
+//	   /api/pull, /api/delete, /api/create, /api/copy and /api/push — downloading
+//	   gigabytes onto the user's disk and deleting models off it. Those are the
+//	   `models` caps family's own powers (anchor.go opFamily, modelswap.go), kept
+//	   deliberately separate from `inference` so that "use what is loaded" and
+//	   "decide what is on this machine" are different permissions. Through a raw
+//	   splice that separation was not enforced at all: anything holding
+//	   `inference` could reach them. A pairing that does not hold `models` is now
+//	   refused those paths here, scoped or not.
+//
+// And one thing is filtered on the way BACK: the model LISTS (/api/tags,
+// /api/ps, /v1/models). Same rule the heartbeat follows (agents.go) — offering
+// a model that will then be refused is a lie about reachability, and the
+// friend's host would put it in a picker.
+//
+// Everything outside the allow-list is refused when a scope is in force. Deny
+// by default, like capsAllow: a path nobody has thought about is not a path
+// this lane serves.
+
+import (
+	"bufio"
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"path"
+	"strings"
+	"time"
+)
+
+const (
+	// maxFilteredReqBody bounds a buffered request. An embedding call can carry
+	// a real batch of text, so this is generous; past it the request is refused
+	// rather than forwarded unread, because forwarding unread is how the model
+	// check gets skipped on exactly the request that wanted to skip it.
+	maxFilteredReqBody = 64 << 20
+	// maxFilteredListBody bounds a response we must REWRITE (the model lists).
+	// Only these are buffered; every other response streams through.
+	maxFilteredListBody = 8 << 20
+	// maxFilteredHeaderBytes bounds one request's line + headers.
+	// http.ReadRequest has no cap of its own; ollama's server uses 1 MB.
+	maxFilteredHeaderBytes = 1 << 20
+	filteredIdleTimeout    = 10 * time.Minute
+)
+
+// inferenceLane is which forward is being served — they have different HTTP
+// surfaces, and only ollama's carries the store-mutating half.
+type inferenceLane int
+
+const (
+	laneOllama inferenceLane = iota
+	laneLLM
+)
+
+func (l inferenceLane) String() string {
+	if l == laneLLM {
+		return "llm"
+	}
+	return "ollama"
+}
+
+// storeMutatingPaths are the ollama endpoints that CHANGE what is on this
+// machine. Exact paths, matched after stripping the query. /api/blobs/ is a
+// prefix: it takes a digest in the path and is how a create uploads layers.
+var storeMutatingPaths = map[string]bool{
+	"/api/pull": true, "/api/push": true, "/api/create": true,
+	"/api/copy": true, "/api/delete": true,
+}
+
+// modelBearingPaths name a model in their JSON body. A scoped pairing's request
+// to one of these is read and checked.
+var modelBearingPaths = map[string]bool{
+	"/api/embed": true, "/api/embeddings": true, "/api/generate": true,
+	"/api/chat": true, "/api/show": true,
+	"/v1/embeddings": true, "/v1/chat/completions": true, "/v1/completions": true,
+}
+
+// modelListPaths return a list of models and are rewritten on the way back.
+var modelListPaths = map[string]bool{
+	"/api/tags": true, "/api/ps": true, "/v1/models": true,
+}
+
+// harmlessReadPaths are the remaining reads a scoped pairing may make: they
+// name no model and change nothing.
+var harmlessReadPaths = map[string]bool{
+	"/api/version": true, "/health": true, "/": true,
+}
+
+// inferenceFilterNeeded decides whether this pairing's connections go through
+// the filter or stay a raw splice. Unscoped AND allowed to manage models means
+// there is nothing this file could enforce, so such a pairing (every anchor,
+// and any member granted both) keeps the byte-for-byte path it had before.
+func inferenceFilterNeeded(self string, lane inferenceLane, scope modelScope) bool {
+	if scope.scoped {
+		return true
+	}
+	// The store-mutating half only exists on ollama's port. On an ANCHORED box
+	// it is filtered for every pairing, the anchor included: changing which
+	// models are on a fleet machine is model:pull / model:drop on the
+	// management lane, where it spends the person's approval (writes.go) — a
+	// raw splice that let the anchor POST /api/pull was a way around that.
+	return lane == laneOllama && (currentAnchor() != nil || capsAllow(self, modelGateOp) != "")
+}
+
+// proxyInference serves ONE accepted connection on an inference forward. The
+// scope is read here, per connection, for the same reason inferenceRefuser
+// re-checks caps per connection (run.go): a grant is re-signed while the daemon
+// runs, and a narrowed re-approval has to take effect without a restart.
+//
+// With nothing to enforce it is the original raw splice, byte for byte.
+func proxyInference(in net.Conn, upstream string, lane inferenceLane, self string) {
+	if inferenceFilterNeeded(self, lane, pairingModelScope(self)) {
+		serveInferenceFiltered(in, upstream, lane, self)
+		return
+	}
+	spliceConn(in, upstream)
+}
+
+// spliceConn is what the forwards did before this file existed, unchanged and
+// in one place now that three call sites share it.
+func spliceConn(in net.Conn, upstream string) {
+	defer in.Close()
+	out, err := net.DialTimeout("tcp", upstream, 5*time.Second)
+	if err != nil {
+		return
+	}
+	defer out.Close()
+	done := make(chan struct{}, 2)
+	go func() { _, _ = io.Copy(out, in); done <- struct{}{} }()
+	go func() { _, _ = io.Copy(in, out); done <- struct{}{} }()
+	<-done
+}
+
+// serveInferenceFiltered proxies one accepted connection with this pairing's
+// scope applied. It owns `in` and closes it. Errors are silent by design: this
+// is a data-plane proxy and a peer that hangs up mid-request is normal.
+//
+// It takes no scope ARGUMENT on purpose: the scope is read from the signed
+// grant on every request below, so there is nowhere for a caller — including a
+// test — to hand it one the grant does not say.
+func serveInferenceFiltered(in net.Conn, upstream string, lane inferenceLane, self string) {
+	defer in.Close()
+	out, err := net.DialTimeout("tcp", upstream, 5*time.Second)
+	if err != nil {
+		writeFilterError(in, nil, http.StatusBadGateway, "the service behind this lane is not answering")
+		return
+	}
+	defer out.Close()
+	capped := &capReader{r: in, n: -1}
+	br := bufio.NewReader(capped)
+	bo := bufio.NewReader(out)
+	for {
+		// The deadline covers the HEADERS AND THE BODY, not just the first line:
+		// a peer that sends a request line and then stops would otherwise park
+		// this goroutine on the body read for as long as the process lives.
+		_ = in.SetReadDeadline(time.Now().Add(filteredIdleTimeout))
+		// The request line and headers get their own allowance, because
+		// http.ReadRequest has none: textproto keeps growing a line buffer, so
+		// one approved peer sending a header it never terminates could push
+		// gigabytes into this process. ollama's own server caps this at 1 MB;
+		// the body is bounded separately, below, so the allowance is lifted
+		// once the headers are parsed.
+		capped.allow(maxFilteredHeaderBytes)
+		req, err := http.ReadRequest(br)
+		capped.allow(-1)
+		if err != nil {
+			return
+		}
+		body, err := io.ReadAll(io.LimitReader(req.Body, maxFilteredReqBody+1))
+		req.Body.Close()
+		_ = in.SetReadDeadline(time.Time{})
+		if err != nil {
+			return
+		}
+		if len(body) > maxFilteredReqBody {
+			writeFilterError(in, req, http.StatusRequestEntityTooLarge,
+				"request body too large for a scoped inference lane")
+			return
+		}
+		// Read the grant AFTER the request arrives, never before waiting for
+		// it. A pairing grant is re-signed while the daemon runs — a narrowed
+		// re-approval arrives as a grant_sync that restarts nothing — and a
+		// keep-alive connection under continuous use never reaches its idle
+		// deadline, so a scope read at accept could outlive the approval it
+		// came from for as long as the peer keeps talking. Reading at the TOP
+		// of the loop has the same bug one request smaller: the value is taken
+		// before the blocking read, so it is always the scope as of the
+		// PREVIOUS request. Same reasoning as inferenceRefuser re-checking caps
+		// per connection (run.go), one level finer, and it costs a file read on
+		// a lane that carries a handful of requests a minute.
+		scope := pairingModelScope(self)
+		// Never through the forward on an anchored box — see inferenceFilterNeeded.
+		mayManageModels := currentAnchor() == nil && capsAllow(self, modelGateOp) == ""
+		path := requestPath(req)
+		if why, code := filterVerdict(path, body, lane, scope, mayManageModels); why != "" {
+			if !writeFilterError(in, req, code, why) {
+				return
+			}
+			continue
+		}
+		// Re-frame with the buffered body: Content-Length, never chunked, so
+		// the upstream reads exactly what was checked. A body of zero bytes
+		// becomes a NIL body rather than an empty one — Go's outgoingLength
+		// returns -1 for a non-nil empty Body, which sends the request chunked,
+		// and a chunked GET is a thing some servers simply refuse.
+		if len(body) == 0 {
+			req.Body = nil
+		} else {
+			req.Body = io.NopCloser(bytes.NewReader(body))
+		}
+		req.ContentLength = int64(len(body))
+		req.TransferEncoding = nil
+		req.URL.Scheme, req.URL.Host = "", ""
+		// The body is already here, in full, so a 100-continue handshake can
+		// only cost a round trip — and it would cost more than that: ReadResponse
+		// hands back the interim 1xx as though it were the answer (skipping them
+		// is Transport's job, not its own), so the real response would be left
+		// unread in the buffer and the next request would read it instead.
+		req.Header.Del("Expect")
+		// A list we are about to REWRITE must come back in a form we can read.
+		// The caller is an ordinary HTTP client and Go's transport adds
+		// Accept-Encoding: gzip by itself, so without this the filter would meet
+		// a gzipped body, fail to parse it, and answer 502 for a model list that
+		// was perfectly fine. Dropped only for the paths that get rewritten —
+		// every other response streams through with its encoding untouched.
+		if scope.scoped && modelListPaths[path] {
+			req.Header.Del("Accept-Encoding")
+		}
+		if err := req.Write(out); err != nil {
+			return
+		}
+		resp, err := http.ReadResponse(bo, req)
+		// Belt for the same 1xx problem, in case an upstream sends one
+		// unprompted: skip interim responses rather than forwarding one as the
+		// answer and desyncing the connection. Bounded, so a server stuck
+		// emitting them cannot hold this goroutine.
+		for i := 0; err == nil && resp.StatusCode/100 == 1 && i < 4; i++ {
+			resp.Body.Close()
+			resp, err = http.ReadResponse(bo, req)
+		}
+		if err != nil {
+			writeFilterError(in, req, http.StatusBadGateway, "the service behind this lane closed the connection")
+			return
+		}
+		if scope.scoped && modelListPaths[path] && resp.StatusCode == http.StatusOK {
+			if !writeFilteredList(in, resp, path, scope) {
+				return
+			}
+		} else if err := resp.Write(in); err != nil {
+			resp.Body.Close()
+			return
+		}
+		resp.Body.Close()
+		if resp.Close || req.Close {
+			return
+		}
+	}
+}
+
+// capReader bounds how much one phase of a request may read. `allow(n)` sets
+// the allowance for what comes next; allow(-1) lifts it. It wraps the CONN, not
+// the bufio above it, so the buffering that makes keep-alive work is untouched.
+type capReader struct {
+	r net.Conn
+	n int64 // remaining; <0 = unlimited
+}
+
+func (c *capReader) allow(n int64) { c.n = n }
+
+func (c *capReader) Read(p []byte) (int, error) {
+	if c.n == 0 {
+		return 0, fmt.Errorf("request headers exceeded %d bytes", maxFilteredHeaderBytes)
+	}
+	if c.n > 0 && int64(len(p)) > c.n {
+		p = p[:c.n]
+	}
+	n, err := c.r.Read(p)
+	if c.n > 0 {
+		c.n -= int64(n)
+	}
+	return n, err
+}
+
+// requestPath is the path the POLICY is decided on: lowercased, cleaned, and
+// without a trailing slash. Cleaning matters — `/api/pull/` and `//api/pull`
+// are the same endpoint to a Go/Gin router but three different map keys, and an
+// exact-match deny list that misses two of them is a deny list with two holes.
+// Query is dropped: nothing here is decided on it.
+func requestPath(req *http.Request) string {
+	if req.URL == nil {
+		return ""
+	}
+	p := path.Clean("/" + strings.TrimPrefix(strings.ToLower(req.URL.Path), "/"))
+	if len(p) > 1 {
+		p = strings.TrimSuffix(p, "/")
+	}
+	return p
+}
+
+// filterVerdict is the whole policy, pure and testable: "" to forward, or the
+// refusal sentence and the status to answer with.
+func filterVerdict(path string, body []byte, lane inferenceLane, scope modelScope, mayManageModels bool) (string, int) {
+	if lane == laneOllama && (storeMutatingPaths[path] || strings.HasPrefix(path, "/api/blobs/")) {
+		// SCOPE BEATS THE CAP HERE, and the order is the whole point. A pairing
+		// holding both `models` and a model list could otherwise POST
+		// /api/copy {"source":"llama3:70b","destination":"qwen4-embed"} and use
+		// the result — renaming any model on the box INTO its scope. A ceiling
+		// you can walk around by renaming is not a ceiling, so a scoped pairing
+		// does not manage models, whatever else it was granted.
+		if scope.scoped {
+			return "refused: " + path + " changes which models are on this machine, and this " +
+				"pairing is scoped to named models (" + scope.describe() + ") — which a rename " +
+				"would walk straight around.", http.StatusForbidden
+		}
+		if mayManageModels {
+			return "", 0
+		}
+		return "refused: " + path + " changes which models are on this machine. That is the " +
+			"'models' capability, which this pairing was not granted — 'inference' is the right " +
+			"to USE what is already loaded.", http.StatusForbidden
+	}
+	if !scope.scoped {
+		return "", 0
+	}
+	if modelListPaths[path] || harmlessReadPaths[path] {
+		return "", 0
+	}
+	if modelBearingPaths[path] {
+		name := modelNameFromBody(body)
+		if name == "" {
+			return "refused: this pairing may use " + scope.describe() + " on this machine, so every " +
+				"request has to name the model it wants; this one names none.", http.StatusForbidden
+		}
+		if !scope.allowsModel(name) {
+			return "refused: model '" + name + "' is not shared with this host. This pairing may use " +
+				scope.describe() + " on this machine.", http.StatusForbidden
+		}
+		return "", 0
+	}
+	return "refused: " + path + " is not served to a pairing scoped to named models (" +
+		scope.describe() + ").", http.StatusForbidden
+}
+
+// modelNameFromBody pulls the model out of a request body. `model` is what both
+// dialects use today; `name` is ollama's older spelling, still accepted by
+// /api/show and /api/delete, so a check that only read `model` would wave
+// through a request that named its model the old way.
+func modelNameFromBody(body []byte) string {
+	var b struct {
+		Model string `json:"model"`
+		Name  string `json:"name"`
+	}
+	if json.Unmarshal(body, &b) != nil {
+		return ""
+	}
+	if s := strings.TrimSpace(b.Model); s != "" {
+		return s
+	}
+	return strings.TrimSpace(b.Name)
+}
+
+// writeFilteredList rewrites a model-list response to the models this scope
+// allows. A list too large to buffer is answered 502 rather than passed
+// through: an unfiltered list is the one thing this function exists to prevent.
+// Returns false when the connection should be dropped.
+func writeFilteredList(in net.Conn, resp *http.Response, path string, scope modelScope) bool {
+	raw, err := io.ReadAll(io.LimitReader(resp.Body, maxFilteredListBody+1))
+	if err != nil || len(raw) > maxFilteredListBody {
+		writeFilterError(in, nil, http.StatusBadGateway, "model list too large to filter")
+		return false
+	}
+	out, err := filterModelListJSON(raw, path, scope)
+	if err != nil {
+		writeFilterError(in, nil, http.StatusBadGateway, "model list could not be filtered")
+		return false
+	}
+	resp.Body = io.NopCloser(bytes.NewReader(out))
+	resp.ContentLength = int64(len(out))
+	resp.TransferEncoding = nil
+	resp.Header.Set("Content-Length", fmt.Sprint(len(out)))
+	resp.Header.Del("Content-Encoding") // we rewrote the bytes; any encoding is now wrong
+	return resp.Write(in) == nil
+}
+
+// filterModelListJSON keeps every field of every entry the scope allows and
+// drops the rest. Deliberately re-marshals the ORIGINAL entries rather than
+// rebuilding them, so fields this code has never heard of (sizes, digests,
+// capabilities, quantization) survive untouched.
+func filterModelListJSON(raw []byte, path string, scope modelScope) ([]byte, error) {
+	var doc map[string]json.RawMessage
+	if err := json.Unmarshal(raw, &doc); err != nil {
+		return nil, err
+	}
+	listKey, nameKeys := "models", []string{"name", "model"}
+	if path == "/v1/models" {
+		listKey, nameKeys = "data", []string{"id"}
+	}
+	rawList, ok := doc[listKey]
+	if !ok {
+		return raw, nil
+	}
+	var entries []json.RawMessage
+	if err := json.Unmarshal(rawList, &entries); err != nil {
+		return nil, err
+	}
+	kept := []json.RawMessage{}
+	for _, e := range entries {
+		var fields map[string]any
+		if json.Unmarshal(e, &fields) != nil {
+			continue // an entry we cannot read is an entry we cannot vouch for
+		}
+		name := ""
+		for _, k := range nameKeys {
+			if s, _ := fields[k].(string); strings.TrimSpace(s) != "" {
+				name = s
+				break
+			}
+		}
+		if scope.allowsModel(name) {
+			kept = append(kept, e)
+		}
+	}
+	nl, err := json.Marshal(kept)
+	if err != nil {
+		return nil, err
+	}
+	doc[listKey] = nl
+	return json.Marshal(doc)
+}
+
+// writeFilterError answers one request with a refusal the caller can read. JSON
+// with an `error` string: both dialects surface that field, so the sentence
+// reaches a person instead of becoming "unexpected status 403". Returns false
+// when the connection cannot continue.
+func writeFilterError(in net.Conn, req *http.Request, code int, why string) bool {
+	body, _ := json.Marshal(map[string]string{"error": why})
+	resp := &http.Response{
+		StatusCode:    code,
+		ProtoMajor:    1,
+		ProtoMinor:    1,
+		Header:        http.Header{"Content-Type": []string{"application/json"}},
+		Body:          io.NopCloser(bytes.NewReader(body)),
+		ContentLength: int64(len(body)),
+		Request:       req,
+	}
+	return resp.Write(in) == nil
+}