ollie/cmd/olliesrv/internal/bypass/broker.go

package bypass

import (
	"fmt"
	"sync"
	"time"

	"golang.org/x/time/rate"
)

const (
	RequestTTL     = 300 * time.Second
	BurstLimit     = 3               // max denied/unanswered requests in burst window
	RefillInterval = 10 * time.Second // one slot recovers every 10s (30s to full)
)

// Broker manages bypass request evaluation and policy.
// It does NOT execute commands - that happens in toolsrv.
type Broker struct {
	policy       *PolicyStore
	notify       NotifyFunc
	logf         func(string, ...any)
	sessionValid func(string) bool

	mu       sync.RWMutex
	pending  map[string]*Request // id -> request
	sessions map[string]*Policy  // sessionID -> session policy

	rateMu   sync.Mutex
	limiters map[string]*rate.Limiter // sessionID -> denial rate limiter

	pendingCh chan *Request // channel for blocking NextPending
}

// BrokerConfig configures the broker.
type BrokerConfig struct {
	PolicyPath   string // disk path for global policy YAML
	Notify       NotifyFunc
	Logf         func(string, ...any)
	SessionValid func(sessionID string) bool // returns true if session exists
}

// NotifyFunc is called when a new elevation request needs user attention.
type NotifyFunc func(*Request)

// NewBroker creates the bypass broker.
func NewBroker(cfg BrokerConfig) (*Broker, error) {
	if cfg.Logf == nil {
		cfg.Logf = func(string, ...any) {}
	}
	if cfg.Notify == nil {
		cfg.Notify = func(*Request) {}
	}

	b := &Broker{
		policy:       NewPolicyStore(cfg.PolicyPath),
		notify:       cfg.Notify,
		logf:         cfg.Logf,
		sessionValid: cfg.SessionValid,
		pending:      make(map[string]*Request),
		sessions:     make(map[string]*Policy),
		limiters:     make(map[string]*rate.Limiter),
		pendingCh:    make(chan *Request, 16),
	}

	cfg.Logf("elevate: broker ready")
	return b, nil
}

// Close shuts down the broker.
func (b *Broker) Close() {
	// No resources to clean up now that socket is gone
}

// Pending returns all pending requests.
func (b *Broker) Pending() []*Request {
	b.mu.RLock()
	defer b.mu.RUnlock()
	out := make([]*Request, 0, len(b.pending))
	for _, r := range b.pending {
		out = append(out, r)
	}
	return out
}

// NextPending blocks until a bypass request is available and returns it.
func (b *Broker) NextPending() *Request {
	return <-b.pendingCh
}

// PendingByID returns a specific pending request.
func (b *Broker) PendingByID(id string) *Request {
	b.mu.RLock()
	defer b.mu.RUnlock()
	return b.pending[id]
}

// Resolve resolves a pending request. Safe to call after timeout (no-op).
func (b *Broker) Resolve(id string, res Resolution) bool {
	b.mu.Lock()
	req, ok := b.pending[id]
	if !ok {
		b.mu.Unlock()
		return false
	}
	delete(b.pending, id)
	b.mu.Unlock()

	// If persist, add rule. Session-scoped if we know the session, global otherwise.
	if res == ResolvePersist {
		if req.SessionID != "" {
			b.addSessionRule(req.SessionID, Rule{Cmd: req.Cmd})
		} else {
			b.policy.AddGlobal(Rule{Cmd: req.Cmd}) //nolint:errcheck
		}
		res = ResolveApprove // treat as approve for execution
	}

	select {
	case req.resolved <- res:
	default:
	}
	return true
}

// SessionPolicy returns the effective policy for a session (global + session).
func (b *Broker) SessionPolicy(sessionID string) *Policy {
	global := b.policy.Global()
	b.mu.RLock()
	sess, ok := b.sessions[sessionID]
	b.mu.RUnlock()

	// Merge: global rules + session rules
	merged := &Policy{Rules: make([]Rule, len(global.Rules))}
	copy(merged.Rules, global.Rules)
	if ok {
		merged.Rules = append(merged.Rules, sess.Rules...)
	}
	return merged
}

