llmgw/internal/store/toolgrants.go

package store

import (
	"context"
	"errors"
	"slices"
	"strings"
)

// Granting one MCP server's tools at a level of the tree, and taking them
// back, as the console and the admin page do. A tool is usable only where
// every level above that names tools names it too (Effective.AllowsTool),
// which makes adding a list where there was none a narrowing: a team that
// saw its organization's jira.* through having no list of its own would see
// only hr.* once given hr.*. So a level given its first list starts from
// the nearest list above it, and keeps what it had.

// ToolGrant is what granting came to: the levels above whose lists do not
// name the tools granted, so that they stay out of reach until those grant
// them too.
type ToolGrant struct {
	BlockedBy []LimitOwner
}

// GrantTools makes a level's grants of server's tools these, keeping its
// grants of other servers' and the rest of its limits.
func (s *DB) GrantTools(ctx context.Context, scope Scope, id, server string, tools []string) (ToolGrant, error) {
	above, err := s.limitsAbove(ctx, scope, id)
	if err != nil {
		return ToolGrant{}, err
	}
	lim, err := s.LimitsFor(ctx, scope, id)
	switch {
	case errors.Is(err, ErrNotFound):
		lim = Limits{Scope: scope, ScopeID: id}
	case err != nil:
		return ToolGrant{}, err
	}
	if lim.AllowedTools == nil {
		// Its first list: what it saw through the levels above, kept.
		for _, a := range above {
			if a.AllowedTools != nil {
				lim.AllowedTools = slices.Clone(a.AllowedTools)
				break
			}
		}
	}
	keep := withoutServer(lim.AllowedTools, server)
	for _, t := range tools {
		if !slices.Contains(keep, t) {
			keep = append(keep, t)
		}
	}
	lim.AllowedTools = keep
	if err := s.PutLimits(ctx, lim); err != nil {
		return ToolGrant{}, err
	}
	var g ToolGrant
	for _, a := range above {
		if a.AllowedTools == nil {
			continue
		}
		for _, t := range tools {
			if !matchesTool(a.AllowedTools, sampleTool(t)) {
				g.BlockedBy = append(g.BlockedBy, LimitOwner{Scope: a.Scope, ID: a.ScopeID})
				break
			}
		}
	}
	return g, nil
}

// RevokeTools takes server's tools away from a level and from every level
// beneath it: taken from an organization, no team, project or key in it
// keeps them by having been given them below. A level left naming nothing
// is closed — it names no tool, so sees none — while any level above or
// beneath it names tools: going back to naming none at all would open it
// to what the levels above allow. With none anywhere, it names none.
func (s *DB) RevokeTools(ctx context.Context, scope Scope, id, server string) error {
	below, err := s.limitsBelow(ctx, scope, id)
	if err != nil {
		return err
	}
	for _, l := range below {
		if l.AllowedTools == nil {
			continue
		}
		if kept := withoutServer(l.AllowedTools, server); len(kept) != len(l.AllowedTools) {
			l.AllowedTools = kept
			if err := s.PutLimits(ctx, l); err != nil {
				return err
			}
		}
	}
	lim, err := s.LimitsFor(ctx, scope, id)
	if errors.Is(err, ErrNotFound) {
		return nil
	}
	if err != nil {
		return err
	}
	lim.AllowedTools = withoutServer(lim.AllowedTools, server)
	if len(lim.AllowedTools) == 0 {
		above, err := s.limitsAbove(ctx, scope, id)
		if err != nil {
			return err
		}
		lim.AllowedTools = nil
		for _, l := range append(above, below...) {
			if l.AllowedTools != nil {
				lim.AllowedTools = []string{}
				break
			}
		}
	}
	return s.PutLimits(ctx, lim)
}

// OpenTools has a level name no tools of its own: it then sees what the
// levels above it allow. For a level left closed by RevokeTools, on
// purpose.
func (s *DB) OpenTools(ctx context.Context, scope Scope, id string) error {
	lim, err := s.LimitsFor(ctx, scope, id)
	if errors.Is(err, ErrNotFound) {
		return nil
	}
	if err != nil {
		return err
	}
	lim.AllowedTools = nil
	return s.PutLimits(ctx, lim)
}

