permission.go

  1package permission
  2
  3import (
  4	"context"
  5	"errors"
  6	"path/filepath"
  7	"slices"
  8	"sync"
  9
 10	"github.com/charmbracelet/crush/internal/pubsub"
 11	"github.com/google/uuid"
 12)
 13
 14var ErrorPermissionDenied = errors.New("permission denied")
 15
 16type CreatePermissionRequest struct {
 17	SessionID   string `json:"session_id"`
 18	ToolCallID  string `json:"tool_call_id"`
 19	ToolName    string `json:"tool_name"`
 20	Description string `json:"description"`
 21	Action      string `json:"action"`
 22	Params      any    `json:"params"`
 23	Path        string `json:"path"`
 24}
 25
 26type PermissionNotification struct {
 27	ToolCallID string `json:"tool_call_id"`
 28	Granted    bool   `json:"granted"`
 29	Denied     bool   `json:"denied"`
 30}
 31
 32type PermissionRequest struct {
 33	ID          string `json:"id"`
 34	SessionID   string `json:"session_id"`
 35	ToolCallID  string `json:"tool_call_id"`
 36	ToolName    string `json:"tool_name"`
 37	Description string `json:"description"`
 38	Action      string `json:"action"`
 39	Params      any    `json:"params"`
 40	Path        string `json:"path"`
 41}
 42
 43type Service interface {
 44	pubsub.Suscriber[PermissionRequest]
 45	GrantPersistent(permission PermissionRequest)
 46	Grant(permission PermissionRequest)
 47	Deny(permission PermissionRequest)
 48	Request(opts CreatePermissionRequest) bool
 49	AutoApproveSession(sessionID string)
 50	SubscribeNotifications(ctx context.Context) <-chan pubsub.Event[PermissionNotification]
 51}
 52
 53type permissionService struct {
 54	*pubsub.Broker[PermissionRequest]
 55
 56	notificationBroker    *pubsub.Broker[PermissionNotification]
 57	workingDir            string
 58	sessionPermissions    []PermissionRequest
 59	sessionPermissionsMu  sync.RWMutex
 60	pendingRequests       sync.Map
 61	autoApproveSessions   map[string]bool
 62	autoApproveSessionsMu sync.RWMutex
 63	skip                  bool
 64	allowedTools          []string
 65
 66	// used to make sure we only process one request at a time
 67	requestMu     sync.Mutex
 68	activeRequest *PermissionRequest
 69}
 70
 71func (s *permissionService) GrantPersistent(permission PermissionRequest) {
 72	s.notificationBroker.Publish(pubsub.CreatedEvent, PermissionNotification{
 73		ToolCallID: permission.ToolCallID,
 74		Granted:    true,
 75	})
 76	respCh, ok := s.pendingRequests.Load(permission.ID)
 77	if ok {
 78		respCh.(chan bool) <- true
 79	}
 80
 81	s.sessionPermissionsMu.Lock()
 82	s.sessionPermissions = append(s.sessionPermissions, permission)
 83	s.sessionPermissionsMu.Unlock()
 84
 85	if s.activeRequest != nil && s.activeRequest.ID == permission.ID {
 86		s.activeRequest = nil
 87	}
 88}
 89
 90func (s *permissionService) Grant(permission PermissionRequest) {
 91	s.notificationBroker.Publish(pubsub.CreatedEvent, PermissionNotification{
 92		ToolCallID: permission.ToolCallID,
 93		Granted:    true,
 94	})
 95	respCh, ok := s.pendingRequests.Load(permission.ID)
 96	if ok {
 97		respCh.(chan bool) <- true
 98	}
 99
100	if s.activeRequest != nil && s.activeRequest.ID == permission.ID {
101		s.activeRequest = nil
102	}
103}
104
105func (s *permissionService) Deny(permission PermissionRequest) {
106	s.notificationBroker.Publish(pubsub.CreatedEvent, PermissionNotification{
107		ToolCallID: permission.ToolCallID,
108		Granted:    false,
109		Denied:     true,
110	})
111	respCh, ok := s.pendingRequests.Load(permission.ID)
112	if ok {
113		respCh.(chan bool) <- false
114	}
115
116	if s.activeRequest != nil && s.activeRequest.ID == permission.ID {
117		s.activeRequest = nil
118	}
119}
120
121func (s *permissionService) Request(opts CreatePermissionRequest) bool {
122	if s.skip {
123		return true
124	}
125
126	// tell the UI that a permission was requested
127	s.notificationBroker.Publish(pubsub.CreatedEvent, PermissionNotification{
128		ToolCallID: opts.ToolCallID,
129	})
130	s.requestMu.Lock()
131	defer s.requestMu.Unlock()
132
133	// Check if the tool/action combination is in the allowlist
134	commandKey := opts.ToolName + ":" + opts.Action
135	if slices.Contains(s.allowedTools, commandKey) || slices.Contains(s.allowedTools, opts.ToolName) {
136		return true
137	}
138
139	s.autoApproveSessionsMu.RLock()
140	autoApprove := s.autoApproveSessions[opts.SessionID]
141	s.autoApproveSessionsMu.RUnlock()
142
143	if autoApprove {
144		return true
145	}
146
147	dir := filepath.Dir(opts.Path)
148	if dir == "." {
149		dir = s.workingDir
150	}
151	permission := PermissionRequest{
152		ID:          uuid.New().String(),
153		Path:        dir,
154		SessionID:   opts.SessionID,
155		ToolCallID:  opts.ToolCallID,
156		ToolName:    opts.ToolName,
157		Description: opts.Description,
158		Action:      opts.Action,
159		Params:      opts.Params,
160	}
161
162	s.sessionPermissionsMu.RLock()
163	for _, p := range s.sessionPermissions {
164		if p.ToolName == permission.ToolName && p.Action == permission.Action && p.SessionID == permission.SessionID && p.Path == permission.Path {
165			s.sessionPermissionsMu.RUnlock()
166			return true
167		}
168	}
169	s.sessionPermissionsMu.RUnlock()
170
171	s.sessionPermissionsMu.RLock()
172	for _, p := range s.sessionPermissions {
173		if p.ToolName == permission.ToolName && p.Action == permission.Action && p.SessionID == permission.SessionID && p.Path == permission.Path {
174			s.sessionPermissionsMu.RUnlock()
175			return true
176		}
177	}
178	s.sessionPermissionsMu.RUnlock()
179
180	s.activeRequest = &permission
181
182	respCh := make(chan bool, 1)
183	s.pendingRequests.Store(permission.ID, respCh)
184	defer s.pendingRequests.Delete(permission.ID)
185
186	// Publish the request
187	s.Publish(pubsub.CreatedEvent, permission)
188
189	return <-respCh
190}
191
192func (s *permissionService) AutoApproveSession(sessionID string) {
193	s.autoApproveSessionsMu.Lock()
194	s.autoApproveSessions[sessionID] = true
195	s.autoApproveSessionsMu.Unlock()
196}
197
198func (s *permissionService) SubscribeNotifications(ctx context.Context) <-chan pubsub.Event[PermissionNotification] {
199	return s.notificationBroker.Subscribe(ctx)
200}
201
202func NewPermissionService(workingDir string, skip bool, allowedTools []string) Service {
203	return &permissionService{
204		Broker:              pubsub.NewBroker[PermissionRequest](),
205		notificationBroker:  pubsub.NewBroker[PermissionNotification](),
206		workingDir:          workingDir,
207		sessionPermissions:  make([]PermissionRequest, 0),
208		autoApproveSessions: make(map[string]bool),
209		skip:                skip,
210		allowedTools:        allowedTools,
211	}
212}