diff --git a/guard/guard.go b/guard/guard.go index c524bab..0307b7b 100644 --- a/guard/guard.go +++ b/guard/guard.go @@ -162,7 +162,7 @@ func (g *Guard) applyIntercept(ctx context.Context, ev hook.Event) hook.Result { if gated { var res intercept.TargetResult if unpinned, res = g.gate.CheckTargets(ctx, scope, ev.ToolName, ev.Input); res.Code != "" { - return g.block(ev.ToolName, targetBlockMessage(res), "") + return g.block(ev.ToolName, TargetBlockMessage(res), "") } } var ask *intercept.Decision @@ -182,7 +182,7 @@ func (g *Guard) applyIntercept(ctx context.Context, ev hook.Event) hook.Result { if len(unpinned) > 0 { res, merged := g.requestTargets(ctx, scope, ev, unpinned, ask) if res.Code != "" { - return g.block(ev.ToolName, targetBlockMessage(res), "") + return g.block(ev.ToolName, TargetBlockMessage(res), "") } if merged { // The approval of the target request answered the ask as well. @@ -235,11 +235,11 @@ var targetGuidance = map[string]string{ "Retry the same call later, or report the failure.", } -// targetBlockMessage is systemBlockMessage for the target gate: the block code, +// TargetBlockMessage is systemBlockMessage for the target gate: the block code, // the reason and the target, then what to do about it. A target block is not // always final (an approval can still arrive, a lookup can succeed later), so // it does not end with "This operation is forbidden." -func targetBlockMessage(res intercept.TargetResult) string { +func TargetBlockMessage(res intercept.TargetResult) string { reason := res.Reason if res.Host != "" && !strings.Contains(reason, res.Host) { target := res.Host diff --git a/guard/target_gate_test.go b/guard/target_gate_test.go index 1859c1a..ec61a86 100644 --- a/guard/target_gate_test.go +++ b/guard/target_gate_test.go @@ -257,7 +257,7 @@ func TestTargetBlockMessage(t *testing.T) { []string{"TARGET_GATE_UNAVAILABLE: target app.example.com:443: the pins could not be read: connection refused.", "failure on the platform side"}}, } for _, c := range cases { - msg := targetBlockMessage(c.res) + msg := TargetBlockMessage(c.res) if !strings.HasPrefix(msg, "[ARTEX platform control, not target defense] "+c.res.Code+": ") { t.Errorf("%s: message = %q", c.res.Code, msg) } diff --git a/intercept/target_rendered.go b/intercept/target_rendered.go new file mode 100644 index 0000000..67921b0 --- /dev/null +++ b/intercept/target_rendered.go @@ -0,0 +1,154 @@ +package intercept + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/url" + "strings" +) + +// The custom tools render their target from a template and the parameters of +// the call, after the hook has seen the call: the hook reads the parameters, not +// the URL or the command they produce, and not a target written in the +// template itself. These functions run the gate inside the tool, on what the +// tool is about to run, before anything is sent or started. + +// renderedPasses bounds the request-and-recheck loop of a rendered target, as +// the hook bounds its own: a request that was decided in between, or addresses +// that changed again after an approval, send the call back to the pin lookup at +// most this many times. +const renderedPasses = 3 + +// GateRenderedURL is the target gate for the URL a custom http tool renders. It +// reads the target as the HTTP client will: the host of the parsed URL and its +// port, else the default port of the scheme. A URL whose target cannot be read +// is blocked with TARGET_GATE_UNAVAILABLE: the call is not executed and no +// approval is requested. +// +// The result is empty when the tool may run. With the gate off, or with no +// scope in ctx (the operator's test endpoint), it returns at once: nothing is +// parsed, resolved, read or written. +func GateRenderedURL(ctx context.Context, toolName, rawURL string) TargetResult { + scope, on := renderedScope(ctx) + if !on { + return TargetResult{} + } + input := renderedInput("url", rawURL) + t, err := urlTarget(rawURL) + if err != nil { + return undeterminedTarget(ctx, scope, toolName, input, "the target of the rendered URL could not be read ("+err.Error()+")") + } + return gateRendered(ctx, scope, toolName, input, []Target{t}) +} + +// GateRenderedCommand is the target gate for the command a custom command tool +// renders, read with the command-string rules of Bash. readings are other +// renderings of the same command that are read as well: the shell quoting of a +// parameter placed inside a host name hides that name from the rules, and the +// shell removes the quotes before the command runs. A command that names no +// target passes, as it does in Bash: what a command hides is advisory (D9). +// +// The result is empty when the tool may run. With the gate off, or with no +// scope in ctx, it returns at once. +func GateRenderedCommand(ctx context.Context, toolName, command string, readings ...string) TargetResult { + scope, on := renderedScope(ctx) + if !on { + return TargetResult{} + } + var l targetList + l.scan(command) + for _, r := range readings { + l.scan(r) + } + return gateRendered(ctx, scope, toolName, renderedInput("command", command), l.targets) +} + +func renderedScope(ctx context.Context) (TargetScope, bool) { + scope, ok := TargetScopeFromContext(ctx) + return scope, ok && scope.Enabled != nil && scope.Enabled() +} + +// renderedInput is the tool input recorded with a rendered target: what the +// tool runs, which is what the human reads on the approval. +func renderedInput(field, value string) []byte { + b, _ := json.Marshal(map[string]string{field: value}) + return b +} + +// gateRendered runs the gate decision on the targets of a rendered call, asks +// for the approval of those no pin covers, and checks again after the +// decisions. It never merges an ask: the rules and the judge decided at the +// hook, before the tool ran. +func gateRendered(ctx context.Context, scope TargetScope, toolName string, input []byte, targets []Target) TargetResult { + if len(targets) == 0 { + return TargetResult{} + } + i := scope.Interceptor + if i == nil { + return TargetResult{Code: TargetGateUnavailable, Host: targets[0].Host, Port: targets[0].Port, + Reason: "the session's target scope carries no gate"} + } + ctx = WithCall(ctx, toolName, input) + check := func() ([]TargetRequest, TargetResult) { + if !scope.Enabled() { + return nil, TargetResult{} + } + c := &targetCheck{i: i, ctx: ctx, scope: scope, tool: toolName, input: input} + return c.run(targets) + } + unpinned, res := check() + for pass := 0; res.Code == "" && len(unpinned) > 0; pass++ { + if pass == renderedPasses { + t := unpinned[0] + return TargetResult{Code: TargetApprovalPending, Host: t.Host, Port: t.Port, + Reason: fmt.Sprintf("the target is still not covered by a pin after %d approval rounds; its addresses changed again", renderedPasses)} + } + if res = i.RequestTargets(ctx, scope, toolName, input, unpinned, nil); res.Code != "" { + return res + } + unpinned, res = check() + } + return res +} + +// undeterminedTarget blocks a call whose target cannot be read, and logs the +// block as a decided row. +func undeterminedTarget(ctx context.Context, scope TargetScope, toolName string, input []byte, reason string) TargetResult { + res := TargetResult{Code: TargetGateUnavailable, Reason: reason} + if scope.Interceptor == nil { + return res + } + ref, err := scope.ref() + if err != nil { + return res + } + c := &targetCheck{i: scope.Interceptor, ctx: WithCall(ctx, toolName, input), scope: scope, ref: ref, tool: toolName, input: input} + _, res = c.block(TargetGateUnavailable, "gate_unavailable", &checkedTarget{}, reason) + return res +} + +// urlTarget reads the target of a URL the way net/http dials it: the host of +// the parsed URL, and its port, else the default port of the scheme. +func urlTarget(raw string) (Target, error) { + u, err := url.Parse(raw) + if err != nil { + return Target{}, err + } + name := u.Hostname() + if name == "" { + return Target{}, errors.New("no host") + } + kind, host, ok := normalizeHost(name, true) + if !ok { + return Target{}, fmt.Errorf("host %q", name) + } + t := Target{Kind: kind, Host: host, Port: schemePorts[strings.ToLower(u.Scheme)]} + if p := u.Port(); p != "" { + if t.Port = validPort(p); t.Port == 0 { + return Target{}, fmt.Errorf("port %q", p) + } + } + return t, nil +} diff --git a/intercept/target_rendered_test.go b/intercept/target_rendered_test.go new file mode 100644 index 0000000..0486c1d --- /dev/null +++ b/intercept/target_rendered_test.go @@ -0,0 +1,177 @@ +package intercept + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" + + "github.com/Autumn-27/artex/db" +) + +// The gate inside the custom tools: GateRenderedURL and GateRenderedCommand +// read what the tool is about to run. Without a database: the cases that reach +// a request are covered by the server tests. + +func (g *gateFixture) ctx() context.Context { + return WithTargetScope(context.Background(), g.scope) +} + +// With the gate off, or without a scope (the operator's test endpoint), the +// rendered target is not even parsed: nothing is read, resolved or logged, and +// a URL the gate could not read passes too, as it did before the gate. +func TestRenderedGateOffTouchesNothing(t *testing.T) { + g := newGate(t) + g.res.set("app.example.com", "203.0.113.7") + for name, ctx := range map[string]context.Context{ + "no scope": context.Background(), + "no flag": WithTargetScope(context.Background(), TargetScope{TaskID: "7", WaitBound: time.Minute, Interceptor: g.ic}), + "flag off": WithTargetScope(context.Background(), TargetScope{TaskID: "7", Enabled: func() bool { return false }, Interceptor: g.ic}), + } { + for _, res := range []TargetResult{ + GateRenderedURL(ctx, "probe", appURL), + GateRenderedURL(ctx, "probe", "http://a b/"), + GateRenderedCommand(ctx, "probe", "curl "+appURL), + } { + if res != (TargetResult{}) { + t.Errorf("%s: %+v, want nothing", name, res) + } + } + } + if g.store.reads != 0 || len(g.store.logged) != 0 || g.res.total() != 0 { + t.Errorf("gate off: %d store reads, %d rows, %d lookups, want none", g.store.reads, len(g.store.logged), g.res.total()) + } +} + +// The target of a rendered URL is read as net/http dials it: the host of the +// parsed URL, the explicit port, else the default of the scheme. +func TestRenderedURLTarget(t *testing.T) { + for _, c := range []struct { + raw string + want Target + }{ + {"https://App.Example.COM./x?next=http://other.example/", Target{TargetDomain, "app.example.com", 443}}, + {"http://user:pw@203.0.113.7:8080/a", Target{TargetIP, "203.0.113.7", 8080}}, + {"http://[::ffff:203.0.113.7]/", Target{TargetIP, "203.0.113.7", 80}}, + {"http://[2001:db8::1]:8443/", Target{TargetIP, "2001:db8::1", 8443}}, + {"http://2130706433/", Target{TargetIP, "127.0.0.1", 80}}, + {"HTTPS://bücher.example/", Target{TargetDomain, "xn--bcher-kva.example", 443}}, + // The client's reading of the authority, not the shell's: a.example; is userinfo. + {"http://a.example;@b.example/", Target{TargetDomain, "b.example", 80}}, + } { + got, err := urlTarget(c.raw) + if err != nil || got != c.want { + t.Errorf("%q: %+v, %v; want %+v", c.raw, got, err, c.want) + } + } + for _, raw := range []string{ + "http://a b/", // not a URL the client accepts + `http://a.example\@b.example/`, // idem + "http://[::1/", // idem + "http:///path", // no host + "app.example.com/path", // no scheme: no host + "http://app.example.com:99999/x", // no valid port + } { + if got, err := urlTarget(raw); err == nil { + t.Errorf("%q: %+v, want no target", raw, got) + } + } +} + +// A rendered URL whose target cannot be read is blocked: no lookup, no +// request, one decided row that records the URL. +func TestRenderedURLWithoutTargetFailsClosed(t *testing.T) { + g := newGate(t) + for n, raw := range []string{"http://a b/", "http:///path", "app.example.com/path", "http://app.example.com:99999/"} { + res := GateRenderedURL(g.ctx(), "probe_http", raw) + if res.Code != TargetGateUnavailable || !strings.Contains(res.Reason, "could not be read") { + t.Fatalf("%q: %+v, want %s", raw, res, TargetGateUnavailable) + } + if len(g.store.logged) != n+1 { + t.Fatalf("%q: %d rows, want %d", raw, len(g.store.logged), n+1) + } + row := g.store.logged[n] + var input map[string]string + _ = json.Unmarshal(row.ToolInput, &input) + if row.Status != "denied" || row.ToolName != "probe_http" || input["url"] != raw || row.Audit.Targets[0].Outcome != "gate_unavailable" { + t.Fatalf("%q: row = %+v (input %s)", raw, row, row.ToolInput) + } + } + if g.res.total() != 0 { + t.Fatalf("%d lookups, want none", g.res.total()) + } +} + +// A gated scope without the interceptor that owns the gate cannot decide: the +// call is blocked, never let through. +func TestRenderedGateWithoutInterceptorFailsClosed(t *testing.T) { + ctx := WithTargetScope(context.Background(), TargetScope{TaskID: "7", Enabled: func() bool { return true }}) + for name, res := range map[string]TargetResult{ + "url": GateRenderedURL(ctx, "probe", appURL), + "unreadable url": GateRenderedURL(ctx, "probe", "http://a b/"), + "command": GateRenderedCommand(ctx, "probe", "curl "+appURL), + } { + if res.Code != TargetGateUnavailable { + t.Errorf("%s: %+v, want %s", name, res, TargetGateUnavailable) + } + } + // A command that names no target needs no gate. + if res := GateRenderedCommand(ctx, "probe", "ls -la"); res != (TargetResult{}) { + t.Errorf("no target: %+v, want nothing", res) + } +} + +// The rendered target goes through the same decision as a target at the hook: +// block rules, resolution, pins. A pinned target passes, a blocked one is +// blocked without a request. +func TestRenderedTargetUsesTheGateDecision(t *testing.T) { + g := newGate(t) + g.res.set("app.example.com", "203.0.113.7") + g.store.approve(TargetRequest{Host: "app.example.com", Port: 443, ResolvedIPs: []string{"203.0.113.7"}}) + g.store.approve(TargetRequest{Host: "198.51.100.9", Port: 8080, ResolvedIPs: []string{"198.51.100.9"}}) + + for name, res := range map[string]TargetResult{ + "url": GateRenderedURL(g.ctx(), "probe", appURL), + "url, ip and port": GateRenderedURL(g.ctx(), "probe", "http://198.51.100.9:8080/x"), + "command": GateRenderedCommand(g.ctx(), "probe", "curl -s '"+appURL+"'"), + "command, no target": GateRenderedCommand(g.ctx(), "probe", "nmap -sV app.example.com"), + "command, ip literal": GateRenderedCommand(g.ctx(), "probe", "nc 198.51.100.9:8080"), + } { + if res != (TargetResult{}) { + t.Errorf("%s: %+v, want the call to pass", name, res) + } + } + + g.res.errs["down.example"] = context.DeadlineExceeded + if res := GateRenderedURL(g.ctx(), "probe", "https://down.example/"); res.Code != TargetResolveFailed || res.Host != "down.example" { + t.Fatalf("unresolvable: %+v", res) + } + g.store.rules = []db.AssetInterceptRule{{Enabled: true, Kind: "exact_domain", Pattern: "app.example.com"}} + for name, res := range map[string]TargetResult{ + "url": GateRenderedURL(g.ctx(), "probe", appURL), + "command": GateRenderedCommand(g.ctx(), "probe", "curl "+appURL), + } { + if res.Code != TargetRuleBlocked || res.Host != "app.example.com" || res.Port != 443 { + t.Errorf("%s: %+v, want %s on the pinned host", name, res, TargetRuleBlocked) + } + } + if !strings.Contains(strings.Join(g.store.outcomes(), " "), "denied/rule_blocked") { + t.Fatalf("rows = %v, want the block logged", g.store.outcomes()) + } +} + +// The shell quoting of a parameter inside a host name hides the name from the +// command-string rules; the reading without the quoting is read as well. +func TestRenderedCommandReadsTheUnquotedRendering(t *testing.T) { + g := newGate(t) + g.store.rules = []db.AssetInterceptRule{{Enabled: true, Kind: "exact_domain", Pattern: "api.example.com"}} + quoted := "curl http://'api'.example.com:8080/" + if res := GateRenderedCommand(g.ctx(), "probe", quoted); res != (TargetResult{}) { + t.Fatalf("quoted only: %+v; the rules were expected to miss the name", res) + } + res := GateRenderedCommand(g.ctx(), "probe", quoted, "curl http://api.example.com:8080/") + if res.Code != TargetRuleBlocked || res.Host != "api.example.com" || res.Port != 8080 { + t.Fatalf("with the unquoted reading: %+v, want %s on api.example.com:8080", res, TargetRuleBlocked) + } +} diff --git a/server/customtool.go b/server/customtool.go index ba29f7f..5ece010 100644 --- a/server/customtool.go +++ b/server/customtool.go @@ -18,6 +18,8 @@ import ( "time" "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/guard" + "github.com/Autumn-27/artex/intercept" "github.com/Autumn-27/norma/permission" actool "github.com/Autumn-27/norma/tool" ) @@ -163,11 +165,11 @@ func (s *Server) pgTestCustomTool(w http.ResponseWriter, r *http.Request) { var res actool.Result switch req.Kind { case "command": - res, _ = s.runCommandTool(ctx, req.Exec, params, tc) + res, _ = s.runCommandTool(ctx, "test", req.Exec, params, tc) case "script": res, _ = s.runScriptTool(ctx, "test", req.Exec, params, tc) case "http": - res, _ = s.runHTTPTool(ctx, req.Exec, params, tc) + res, _ = s.runHTTPTool(ctx, "test", req.Exec, params, tc) case "shell": writeErr(w, 400, "shell-kind tools are bash environment declarations with no executable content") return @@ -275,11 +277,11 @@ func (s *Server) buildCustomTool(t *db.Tool) actool.CoreTool { } switch kind { case "command": - return s.runCommandTool(ctx, execRaw, params, tc) + return s.runCommandTool(ctx, key, execRaw, params, tc) case "script": return s.runScriptTool(ctx, key, execRaw, params, tc) case "http": - return s.runHTTPTool(ctx, execRaw, params, tc) + return s.runHTTPTool(ctx, key, execRaw, params, tc) default: return actool.Errorf("unknown custom tool kind: " + kind), nil } @@ -328,13 +330,18 @@ func ensureSchema(raw json.RawMessage) map[string]any { // ---------- command: render the command → reuse the Bash low-level run ---------- -func (s *Server) runCommandTool(ctx context.Context, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) { +func (s *Server) runCommandTool(ctx context.Context, key string, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) { var spec commandExec _ = json.Unmarshal(execRaw, &spec) if strings.TrimSpace(spec.Command) == "" { return actool.Errorf("command is empty"), nil } cmd := renderTemplate(spec.Command, params, shellQuote) + // The hook saw the parameters, not the command they render: the target gate + // reads the rendered command here, also without the quoting of the parameters. + if res := intercept.GateRenderedCommand(ctx, key, cmd, renderTemplate(spec.Command, params, identity)); res.Code != "" { + return actool.Errorf(guard.TargetBlockMessage(res)), nil + } // Reuse the low-level run Bash also uses (via the Bash CoreTool.Call): it // automatically inherits the safety floor/timeout/proxy env/output overflow. The tool is a peer of Bash sharing that low level, not routed through the Bash tool for the model to call. bashIn, _ := json.Marshal(map[string]any{"command": cmd}) @@ -417,7 +424,7 @@ func execPython(ctx context.Context, interp, key, code string, params map[string // ---------- http: native request + proxy ---------- -func (s *Server) runHTTPTool(ctx context.Context, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) { +func (s *Server) runHTTPTool(ctx context.Context, key string, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) { var spec httpExec _ = json.Unmarshal(execRaw, &spec) method := strings.ToUpper(strings.TrimSpace(spec.Method)) @@ -428,6 +435,12 @@ func (s *Server) runHTTPTool(ctx context.Context, execRaw json.RawMessage, param if strings.TrimSpace(rawURL) == "" { return actool.Errorf("http url is empty"), nil } + // The hook saw the parameters, not the URL they render (nor a URL written in + // the template): the target gate reads the rendered URL here, before the + // tool's own timeout starts, so a wait for an approval does not consume it. + if res := intercept.GateRenderedURL(ctx, key, rawURL); res.Code != "" { + return actool.Errorf(guard.TargetBlockMessage(res)), nil + } var bodyReader io.Reader if spec.Body != "" { bodyReader = strings.NewReader(renderTemplate(spec.Body, params, identity)) diff --git a/server/customtool_target_gate_test.go b/server/customtool_target_gate_test.go new file mode 100644 index 0000000..47b1723 --- /dev/null +++ b/server/customtool_target_gate_test.go @@ -0,0 +1,619 @@ +package server + +import ( + "context" + "encoding/json" + "fmt" + "iter" + "net" + "net/http" + "net/http/httptest" + "net/netip" + "net/url" + "os" + "path/filepath" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/intercept" + "github.com/Autumn-27/norma/llm" + actool "github.com/Autumn-27/norma/tool" +) + +// Step 1.8: the targets the custom tools render. The hook reads the parameters +// of a custom tool, not the URL or the command they produce: the tool runs the +// gate on what it is about to run. Plus the end-to-end run of a gated worker +// and the run of a task with the gate off. + +// countingServer is an HTTP server on 127.0.0.1 that counts the requests it +// receives. Its host:port is an IP literal: the gate never resolves it. +func countingServer(t *testing.T) (srv *httptest.Server, hits *atomic.Int64, host, port string) { + t.Helper() + hits = &atomic.Int64{} + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.Header().Set("Content-Type", "text/plain") + _, _ = w.Write([]byte("probe-reached")) + })) + t.Cleanup(srv.Close) + u, err := url.Parse(srv.URL) + if err != nil { + t.Fatal(err) + } + host, port, err = net.SplitHostPort(u.Host) + if err != nil { + t.Fatal(err) + } + return srv, hits, host, port +} + +// gatedCtx is the context of a worker tool call of task: its approval scope +// with the given wait bound. +func gatedCtx(task *Task, wait time.Duration) context.Context { + ctx := intercept.WithTaskContext(context.Background(), task.ID, "worker", nil) + return intercept.WithTargetScope(ctx, task.targetScope(wait)) +} + +func execJSON(t *testing.T, v any) json.RawMessage { + t.Helper() + b, err := json.Marshal(v) + if err != nil { + t.Fatal(err) + } + return b +} + +// callTool runs a custom tool the way the harness does and returns its text. +func callTool(t *testing.T, tool actool.CoreTool, ctx context.Context, input string, tc *actool.ToolContext) (string, bool) { + t.Helper() + res, err := tool.Call(ctx, json.RawMessage(input), tc) + if err != nil { + t.Fatalf("%s: %v", tool.Name(), err) + } + return res.Flatten(), res.IsError +} + +// openTargetRequest returns the open target request of the task on key. +func openTargetRequest(t *testing.T, pg *db.DB, taskID, key string) (id int64, toolName, toolInput string, ok bool) { + t.Helper() + err := pg.QueryRow(`SELECT id, tool_name, tool_input::text FROM intercept_pending +WHERE task_id=$1 AND status='pending' AND target_key=$2`, taskID, key).Scan(&id, &toolName, &toolInput) + if err != nil { + return 0, "", "", false + } + return id, toolName, toolInput, true +} + +func targetRowCount(t *testing.T, pg *db.DB, taskID string) int { + t.Helper() + var n int + if err := pg.QueryRow(`SELECT count(*) FROM intercept_pending WHERE task_id=$1 AND (decision_source='target' OR target_key IS NOT NULL)`, taskID).Scan(&n); err != nil { + t.Fatal(err) + } + return n +} + +// A custom http tool is gated on the URL it renders, whether the target comes +// from a parameter (here split over two parameters, so the hook reads no +// target in them) or is written in the template. Nothing is sent before an +// approval; an approval lets both through; a revoked pin asks again. +func TestCustomHTTPRenderedTargetGated(t *testing.T) { + g := newGateServer(t) + on := true + task := g.newTask(t, db.TaskCreateOptions{TargetGate: &on}) + pg := g.m.pg + srv, hits, host, port := countingServer(t) + key := net.JoinHostPort(host, port) + dot := strings.LastIndexByte(host, '.') + + rendered := g.s.buildCustomTool(&db.Tool{Key: "probe_http", Kind: "http", + Schema: json.RawMessage(`{"type":"object","properties":{"last":{"type":"string"},"port":{"type":"string"}}}`), + Exec: execJSON(t, map[string]any{"method": "GET", "url": "http://" + host[:dot+1] + "{last}:{port}/rendered"})}) + params := fmt.Sprintf(`{"last":%q,"port":%q}`, host[dot+1:], port) + fixed := g.s.buildCustomTool(&db.Tool{Key: "probe_fixed", Kind: "http", + Exec: execJSON(t, map[string]any{"method": "GET", "url": srv.URL + "/fixed"})}) + // The hook cannot see either target: neither input names one. + for _, in := range []struct{ tool, input string }{{"probe_http", params}, {"probe_fixed", `{}`}} { + if got := intercept.ExtractTargets(in.tool, []byte(in.input)); len(got) != 0 { + t.Fatalf("%s: the hook reads %+v; the test needs a target only the rendering shows", in.tool, got) + } + } + + // The template's own URL, wait bound 0: blocked at once, nothing sent, the + // request open with the rendered URL as its input. + out, isErr := callTool(t, fixed, gatedCtx(task, 0), `{}`, nil) + if !isErr || !strings.Contains(out, "TARGET_APPROVAL_PENDING") || !strings.Contains(out, key) || hits.Load() != 0 { + t.Fatalf("hard-coded URL: error=%v out=%q hits=%d", isErr, out, hits.Load()) + } + id, toolName, toolInput, ok := openTargetRequest(t, pg, task.ID, key) + if !ok || toolName != "probe_fixed" || !strings.Contains(toolInput, srv.URL+"/fixed") { + t.Fatalf("request row: ok=%v tool=%q input=%s", ok, toolName, toolInput) + } + + // The rendered URL names the same host and port: the call attaches to the + // open request and waits. One click lets it through. + type result struct { + out string + isErr bool + } + done := make(chan result, 1) + go func() { + res, err := rendered.Call(gatedCtx(task, 30*time.Second), json.RawMessage(params), nil) + if err != nil { + done <- result{err.Error(), true} + return + } + done <- result{res.Flatten(), res.IsError} + }() + time.Sleep(300 * time.Millisecond) + if hits.Load() != 0 { + t.Fatal("the rendered URL was fetched before the approval") + } + if n := targetRowCount(t, pg, task.ID); n != 1 { + t.Fatalf("%d target rows, want the one request", n) + } + if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"allowed"}`); code != 200 { + t.Fatalf("decide: %d %v", code, body) + } + select { + case r := <-done: + if r.isErr || !strings.Contains(r.out, "probe-reached") || hits.Load() != 1 { + t.Fatalf("after the approval: error=%v out=%q hits=%d", r.isErr, r.out, hits.Load()) + } + case <-time.After(30 * time.Second): + t.Fatal("the waiting call did not end after the approval") + } + if out, isErr := callTool(t, fixed, gatedCtx(task, 0), `{}`, nil); isErr || hits.Load() != 2 { + t.Fatalf("pinned, hard-coded URL: error=%v out=%q hits=%d", isErr, out, hits.Load()) + } + + // Revoked: the next call asks again and sends nothing. + pins, err := pg.ListTargetPins(taskRef(t, task)) + if err != nil || len(pins) != 1 || pins[0].Host != host || fmt.Sprint(pins[0].Port) != port { + t.Fatalf("pins = %+v, %v", pins, err) + } + if code, _ := g.do(t, "POST", fmt.Sprintf("/api/tasks/%s/target-pins/%d/revoke", task.ID, pins[0].ID), ""); code != 200 { + t.Fatalf("revoke: %d", code) + } + out, isErr = callTool(t, rendered, gatedCtx(task, 0), params, nil) + if !isErr || !strings.Contains(out, "TARGET_APPROVAL_PENDING") || hits.Load() != 2 { + t.Fatalf("revoked: error=%v out=%q hits=%d", isErr, out, hits.Load()) + } + if again, toolName, _, ok := openTargetRequest(t, pg, task.ID, key); !ok || again == id || toolName != "probe_http" { + t.Fatalf("after the revoke: open request %d (ok=%v, tool %q), want a new one", again, ok, toolName) + } + + // A URL whose target cannot be read is blocked, never sent. + broken := g.s.buildCustomTool(&db.Tool{Key: "probe_broken", Kind: "http", + Schema: json.RawMessage(`{"type":"object","properties":{"host":{"type":"string"}}}`), + Exec: execJSON(t, map[string]any{"url": "http://{host}/x"})}) + out, isErr = callTool(t, broken, gatedCtx(task, 0), `{"host":"a b"}`, nil) + if !isErr || !strings.Contains(out, "TARGET_GATE_UNAVAILABLE") { + t.Fatalf("unreadable URL: error=%v out=%q", isErr, out) + } +} + +// A custom command tool is gated on the command it renders, with a target +// written in the template or formed by a parameter inside a host name. The +// command does not run before an approval. +func TestCustomCommandRenderedTargetGated(t *testing.T) { + g := newGateServer(t) + on := true + task := g.newTask(t, db.TaskCreateOptions{TargetGate: &on}) + pg := g.m.pg + res := &countingResolver{answers: map[string][]string{"api.example.com": {"198.51.100.7"}}} + g.m.interceptor.SetResolver(res) + dir := t.TempDir() + marker := filepath.Join(dir, "marker") + ran := func() bool { + _, err := os.Stat(marker) + return err == nil + } + + fixed := g.s.buildCustomTool(&db.Tool{Key: "probe_cmd_fixed", Kind: "command", + Exec: execJSON(t, map[string]any{"command": "echo reached > '" + marker + "' && : http://198.51.100.7:8080/"})}) + glued := g.s.buildCustomTool(&db.Tool{Key: "probe_cmd", Kind: "command", + Schema: json.RawMessage(`{"type":"object","properties":{"sub":{"type":"string"}}}`), + Exec: execJSON(t, map[string]any{"command": "echo reached > '" + marker + "' && : http://{sub}.example.com:8080/"})}) + if got := intercept.ExtractTargets("probe_cmd", []byte(`{"sub":"api"}`)); len(got) != 0 { + t.Fatalf("the hook reads %+v; the test needs a target only the rendering shows", got) + } + tc := &actool.ToolContext{WorkingDir: dir} + + for _, c := range []struct { + name string + tool actool.CoreTool + input string + key string + }{ + {"template", fixed, `{}`, "198.51.100.7:8080"}, + {"parameter inside the host", glued, `{"sub":"api"}`, "api.example.com:8080"}, + } { + out, isErr := callTool(t, c.tool, gatedCtx(task, 0), c.input, tc) + if !isErr || !strings.Contains(out, "TARGET_APPROVAL_PENDING") || ran() { + t.Fatalf("%s: error=%v out=%q ran=%v, want a block before the command runs", c.name, isErr, out, ran()) + } + id, toolName, _, ok := openTargetRequest(t, pg, task.ID, c.key) + if !ok || toolName != c.tool.Name() { + t.Fatalf("%s: no open request on %s for %s", c.name, c.key, c.tool.Name()) + } + if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"allowed"}`); code != 200 { + t.Fatalf("%s: decide: %d %v", c.name, code, body) + } + if out, isErr := callTool(t, c.tool, gatedCtx(task, 0), c.input, tc); isErr || !ran() { + t.Fatalf("%s, approved: error=%v out=%q ran=%v", c.name, isErr, out, ran()) + } + _ = os.Remove(marker) + } + if res.lookups() == 0 { + t.Fatal("the name formed by the parameter was never resolved") + } +} + +// The custom-tool test endpoint is the operator's own action (D14): it has no +// approval scope, so the gate is not on there, whatever the tasks say. +func TestCustomToolTestEndpointUngated(t *testing.T) { + g := newGateServer(t) + on := true + g.newTask(t, db.TaskCreateOptions{TargetGate: &on}) + _, hits, host, port := countingServer(t) + before := 0 + if err := g.m.pg.QueryRow(`SELECT count(*) FROM intercept_pending WHERE decision_source='target' OR target_key IS NOT NULL`).Scan(&before); err != nil { + t.Fatal(err) + } + body := fmt.Sprintf(`{"kind":"http","exec":{"url":"http://%s:{port}/test"},"params":{"port":%q}}`, host, port) + code, out := g.do(t, "POST", "/api/tools/custom/test", body) + if code != 200 || out["is_error"] != false || !strings.Contains(fmt.Sprint(out["output"]), "probe-reached") || hits.Load() != 1 { + t.Fatalf("http: %d %v hits=%d", code, out, hits.Load()) + } + code, out = g.do(t, "POST", "/api/tools/custom/test", `{"kind":"command","exec":{"command":": http://198.51.100.7:8080/ && echo ran-ungated"}}`) + if code != 200 || !strings.Contains(fmt.Sprint(out["output"]), "ran-ungated") { + t.Fatalf("command: %d %v", code, out) + } + after := 0 + if err := g.m.pg.QueryRow(`SELECT count(*) FROM intercept_pending WHERE decision_source='target' OR target_key IS NOT NULL`).Scan(&after); err != nil { + t.Fatal(err) + } + if after != before { + t.Fatalf("the test endpoint wrote %d target rows", after-before) + } +} + +// countingResolver answers from a table and counts every lookup. +type countingResolver struct { + mu sync.Mutex + answers map[string][]string + n int +} + +func (r *countingResolver) Name() string { return "counting" } + +func (r *countingResolver) LookupIPs(_ context.Context, host string) ([]netip.Addr, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.n++ + var out []netip.Addr + for _, ip := range r.answers[host] { + out = append(out, netip.MustParseAddr(ip)) + } + if len(out) == 0 { + return nil, fmt.Errorf("no answer for %s", host) + } + return out, nil +} + +func (r *countingResolver) lookups() int { + r.mu.Lock() + defer r.mu.Unlock() + return r.n +} + +// scriptedCall is one tool call of a toolScript. +type scriptedCall struct{ tool, input string } + +func fetches(urls ...string) []scriptedCall { + calls := make([]scriptedCall, len(urls)) + for n, u := range urls { + calls[n] = scriptedCall{"WebFetch", fmt.Sprintf(`{"url":%q}`, u)} + } + return calls +} + +// toolScript is a model that makes each of its calls in turn, one per turn, +// then ends the run. It hands the result of every call to the test, and waits +// before turn n while hold[n] is open. +type toolScript struct { + calls []scriptedCall + hold map[int64]chan struct{} + results chan string + turn atomic.Int64 +} + +func (p *toolScript) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { + n := p.turn.Add(1) + if n > 1 { + select { + case p.results <- lastToolResult(req.Messages): + default: + } + } + if h := p.hold[n]; h != nil { + select { + case <-h: + case <-ctx.Done(): + } + } + events := []llm.StreamEvent{{Type: llm.SETextDelta, Text: "done"}, {Type: llm.SEMessageDelta, StopReason: "end_turn"}} + if int(n) <= len(p.calls) { + call := p.calls[n-1] + events = []llm.StreamEvent{ + {Type: llm.SEToolUseStart, ToolID: fmt.Sprintf("call-%d", n), ToolName: call.tool}, + {Type: llm.SEToolInputJSON, Text: call.input}, + {Type: llm.SEMessageDelta, StopReason: "tool_use"}, + } + } + return func(yield func(llm.StreamEvent, error) bool) { + for _, event := range events { + if !yield(event, nil) { + return + } + } + } +} + +func (p *toolScript) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + return accumulateStreamForTest(ctx, p.Stream, req) +} + +func (p *toolScript) next(t *testing.T) string { + t.Helper() + select { + case r := <-p.results: + return r + case <-time.After(60 * time.Second): + t.Fatal("the worker did not hand back a tool result") + return "" + } +} + +// lastToolResult is the text of the last tool result the model is shown. +func lastToolResult(msgs []llm.Message) string { + for i := len(msgs) - 1; i >= 0; i-- { + for _, b := range msgs[i].Content { + if b.Type != llm.BlockToolResult { + continue + } + var s strings.Builder + if b.IsError { + s.WriteString("[error] ") + } + for _, c := range b.Content { + s.WriteString(c.Text) + } + return s.String() + } + } + return "" +} + +// runScriptedWorker claims one intent of task and runs it with a worker driven +// by p, in the background. extra are tools the worker has besides its own. +func runScriptedWorker(t *testing.T, g *gateServer, task *Task, p llm.Provider, extra ...actool.CoreTool) <-chan struct{} { + t.Helper() + if _, err := task.Store.AddIntent(map[string]any{"summary": "fetch the probe"}, 1, nil, "planner"); err != nil { + t.Fatal(err) + } + intent, token := g.s.engine.claimNext(task, "work#1") + if intent == nil { + t.Fatal("no intent claimed") + } + worker := agent.NewWorker(p, "test-model", g.m.dir, nil, 0, 8, extra...) + done := make(chan struct{}) + go func() { + defer close(done) + g.s.engine.runIntent(context.Background(), task, "work#1", worker, intent, token, "", "") + }() + t.Cleanup(func() { + select { + case <-done: + case <-time.After(60 * time.Second): + t.Error("the worker run did not end") + } + }) + return done +} + +// A gated task end to end: a worker's WebFetch waits on a request, a click +// through the API pins the address and the call is made, the next call passes +// on the pin, the pin is listed and revoked, and the call after asks again. +func TestTargetGateEndToEnd(t *testing.T) { + g := newGateServer(t) + on := true + task := g.newTask(t, db.TaskCreateOptions{TargetGate: &on}) + pg := g.m.pg + srv, hits, host, port := countingServer(t) + key := net.JoinHostPort(host, port) + target := srv.URL + "/e2e" + p := &toolScript{calls: fetches(target, target, target), results: make(chan string, 16), + hold: map[int64]chan struct{}{3: make(chan struct{})}} + done := runScriptedWorker(t, g, task, p) + + var id int64 + waitUntil(t, "the first request", func() bool { + var ok bool + id, _, _, ok = openTargetRequest(t, pg, task.ID, key) + return ok + }) + if hits.Load() != 0 { + t.Fatal("fetched before the approval") + } + if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"allowed"}`); code != 200 { + t.Fatalf("decide: %d %v", code, body) + } + if r := p.next(t); !strings.Contains(r, "probe-reached") || hits.Load() != 1 { + t.Fatalf("first call after the approval: %q hits=%d", r, hits.Load()) + } + if r := p.next(t); !strings.Contains(r, "probe-reached") || hits.Load() != 2 { + t.Fatalf("second call, pinned: %q hits=%d", r, hits.Load()) + } + var requests int + if err := pg.QueryRow(`SELECT count(*) FROM intercept_pending WHERE task_id=$1 AND target_key IS NOT NULL`, task.ID).Scan(&requests); err != nil || requests != 1 { + t.Fatalf("%d requests, %v; want the one that was approved", requests, err) + } + + code, body := g.do(t, "GET", "/api/tasks/"+task.ID+"/target-pins", "") + pins, _ := body["pins"].([]any) + if code != 200 || len(pins) != 1 { + t.Fatalf("pins: %d %v", code, body) + } + pin := pins[0].(map[string]any) + if pin["host"] != host || fmt.Sprint(pin["port"]) != port || pin["status"] != "active" || fmt.Sprint(pin["pinned_ips"]) != "["+host+"]" { + t.Fatalf("listed pin = %v", pin) + } + pinID, _ := pin["id"].(float64) + if code, _ := g.do(t, "POST", fmt.Sprintf("/api/tasks/%s/target-pins/%d/revoke", task.ID, int64(pinID)), ""); code != 200 { + t.Fatalf("revoke: %d", code) + } + close(p.hold[3]) + + var again int64 + waitUntil(t, "the request after the revoke", func() bool { + var ok bool + again, _, _, ok = openTargetRequest(t, pg, task.ID, key) + return ok && again != id + }) + if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", again), `{"decision":"denied"}`); code != 200 { + t.Fatalf("deny: %d %v", code, body) + } + if r := p.next(t); !strings.Contains(r, "TARGET_DENIED") || hits.Load() != 2 { + t.Fatalf("third call, revoked then denied: %q hits=%d", r, hits.Load()) + } + select { + case <-done: + case <-time.After(60 * time.Second): + t.Fatal("the worker run did not end") + } +} + +// A custom tool in a gated worker session. A deny rule on the call decides +// first, at the hook: the tool never runs, so its rendered target is never +// read, resolved or requested. The next call, which no rule refuses, is gated +// inside the tool on the URL it renders. +func TestDenyRuleWinsBeforeRenderedTarget(t *testing.T) { + g := newGateServer(t) + on := true + task := g.newTask(t, db.TaskCreateOptions{TargetGate: &on}) + pg, ic := g.m.pg, g.m.interceptor + _, hits, _, port := countingServer(t) + res := &countingResolver{answers: map[string][]string{"api.example.com": {"127.0.0.1"}}} + ic.SetResolver(res) + + prior, err := ic.GetEnabledTools() + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ic.SetEnabledTools(prior) }) + if err := ic.SetEnabledTools(append(slices.Clone(prior), "probe_http")); err != nil { + t.Fatal(err) + } + const sentinel = "ARTEX_RENDERED_DENY_PROBE" + rule, err := pg.CreateInterceptRule("rendered deny probe", "tool_input", "string", sentinel, "deny", "", 1000000, true, false, 0, "deny") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = pg.DeleteInterceptRule(rule.ID) + ic.Invalidate() + }) + ic.Invalidate() + + tool := g.s.buildCustomTool(&db.Tool{Key: "probe_http", Kind: "http", + Schema: json.RawMessage(`{"type":"object","properties":{"sub":{"type":"string"},"note":{"type":"string"}}}`), + Exec: execJSON(t, map[string]any{"url": "http://{sub}.example.com:" + port + "/deny"})}) + p := &toolScript{calls: []scriptedCall{ + {"probe_http", `{"sub":"api","note":"` + sentinel + `"}`}, + {"probe_http", `{"sub":"api"}`}, + }, results: make(chan string, 16), hold: map[int64]chan struct{}{2: make(chan struct{})}} + done := runScriptedWorker(t, g, task, p, tool) + + if r := p.next(t); !strings.Contains(r, "intercepted by the platform") || strings.Contains(r, "TARGET_") { + t.Fatalf("call refused by the rule: %q", r) + } + if n, rows := res.lookups(), targetRowCount(t, pg, task.ID); n != 0 || rows != 0 || hits.Load() != 0 { + t.Fatalf("after the deny: %d lookups, %d target rows, %d requests sent; want none", n, rows, hits.Load()) + } + close(p.hold[2]) + + key := "api.example.com:" + port + var id int64 + waitUntil(t, "the request for the rendered target", func() bool { + var ok bool + var tool string + id, tool, _, ok = openTargetRequest(t, pg, task.ID, key) + return ok && tool == "probe_http" + }) + if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"denied"}`); code != 200 { + t.Fatalf("deny: %d %v", code, body) + } + if r := p.next(t); !strings.Contains(r, "TARGET_DENIED") || hits.Load() != 0 || res.lookups() == 0 { + t.Fatalf("rendered target: %q, %d requests sent, %d lookups", r, hits.Load(), res.lookups()) + } + select { + case <-done: + case <-time.After(60 * time.Second): + t.Fatal("the worker run did not end") + } +} + +// A task with the gate off (the default) runs as before: a worker's WebFetch +// and the custom tools reach their targets, and the gate writes no row and +// resolves no name. +func TestGateOffTaskIsUnchanged(t *testing.T) { + g := newGateServer(t) + task := g.newTask(t, db.TaskCreateOptions{}) + res := &countingResolver{answers: map[string][]string{"localhost": {"127.0.0.1"}}} + g.m.interceptor.SetResolver(res) + _, hits, _, port := countingServer(t) + byName := "http://localhost:" + port + + p := &toolScript{calls: fetches(byName+"/a", byName+"/b"), results: make(chan string, 16)} + done := runScriptedWorker(t, g, task, p) + for n := 1; n <= 2; n++ { + if r := p.next(t); !strings.Contains(r, "probe-reached") { + t.Fatalf("WebFetch %d: %q", n, r) + } + } + select { + case <-done: + case <-time.After(60 * time.Second): + t.Fatal("the worker run did not end") + } + + httpTool := g.s.buildCustomTool(&db.Tool{Key: "probe_http", Kind: "http", + Schema: json.RawMessage(`{"type":"object","properties":{"port":{"type":"string"}}}`), + Exec: execJSON(t, map[string]any{"url": "http://localhost:{port}/c"})}) + if out, isErr := callTool(t, httpTool, gatedCtx(task, intercept.DefaultTargetWaitBound), fmt.Sprintf(`{"port":%q}`, port), nil); isErr || !strings.Contains(out, "probe-reached") { + t.Fatalf("custom http tool: error=%v out=%q", isErr, out) + } + cmdTool := g.s.buildCustomTool(&db.Tool{Key: "probe_cmd", Kind: "command", + Exec: execJSON(t, map[string]any{"command": ": http://localhost:" + port + "/ && echo ran-gate-off"})}) + if out, isErr := callTool(t, cmdTool, gatedCtx(task, intercept.DefaultTargetWaitBound), `{}`, &actool.ToolContext{WorkingDir: t.TempDir()}); isErr || !strings.Contains(out, "ran-gate-off") { + t.Fatalf("custom command tool: error=%v out=%q", isErr, out) + } + + if hits.Load() != 3 { + t.Fatalf("%d requests reached the server, want 3", hits.Load()) + } + if n := targetRowCount(t, g.m.pg, task.ID); n != 0 { + t.Fatalf("%d target rows, want none", n) + } + if n := res.lookups(); n != 0 { + t.Fatalf("%d lookups, want none", n) + } +} diff --git a/server/customtool_test.go b/server/customtool_test.go index cead4fe..088651a 100644 --- a/server/customtool_test.go +++ b/server/customtool_test.go @@ -128,7 +128,7 @@ func TestRunHTTPTool(t *testing.T) { "headers": map[string]string{"Authorization": "Bearer {token}"}, "body": `{"flag":"{flag}"}`, }) - res, err := (&Server{}).runHTTPTool(context.Background(), execRaw, + res, err := (&Server{}).runHTTPTool(context.Background(), "probe", execRaw, map[string]any{"flag": "CTF{x}", "token": "sekret"}, nil) if err != nil { t.Fatalf("runHTTPTool: %v", err) @@ -176,7 +176,7 @@ func TestHTTPToolTruncatesLargeBody(t *testing.T) { defer srv.Close() execRaw, _ := json.Marshal(map[string]any{"method": "GET", "url": srv.URL}) - res, err := (&Server{}).runHTTPTool(context.Background(), execRaw, map[string]any{}, nil) + res, err := (&Server{}).runHTTPTool(context.Background(), "probe", execRaw, map[string]any{}, nil) if err != nil { t.Fatalf("runHTTPTool: %v", err) }