// limitsBelow are the limits set beneath a level: an organization's teams,
// projects and their keys, a team's projects and keys, a project's keys,
// and the member scopes of each.
func (s *DB) limitsBelow(ctx context.Context, scope Scope, id string) ([]Limits, error) {
	var owners []LimitOwner
	addProject := func(pid string) error {
		owners = append(owners, LimitOwner{Scope: ScopeProject, ID: pid}, LimitOwner{Scope: MemberScope(ScopeProject), ID: pid})
		keys, err := s.Keys(ctx, pid)
		if err != nil {
			return err
		}
		for _, k := range keys {
			owners = append(owners, LimitOwner{Scope: ScopeKey, ID: k.ID})
		}
		return nil
	}
	switch scope {
	case ScopeOrg, ScopeTeam:
		tree, err := s.Tree(ctx)
		if err != nil {
			return nil, err
		}
		for _, o := range tree {
			if scope == ScopeOrg && o.ID != id {
				continue
			}
			if scope == ScopeOrg {
				owners = append(owners, LimitOwner{Scope: MemberScope(ScopeOrg), ID: o.ID})
			}
			for _, t := range o.Teams {
				if scope == ScopeTeam && t.ID != id {
					continue
				}
				if scope == ScopeOrg {
					owners = append(owners, LimitOwner{Scope: ScopeTeam, ID: t.ID})
				}
				owners = append(owners, LimitOwner{Scope: MemberScope(ScopeTeam), ID: t.ID})
				for _, p := range t.Projects {
					if err := addProject(p.ID); err != nil {
						return nil, err
					}
				}
			}
		}
	case ScopeProject:
		if err := addProject(id); err != nil {
			return nil, err
		}
		owners = owners[1:] // the project itself is not beneath it
	}
	var out []Limits
	for _, o := range owners {
		if o.Scope == "" {
			continue
		}
		l, err := s.LimitsFor(ctx, o.Scope, o.ID)
		if errors.Is(err, ErrNotFound) {
			continue
		}
		if err != nil {
			return nil, err
		}
		out = append(out, l)
	}
	return out, nil
}

func withoutServer(tools []string, server string) []string {
	return slices.DeleteFunc(slices.Clone(tools), func(t string) bool { return strings.HasPrefix(t, server+".") })
}

// sampleTool is a tool a pattern grants, for asking whether a list above
// grants it too: jira.* stands for any of jira's.
func sampleTool(pattern string) string {
	if before, ok := strings.CutSuffix(pattern, "*"); ok {
		return before + "x"
	}
	return pattern
}

// limitsAbove are the limits of the levels above one, nearest first: a
// project's team and organization, a team's organization, a key's project
// and those above it. None for an organization or a person.
func (s *DB) limitsAbove(ctx context.Context, scope Scope, id string) ([]Limits, error) {
	var owners []LimitOwner
	if scope == ScopeKey {
		k, err := s.KeyByID(ctx, id)
		if errors.Is(err, ErrNotFound) {
			return nil, nil
		}
		if err != nil {
			return nil, err
		}
		if k.UserID != "" {
			owners = append(owners, LimitOwner{Scope: ScopeUser, ID: k.UserID})
		}
		scope, id = ScopeProject, k.ProjectID
		owners = append(owners, LimitOwner{Scope: ScopeProject, ID: id})
	}
	if scope == ScopeTeam || scope == ScopeProject {
		tree, err := s.Tree(ctx)
		if err != nil {
			return nil, err
		}
		for _, o := range tree {
			for _, t := range o.Teams {
				if scope == ScopeTeam && t.ID == id {
					owners = append(owners, LimitOwner{Scope: ScopeOrg, ID: o.ID})
				}
				for _, p := range t.Projects {
					if scope == ScopeProject && p.ID == id {
						owners = append(owners, LimitOwner{Scope: ScopeTeam, ID: t.ID}, LimitOwner{Scope: ScopeOrg, ID: o.ID})
					}
				}
			}
		}
	}
	var out []Limits
	for _, o := range owners {
		l, err := s.LimitsFor(ctx, o.Scope, o.ID)
		if errors.Is(err, ErrNotFound) {
			continue
		}
		if err != nil {
			return nil, err
		}
		out = append(out, l)
	}
	return out, nil
}