Add thread-safe probe manager with attach lifecycle

2 files changed, 506 insertions, 0 deletions

diff --git a/internal/probemanager/manager.go b/internal/probemanager/manager.go
new file mode 100644
index 0000000..5a7e85d
--- /dev/null
+++ b/internal/probemanager/manager.go
@@ -0,0 +1,353 @@
+package probemanager
+
+import (
+	"errors"
+	"fmt"
+	"sort"
+	"sync"
+)
+
+// Link abstracts an attached tracepoint link.
+type Link interface {
+	Destroy() error
+}
+
+// Program abstracts a loadable BPF program that can attach to a tracepoint.
+type Program interface {
+	AttachTracepoint(category, name string) (Link, error)
+}
+
+// Attacher resolves BPF programs by name.
+type Attacher interface {
+	GetProgram(name string) (Program, error)
+}
+
+// ProbeState is an immutable view used by callers/UI.
+type ProbeState struct {
+	Syscall string
+	Active  bool
+	Error   string
+}
+
+type probeEntry struct {
+	syscall string
+	enterTP string
+	exitTP  string
+
+	enterLink Link
+	exitLink  Link
+
+	active  bool
+	lastErr error
+}
+
+// Manager tracks probe attach/detach state for grouped syscall tracepoints.
+type Manager struct {
+	mu       sync.Mutex
+	attacher Attacher
+	probes   map[string]*probeEntry
+	closed   bool
+}
+
+func NewManager(attacher Attacher) *Manager {
+	return &Manager{
+		attacher: attacher,
+		probes:   make(map[string]*probeEntry),
+	}
+}
+
+func (m *Manager) Register(syscall string, pair TracepointPair) {
+	if m == nil || syscall == "" {
+		return
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	entry, ok := m.probes[syscall]
+	if !ok {
+		entry = &probeEntry{syscall: syscall}
+		m.probes[syscall] = entry
+	}
+	entry.enterTP = pair.Enter
+	entry.exitTP = pair.Exit
+}
+
+func (m *Manager) AttachAll(shouldAttach func(string) bool, tpNames []string) error {
+	if m == nil {
+		return errors.New("probe manager is nil")
+	}
+	if shouldAttach == nil {
+		shouldAttach = func(string) bool { return true }
+	}
+
+	groups := GroupTracepoints(tpNames)
+	for syscall, pair := range groups {
+		m.Register(syscall, pair)
+		if !shouldAttach(pair.Enter) && !shouldAttach(pair.Exit) {
+			continue
+		}
+		if err := m.Attach(syscall); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (m *Manager) Toggle(syscall string) error {
+	if m == nil {
+		return errors.New("probe manager is nil")
+	}
+	if syscall == "" {
+		return errors.New("syscall is required")
+	}
+
+	m.mu.Lock()
+	entry, err := m.entryLocked(syscall)
+	if err != nil {
+		m.mu.Unlock()
+		return err
+	}
+	active := entry.active
+	m.mu.Unlock()
+
+	if active {
+		return m.Detach(syscall)
+	}
+	return m.Attach(syscall)
+}
+
+func (m *Manager) Attach(syscall string) error {
+	if syscall == "" {
+		return errors.New("syscall is required")
+	}
+
+	m.mu.Lock()
+	entry, err := m.entryLocked(syscall)
+	if err != nil {
+		m.mu.Unlock()
+		return err
+	}
+	if entry.active {
+		m.mu.Unlock()
+		return nil
+	}
+
+	enterTP := entry.enterTP
+	exitTP := entry.exitTP
+	attacher := m.attacher
+	m.mu.Unlock()
+
+	enterLink, exitLink, attachErr := attachPair(attacher, enterTP, exitTP)
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	entry, err = m.entryLocked(syscall)
+	if err != nil {
+		if enterLink != nil {
+			_ = enterLink.Destroy()
+		}
+		if exitLink != nil {
+			_ = exitLink.Destroy()
+		}
+		return err
+	}
+
+	if attachErr != nil {
+		entry.lastErr = attachErr
+		entry.active = entry.enterLink != nil || entry.exitLink != nil
+		return attachErr
+	}
+
+	entry.enterLink = enterLink
+	entry.exitLink = exitLink
+	entry.lastErr = nil
+	entry.active = enterLink != nil || exitLink != nil
+	return nil
+}
+
+func (m *Manager) Detach(syscall string) error {
+	if syscall == "" {
+		return errors.New("syscall is required")
+	}
+
+	m.mu.Lock()
+	entry, err := m.entryLocked(syscall)
+	if err != nil {
+		m.mu.Unlock()
+		return err
+	}
+	enterLink := entry.enterLink
+	exitLink := entry.exitLink
+	entry.enterLink = nil
+	entry.exitLink = nil
+	entry.active = false
+	m.mu.Unlock()
+
+	if enterLink != nil {
+		if err := enterLink.Destroy(); err != nil {
+			m.setLastError(syscall, fmt.Errorf("detach enter %s: %w", syscall, err))
+			return err
+		}
+	}
+	if exitLink != nil {
+		if err := exitLink.Destroy(); err != nil {
+			m.setLastError(syscall, fmt.Errorf("detach exit %s: %w", syscall, err))
+			return err
+		}
+	}
+	m.setLastError(syscall, nil)
+	return nil
+}
+
+func (m *Manager) States() []ProbeState {
+	if m == nil {
+		return nil
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	out := make([]ProbeState, 0, len(m.probes))
+	for syscall, entry := range m.probes {
+		state := ProbeState{
+			Syscall: syscall,
+			Active:  entry.active,
+		}
+		if entry.lastErr != nil {
+			state.Error = entry.lastErr.Error()
+		}
+		out = append(out, state)
+	}
+	sort.Slice(out, func(i, j int) bool { return out[i].Syscall < out[j].Syscall })
+	return out
+}
+
+func (m *Manager) ActiveCount() (active, total int) {
+	if m == nil {
+		return 0, 0
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	total = len(m.probes)
+	for _, entry := range m.probes {
+		if entry.active {
+			active++
+		}
+	}
+	return active, total
+}
+
+func (m *Manager) Close() error {
+	if m == nil {
+		return nil
+	}
+
+	m.mu.Lock()
+	if m.closed {
+		m.mu.Unlock()
+		return nil
+	}
+	type pairLinks struct {
+		syscall   string
+		enterLink Link
+		exitLink  Link
+	}
+	links := make([]pairLinks, 0, len(m.probes))
+	for syscall, entry := range m.probes {
+		links = append(links, pairLinks{
+			syscall:   syscall,
+			enterLink: entry.enterLink,
+			exitLink:  entry.exitLink,
+		})
+		entry.enterLink = nil
+		entry.exitLink = nil
+		entry.active = false
+		entry.lastErr = nil
+	}
+	m.closed = true
+	m.mu.Unlock()
+
+	var firstErr error
+	for _, l := range links {
+		var errForSyscall error
+		if l.enterLink != nil {
+			if err := l.enterLink.Destroy(); err != nil {
+				errForSyscall = err
+				if firstErr == nil {
+					firstErr = err
+				}
+			}
+		}
+		if l.exitLink != nil {
+			if err := l.exitLink.Destroy(); err != nil {
+				if errForSyscall == nil {
+					errForSyscall = err
+				}
+				if firstErr == nil {
+					firstErr = err
+				}
+			}
+		}
+		m.setLastError(l.syscall, errForSyscall)
+	}
+	return firstErr
+}
+
+func (m *Manager) entryLocked(syscall string) (*probeEntry, error) {
+	if m.closed {
+		return nil, errors.New("probe manager is closed")
+	}
+	if m.attacher == nil {
+		return nil, errors.New("probe manager has no attacher")
+	}
+	entry, ok := m.probes[syscall]
+	if !ok {
+		return nil, fmt.Errorf("unknown syscall %q", syscall)
+	}
+	return entry, nil
+}
+
+func (m *Manager) setLastError(syscall string, err error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	entry, ok := m.probes[syscall]
+	if !ok {
+		return
+	}
+	entry.lastErr = err
+}
+
+func attachPair(attacher Attacher, enterTP, exitTP string) (Link, Link, error) {
+	enterLink, err := attachOne(attacher, enterTP)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	exitLink, err := attachOne(attacher, exitTP)
+	if err != nil {
+		if enterLink != nil {
+			_ = enterLink.Destroy()
+		}
+		return nil, nil, err
+	}
+	return enterLink, exitLink, nil
+}
+
+func attachOne(attacher Attacher, tracepoint string) (Link, error) {
+	if tracepoint == "" {
+		return nil, nil
+	}
+	progName := "handle_" + tracepoint
+	prog, err := attacher.GetProgram(progName)
+	if err != nil {
+		return nil, fmt.Errorf("get program %s: %w", progName, err)
+	}
+	link, err := prog.AttachTracepoint("syscalls", tracepoint)
+	if err != nil {
+		return nil, fmt.Errorf("attach %s: %w", tracepoint, err)
+	}
+	return link, nil
+}
diff --git a/internal/probemanager/manager_test.go b/internal/probemanager/manager_test.go
new file mode 100644
index 0000000..44dc259
--- /dev/null
+++ b/internal/probemanager/manager_test.go
@@ -0,0 +1,153 @@
+package probemanager
+
+import (
+	"errors"
+	"testing"
+)
+
+type fakeLink struct {
+	destroyed int
+	err       error
+}
+
+func (l *fakeLink) Destroy() error {
+	l.destroyed++
+	return l.err
+}
+
+type fakeProgram struct {
+	tracepoint string
+	link       *fakeLink
+	err        error
+}
+
+func (p *fakeProgram) AttachTracepoint(_, name string) (Link, error) {
+	p.tracepoint = name
+	if p.err != nil {
+		return nil, p.err
+	}
+	if p.link == nil {
+		p.link = &fakeLink{}
+	}
+	return p.link, nil
+}
+
+type fakeAttacher struct {
+	programs map[string]*fakeProgram
+	errs     map[string]error
+}
+
+func (a *fakeAttacher) GetProgram(name string) (Program, error) {
+	if err, ok := a.errs[name]; ok {
+		return nil, err
+	}
+	p, ok := a.programs[name]
+	if !ok {
+		return nil, errors.New("missing program")
+	}
+	return p, nil
+}
+
+func TestManagerAttachAllToggleAndCounts(t *testing.T) {
+	attacher := &fakeAttacher{
+		programs: map[string]*fakeProgram{
+			"handle_sys_enter_read":  {},
+			"handle_sys_exit_read":   {},
+			"handle_sys_enter_write": {},
+			"handle_sys_exit_write":  {},
+		},
+		errs: map[string]error{},
+	}
+	mgr := NewManager(attacher)
+
+	err := mgr.AttachAll(func(tp string) bool { return tp == "sys_enter_read" || tp == "sys_exit_read" }, []string{
+		"sys_enter_read", "sys_exit_read", "sys_enter_write", "sys_exit_write",
+	})
+	if err != nil {
+		t.Fatalf("AttachAll returned error: %v", err)
+	}
+
+	active, total := mgr.ActiveCount()
+	if active != 1 || total != 2 {
+		t.Fatalf("unexpected counts active=%d total=%d", active, total)
+	}
+
+	states := mgr.States()
+	if len(states) != 2 {
+		t.Fatalf("expected 2 states, got %d", len(states))
+	}
+	if states[0].Syscall != "read" || !states[0].Active {
+		t.Fatalf("expected read active first, got %+v", states[0])
+	}
+	if states[1].Syscall != "write" || states[1].Active {
+		t.Fatalf("expected write inactive second, got %+v", states[1])
+	}
+
+	if err := mgr.Toggle("write"); err != nil {
+		t.Fatalf("Toggle(write) returned error: %v", err)
+	}
+	active, total = mgr.ActiveCount()
+	if active != 2 || total != 2 {
+		t.Fatalf("unexpected counts after toggle active=%d total=%d", active, total)
+	}
+}
+
+func TestManagerDetachDestroysLinks(t *testing.T) {
+	enter := &fakeLink{}
+	exit := &fakeLink{}
+	attacher := &fakeAttacher{
+		programs: map[string]*fakeProgram{
+			"handle_sys_enter_close": {link: enter},
+			"handle_sys_exit_close":  {link: exit},
+		},
+		errs: map[string]error{},
+	}
+	mgr := NewManager(attacher)
+	if err := mgr.AttachAll(nil, []string{"sys_enter_close", "sys_exit_close"}); err != nil {
+		t.Fatalf("AttachAll returned error: %v", err)
+	}
+	if err := mgr.Detach("close"); err != nil {
+		t.Fatalf("Detach returned error: %v", err)
+	}
+	if enter.destroyed != 1 || exit.destroyed != 1 {
+		t.Fatalf("expected both links destroyed once, got enter=%d exit=%d", enter.destroyed, exit.destroyed)
+	}
+}
+
+func TestManagerClosePreventsFurtherOperations(t *testing.T) {
+	attacher := &fakeAttacher{
+		programs: map[string]*fakeProgram{
+			"handle_sys_enter_open": {},
+			"handle_sys_exit_open":  {},
+		},
+		errs: map[string]error{},
+	}
+	mgr := NewManager(attacher)
+	if err := mgr.AttachAll(nil, []string{"sys_enter_open", "sys_exit_open"}); err != nil {
+		t.Fatalf("AttachAll returned error: %v", err)
+	}
+	if err := mgr.Close(); err != nil {
+		t.Fatalf("Close returned error: %v", err)
+	}
+	if err := mgr.Toggle("open"); err == nil {
+		t.Fatalf("expected Toggle to fail after Close")
+	}
+}
+
+func TestManagerAttachAllReturnsProgramError(t *testing.T) {
+	attacher := &fakeAttacher{
+		programs: map[string]*fakeProgram{},
+		errs: map[string]error{
+			"handle_sys_enter_read": errors.New("boom"),
+		},
+	}
+	mgr := NewManager(attacher)
+	err := mgr.AttachAll(nil, []string{"sys_enter_read", "sys_exit_read"})
+	if err == nil {
+		t.Fatalf("expected attach error")
+	}
+	states := mgr.States()
+	if len(states) != 1 || states[0].Error == "" {
+		t.Fatalf("expected state to capture attach error, got %+v", states)
+	}
+}