Vectrum
Sign In
143 lines · 4.9 KB · Go
package approval

import (
	"context"
	"encoding/json"
	"reflect"
	"strings"
	"testing"
)

// fixedRules rules the same way on every call to one tool.
type fixedRules struct {
	tool    string
	verdict Verdict
}

func (f fixedRules) Check(tool string, _ json.RawMessage) Ruling {
	if tool != f.tool {
		return Ruling{}
	}
	return Ruling{Verdict: f.verdict, Rule: "Rule(x)", Source: "settings.json"}
}

func (f fixedRules) Lines() []string { return []string{"settings.json", "  rule"} }

func TestStandingRules(t *testing.T) {
	tests := []struct {
		name       string
		tool       string
		verdict    Verdict
		yolo       bool
		remembered bool
		wantErr    bool
		wantPrompt int
	}{
		{name: "deny holds under yolo", tool: "bash", verdict: VerdictDeny, yolo: true, wantErr: true},
		{name: "deny holds for a read-only tool", tool: "read_file", verdict: VerdictDeny, wantErr: true},
		{name: "deny beats session memory", tool: "bash", verdict: VerdictDeny, remembered: true, wantErr: true},
		{name: "allow skips the prompt", tool: "bash", verdict: VerdictAllow},
		{name: "ask prompts under yolo", tool: "bash", verdict: VerdictAsk, yolo: true, wantPrompt: 1},
		{name: "ask prompts for a read-only tool", tool: "read_file", verdict: VerdictAsk, wantPrompt: 1},
		{name: "ask beats session memory", tool: "bash", verdict: VerdictAsk, remembered: true, wantPrompt: 1},
		{name: "no verdict prompts as usual", tool: "bash", verdict: NoVerdict, wantPrompt: 1},
		{name: "no verdict still honours memory", tool: "bash", verdict: NoVerdict, remembered: true},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			var asked int
			var forced string
			p := New(func(_ context.Context, r Request) (Decision, error) {
				asked++
				forced = r.Forced
				return Once, nil
			})
			p.YOLO = tc.yolo
			p.Rules = fixedRules{tool: tc.tool, verdict: tc.verdict}
			if tc.remembered {
				p.tools[tc.tool] = true
			}
			err := approve(t, p, tc.tool, `{"command":"ls","path":"a"}`)
			if (err != nil) != tc.wantErr {
				t.Fatalf("err = %v, want error %v", err, tc.wantErr)
			}
			if err != nil && !strings.Contains(err.Error(), "Rule(x) in settings.json") {
				t.Errorf("denial %q should name the rule and its file", err)
			}
			if asked != tc.wantPrompt {
				t.Errorf("prompted %d times, want %d", asked, tc.wantPrompt)
			}
			if tc.verdict == VerdictAsk && forced == "" {
				t.Error("a prompt an ask rule forced should say so")
			}
		})
	}
}

// Headless, an ask rule is a refusal with nobody to put the question to —
// which is what the rule asked for.
func TestAskRuleHeadlessDenies(t *testing.T) {
	p := New(nil)
	p.YOLO = true
	p.Rules = fixedRules{tool: "bash", verdict: VerdictAsk}
	if err := approve(t, p, "bash", `{"command":"ls"}`); err == nil {
		t.Error("an ask rule with no prompt should not run the call")
	}
}

func TestRuleLines(t *testing.T) {
	var p *Policy
	if p.RuleLines() != nil {
		t.Error("nil policy")
	}
	p = New(nil)
	if p.RuleLines() != nil {
		t.Error("no rules")
	}
	p.Rules = fixedRules{}
	if len(p.RuleLines()) != 2 {
		t.Error("lines not passed through")
	}
}

func TestSimpleCommands(t *testing.T) {
	tests := []struct {
		script string
		want   []string
		pinned []bool
		opaque bool
	}{
		{"git status", []string{"git status"}, []bool{true}, false},
		{"go build ./... && go test ./...", []string{"go build ./...", "go test ./..."}, []bool{true, true}, false},
		{"ls | grep x; echo done", []string{"ls", "grep x", "echo done"}, []bool{true, true, true}, false},
		{"(cd sub && make)", []string{"cd sub", "make"}, []bool{true, true}, false},
		{`git commit -m "fix: the thing"`, []string{"git commit -m fix: the thing"}, []bool{true}, false},
		{"rm $FILE", []string{"rm <expanded>"}, []bool{true}, false},
		{"ls *.go", []string{"ls <expanded>"}, []bool{true}, false},
		{"FOO=1 make", []string{"make"}, []bool{false}, false},
		{"echo x > out.txt", []string{"echo x"}, []bool{false}, false},
		{"echo x >> out.txt", []string{"echo x"}, []bool{false}, false},
		{"echo x 2>&1 >/dev/null", []string{"echo x"}, []bool{true}, false},
		{"export A=1", []string{"export"}, []bool{false}, false},
		{"echo $(whoami)", []string{"echo <expanded>", "whoami"}, []bool{true, true}, true},
		{"diff <(ls a) <(ls b)", []string{"diff <expanded> <expanded>", "ls a", "ls b"}, []bool{true, true, true}, true},
	}
	for _, tc := range tests {
		cmds, info := SimpleCommands(tc.script)
		var got []string
		var pinned []bool
		for _, c := range cmds {
			got = append(got, c.Text())
			pinned = append(pinned, c.Pinned)
		}
		if !reflect.DeepEqual(got, tc.want) || !reflect.DeepEqual(pinned, tc.pinned) {
			t.Errorf("%q: got %q %v, want %q %v", tc.script, got, pinned, tc.want, tc.pinned)
		}
		if info.Opaque != tc.opaque {
			t.Errorf("%q: opaque %v, want %v", tc.script, info.Opaque, tc.opaque)
		}
	}
	if cmds, info := SimpleCommands("echo 'open"); cmds != nil || info.ParseError == nil {
		t.Error("a parse error should yield nothing")
	}
}