// SetSessionPolicy replaces the session policy.
func (b *Broker) SetSessionPolicy(sessionID string, p Policy) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.sessions[sessionID] = &p
}

// RemoveSession cleans up session policy when a session is killed.
func (b *Broker) RemoveSession(sessionID string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	delete(b.sessions, sessionID)

	b.rateMu.Lock()
	delete(b.limiters, sessionID)
	b.rateMu.Unlock()
}

// limiterFor returns (or creates) the denial rate limiter for a session.
func (b *Broker) limiterFor(sessionID string) *rate.Limiter {
	b.rateMu.Lock()
	defer b.rateMu.Unlock()
	lim, ok := b.limiters[sessionID]
	if !ok {
		lim = rate.NewLimiter(rate.Every(RefillInterval), BurstLimit)
		b.limiters[sessionID] = lim
	}
	return lim
}

// checkRateLimit returns true if the session still has denial budget.
func (b *Broker) checkRateLimit(sessionID string) bool {
	return b.limiterFor(sessionID).Tokens() >= 1
}

// recordDenial consumes a token from the session's denial budget.
func (b *Broker) recordDenial(sessionID string) {
	if sessionID == "" {
		return
	}
	b.limiterFor(sessionID).Allow()
}

// EvaluateRequest evaluates a bypass request from toolsrv and returns the decision.
// Execution happens in toolsrv, not here.
// Returns (approved, error). If needs user approval, blocks until resolved or timeout.
func (b *Broker) EvaluateRequest(sessionID, cmd, cwd string, env map[string]string) (bool, error) {
	// Validate session
	if sessionID == "" {
		return false, fmt.Errorf("no session identity")
	}
	if b.sessionValid != nil && !b.sessionValid(sessionID) {
		return false, fmt.Errorf("session not found")
	}

	// Rate limit check
	if !b.checkRateLimit(sessionID) {
		b.logf("elevate: rate-limited session=%s cmd=%q", sessionID, cmd)
		return false, fmt.Errorf("rate-limited")
	}

	// Check policy auto-approve
	effective := b.SessionPolicy(sessionID)
	if effective.Matches(cmd, "") {
		b.logf("elevate: auto-approved cmd=%q session=%s", cmd, sessionID)
		return true, nil
	}

	// Create pending request for user approval
	req := &Request{
		ID:        nextRequestID(),
		Cmd:       cmd,
		Cwd:       cwd,
		Env:       env,
		SessionID: sessionID,
		CreatedAt: time.Now(),
		resolved:  make(chan Resolution, 1),
	}

	b.mu.Lock()
	b.pending[req.ID] = req
	b.mu.Unlock()

	// Send to pending channel for NextPending readers
	b.pendingCh <- req

	// Notify user (D-Bus notification, etc.)
	b.notify(req)
	b.logf("elevate: pending id=%s cmd=%q session=%s", req.ID, cmd, sessionID)

	// Wait for resolution or timeout
	timer := time.NewTimer(RequestTTL)
	defer timer.Stop()

	var res Resolution
	select {
	case res = <-req.Resolved():
	case <-timer.C:
		b.mu.Lock()
		delete(b.pending, req.ID)
		b.mu.Unlock()
		res = ResolveTimeout
	}

	switch res {
	case ResolveApprove, ResolvePersist:
		b.logf("elevate: approved id=%s cmd=%q", req.ID, cmd)
		return true, nil
	default:
		b.recordDenial(sessionID)
		b.logf("elevate: denied id=%s cmd=%q reason=%s", req.ID, cmd, res)
		return false, nil
	}
}

// GlobalPolicy returns the policy store for 9P access.
func (b *Broker) GlobalPolicy() *PolicyStore {
	return b.policy
}

func (b *Broker) addSessionRule(sessionID string, r Rule) {
	if sessionID == "" {
		return
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	p, ok := b.sessions[sessionID]
	if !ok {
		p = &Policy{}
		b.sessions[sessionID] = p
	}
	p.Add(r)
}