llmgw/internal/store/toolgrants_test.go

package store

import (
	"context"
	"path/filepath"
	"slices"
	"testing"
)

func TestGrantingToolsAtALevelKeepsWhatItSawFromAbove(t *testing.T) {
	ctx := context.Background()
	s, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "s.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	org, _ := s.CreateOrg(ctx, "Acme")
	team, _ := s.CreateTeam(ctx, org, "eng")
	tools := func(scope Scope, id string) []string {
		l, _ := s.LimitsFor(ctx, scope, id)
		return l.AllowedTools
	}
	sees := func(tool string) bool {
		e := Effective{}
		for _, l := range []struct {
			sc Scope
			id string
		}{{ScopeOrg, org}, {ScopeTeam, team}} {
			if ts := tools(l.sc, l.id); ts != nil {
				e.ToolGrants = append(e.ToolGrants, ts)
			}
		}
		return e.AllowsTool(tool)
	}

	// The organization grants hr's tools; the team names none and sees them.
	if _, err := s.GrantTools(ctx, ScopeOrg, org, "hr", []string{"hr.*"}); err != nil {
		t.Fatal(err)
	}
	if !sees("hr.search") {
		t.Fatal("the team does not see what its organization grants")
	}
	// Granting jira's to the team keeps hr's, and says the organization
	// still keeps jira's out of reach.
	g, err := s.GrantTools(ctx, ScopeTeam, team, "jira", []string{"jira.*"})
	if err != nil {
		t.Fatal(err)
	}
	if !sees("hr.search") {
		t.Errorf("granting jira's tools to the team took hr's away: %v", tools(ScopeTeam, team))
	}
	if len(g.BlockedBy) != 1 || g.BlockedBy[0].ID != org {
		t.Errorf("the organization, which does not grant jira's tools, should be named: %+v", g)
	}
	if sees("jira.search") {
		t.Error("the team sees jira's tools its organization does not grant")
	}
	// The organization grants them too: now the team sees both.
	_, _ = s.GrantTools(ctx, ScopeOrg, org, "jira", []string{"jira.*"})
	if !sees("jira.search") || !sees("hr.get") {
		t.Errorf("org %v team %v", tools(ScopeOrg, org), tools(ScopeTeam, team))
	}
	// Taking jira's back from the team leaves it hr's.
	if err := s.RevokeTools(ctx, ScopeTeam, team, "jira"); err != nil {
		t.Fatal(err)
	}
	if sees("jira.search") || !sees("hr.get") {
		t.Errorf("after the revoke: team %v", tools(ScopeTeam, team))
	}
}

func TestRevokingALevelsLastToolsNeitherClosesNorOpensIt(t *testing.T) {
	ctx := context.Background()
	s, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "s.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	org, _ := s.CreateOrg(ctx, "Acme")
	team, _ := s.CreateTeam(ctx, org, "eng")

	// Nothing above names tools: the team's last grant taken leaves it
	// naming none at all, not closed to every tool.
	_, _ = s.GrantTools(ctx, ScopeTeam, team, "jira", []string{"jira.*"})
	_ = s.RevokeTools(ctx, ScopeTeam, team, "jira")
	if l, _ := s.LimitsFor(ctx, ScopeTeam, team); l.AllowedTools != nil {
		t.Errorf("left %v; want no list", l.AllowedTools)
	}
	// The organization grants every tool; the team was kept to jira's.
	// Taking jira's away must not open it to everything above.
	_ = s.PutLimits(ctx, Limits{Scope: ScopeOrg, ScopeID: org, AllowedTools: []string{"*"}})
	_ = s.PutLimits(ctx, Limits{Scope: ScopeTeam, ScopeID: team, AllowedTools: []string{"jira.*"}})
	_ = s.RevokeTools(ctx, ScopeTeam, team, "jira")
	if l, _ := s.LimitsFor(ctx, ScopeTeam, team); l.AllowedTools == nil || len(l.AllowedTools) != 0 {
		t.Errorf("left %v; want an empty list, closed", l.AllowedTools)
	}
	if !slices.Equal(withoutServer([]string{"a.x", "b.*", "a.*"}, "a"), []string{"b.*"}) {
		t.Error("withoutServer")
	}
}

func TestRevokingFromAnOrganizationTakesItFromEverythingBeneath(t *testing.T) {
	ctx := context.Background()
	s, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "s.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	org, _ := s.CreateOrg(ctx, "Acme")
	team, _ := s.CreateTeam(ctx, org, "eng")
	tools := func(scope Scope, id string) []string {
		l, _ := s.LimitsFor(ctx, scope, id)
		return l.AllowedTools
	}
	sees := func(tool string) bool {
		e := Effective{}
		for _, ts := range [][]string{tools(ScopeOrg, org), tools(ScopeTeam, team)} {
			if ts != nil {
				e.ToolGrants = append(e.ToolGrants, ts)
			}
		}
		return e.AllowsTool(tool)
	}
	_, _ = s.GrantTools(ctx, ScopeOrg, org, "hr", []string{"hr.*"})
	_, _ = s.GrantTools(ctx, ScopeTeam, team, "jira", []string{"jira.*"}) // keeps hr.*, copied
	// The organization takes hr's back: the team, which held a copy, loses
	// them too; and the organization, naming nothing now, is closed rather
	// than opened — its team's jira.* was kept out of reach by it before.
	if err := s.RevokeTools(ctx, ScopeOrg, org, "hr"); err != nil {
		t.Fatal(err)
	}
	if sees("hr.search") {
		t.Errorf("the team kept hr after its organization's revoke: org %v team %v", tools(ScopeOrg, org), tools(ScopeTeam, team))
	}
	if sees("jira.search") {
		t.Errorf("the revoke opened the organization to jira: org %v", tools(ScopeOrg, org))
	}
	if got := tools(ScopeOrg, org); got == nil || len(got) != 0 {
		t.Errorf("the organization is %v; want closed, []", got)
	}
	// Opened on purpose, it names none and bounds nothing: the team's own
	// list decides.
	if err := s.OpenTools(ctx, ScopeOrg, org); err != nil {
		t.Fatal(err)
	}
	if tools(ScopeOrg, org) != nil || !sees("jira.search") {
		t.Errorf("opened: org %v team %v", tools(ScopeOrg, org), tools(ScopeTeam, team))
	}
}