on main help sign in
fix(permission): fix publish-before-lock race and use O(1) session permission lookups
02b783aKieran Klukas
internal/permission/permission.go+30−19
@@ -73,13 +73,20 @@ type Service interface {
73 SubscribeNotifications(ctx context.Context) <-chan pubsub.Event[PermissionNotification]
73 SubscribeNotifications(ctx context.Context) <-chan pubsub.Event[PermissionNotification]
74}
74}
75
75
76+// PermissionKey is a composite key for session permission lookups.
77+type PermissionKey struct {
78+ SessionID string
79+ ToolName string
80+ Action string
81+ Path string
82+}
83+
76type permissionService struct {
84type permissionService struct {
77 *pubsub.Broker[PermissionRequest]
85 *pubsub.Broker[PermissionRequest]
78
86
79 notificationBroker *pubsub.Broker[PermissionNotification]
87 notificationBroker *pubsub.Broker[PermissionNotification]
80 workingDir string
88 workingDir string
81− sessionPermissions []PermissionRequest
89+ sessionPermissions *csync.Map[PermissionKey, bool]
82− sessionPermissionsMu sync.RWMutex
83 pendingRequests *csync.Map[string, chan bool]
90 pendingRequests *csync.Map[string, chan bool]
84 autoApproveSessions map[string]bool
91 autoApproveSessions map[string]bool
85 autoApproveSessionsMu sync.RWMutex
92 autoApproveSessionsMu sync.RWMutex
@@ -102,9 +109,12 @@ func (s *permissionService) GrantPersistent(permission PermissionRequest) {
102 respCh <- true
109 respCh <- true
103 }
110 }
104
111
105− s.sessionPermissionsMu.Lock()
112+ s.sessionPermissions.Set(PermissionKey{
106− s.sessionPermissions = append(s.sessionPermissions, permission)
113+ SessionID: permission.SessionID,
107− s.sessionPermissionsMu.Unlock()
114+ ToolName: permission.ToolName,
115+ Action: permission.Action,
116+ Path: permission.Path,
117+ }, true)
108
118
109 s.activeRequestMu.Lock()
119 s.activeRequestMu.Lock()
110 if s.activeRequest != nil && s.activeRequest.ID == permission.ID {
120 if s.activeRequest != nil && s.activeRequest.ID == permission.ID {
@@ -171,12 +181,13 @@ func (s *permissionService) Request(ctx context.Context, opts CreatePermissionRe
171 return true, nil
181 return true, nil
172 }
182 }
173
183
184+ s.requestMu.Lock()
185+ defer s.requestMu.Unlock()
186+
174 // tell the UI that a permission was requested
187 // tell the UI that a permission was requested
175 s.notificationBroker.Publish(pubsub.CreatedEvent, PermissionNotification{
188 s.notificationBroker.Publish(pubsub.CreatedEvent, PermissionNotification{
176 ToolCallID: opts.ToolCallID,
189 ToolCallID: opts.ToolCallID,
177 })
190 })
178− s.requestMu.Lock()
179− defer s.requestMu.Unlock()
180
191
181 s.autoApproveSessionsMu.RLock()
192 s.autoApproveSessionsMu.RLock()
182 autoApprove := s.autoApproveSessions[opts.SessionID]
193 autoApprove := s.autoApproveSessions[opts.SessionID]
@@ -214,18 +225,18 @@ func (s *permissionService) Request(ctx context.Context, opts CreatePermissionRe
214 Params: opts.Params,
225 Params: opts.Params,
215 }
226 }
216
227
217− s.sessionPermissionsMu.RLock()
218− for _, p := range s.sessionPermissions {
228+ if _, ok := s.sessionPermissions.Get(PermissionKey{
229+ SessionID: permission.SessionID,
219− if p.ToolName == permission.ToolName && p.Action == permission.Action && p.SessionID == permission.SessionID && p.Path == permission.Path {
230+ ToolName: permission.ToolName,
220− s.sessionPermissionsMu.RUnlock()
231+ Action: permission.Action,
232+ Path: permission.Path,
233+ }); ok {
221− s.notificationBroker.Publish(pubsub.CreatedEvent, PermissionNotification{
234+ s.notificationBroker.Publish(pubsub.CreatedEvent, PermissionNotification{
222− ToolCallID: opts.ToolCallID,
235+ ToolCallID: opts.ToolCallID,
223− Granted: true,
236+ Granted: true,
224− })
237+ })
225− return true, nil
238+ return true, nil
226− }
227 }
239 }
228− s.sessionPermissionsMu.RUnlock()
229
240
230 s.activeRequestMu.Lock()
241 s.activeRequestMu.Lock()
231 s.activeRequest = &permission
242 s.activeRequest = &permission
@@ -269,7 +280,7 @@ func NewPermissionService(workingDir string, skip bool, allowedTools []string) S
269 Broker: pubsub.NewBroker[PermissionRequest](),
280 Broker: pubsub.NewBroker[PermissionRequest](),
270 notificationBroker: pubsub.NewBroker[PermissionNotification](),
281 notificationBroker: pubsub.NewBroker[PermissionNotification](),
271 workingDir: workingDir,
282 workingDir: workingDir,
272− sessionPermissions: make([]PermissionRequest, 0),
283+ sessionPermissions: csync.NewMap[PermissionKey, bool](),
273 autoApproveSessions: make(map[string]bool),
284 autoApproveSessions: make(map[string]bool),
274 skip: skip,
285 skip: skip,
275 allowedTools: allowedTools,
286 allowedTools: allowedTools,