Improve Enhanced Session Recording flakiness (#64583)

* ensure all bpf events are processed before stopping, refactor tests a bit to make the stress test much more consistent

* prevent startSession/endSession blocking during shutdown

* make lost event tracking in tests work correctly between sub-tests

* Remove accidental extra wg.Wait call

* simplify LostEventTracker

* improved log message and switched to use a time.Timer

* fix not building without bpf tag

* drop events if the buffer is full

bpf/enhancedrecording/command.bpf.c

@@ -8,7 +8,7 @@
 // Size, in bytes, of the ring buffer used to report
 // audit events to userspace. This is the default,
 // the userspace can adjust this value based on config.
-#define EVENTS_BUF_SIZE (4096*8)
+#define EVENTS_BUF_SIZE (4096 * 64)
 
 char LICENSE[] SEC("license") = "Dual BSD/GPL";
 

lib/bpf/bpf.go

@@ -29,6 +29,7 @@ import (
 	"net"
 	"slices"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/cilium/ebpf/ringbuf"
@@ -44,9 +45,15 @@ import (
 	"github.com/gravitational/teleport/lib/utils"
 )
 
-// ArgsCacheSize is the number of args events to store before dropping args
-// events.
-const ArgsCacheSize = len(commandDataT{}.Argv)
+const (
+	// ArgsCacheSize is the number of args events to store before dropping args
+	// events.
+	ArgsCacheSize = len(commandDataT{}.Argv)
+
+	// eventSendTimeout is the maximum time to wait for an event to be sent
+	// to be emitted to the Audit log.
+	eventSendTimeout = 10 * time.Second
+)
 
 type sessionHandler interface {
 	startSession(auditSessionID uint32) error
@@ -82,6 +89,10 @@ type Service struct {
 	// conn is a BPF programs that hooks connect.
 	// conn is set only when restricted sessions are enabled.
 	conn *conn
+
+	lostEvents EventCount
+
+	wg sync.WaitGroup
 }
 
 // New creates a BPF service.
@@ -152,11 +163,31 @@ func New(config *servicecfg.BPFConfig) (bpf BPF, err error) {
 		"elapsed", time.Since(start),
 	)
 
-	go s.processNetworkEvents()
-
 	// Start pulling events off the perf buffers and emitting them to the
 	// Audit Log.
-	go s.processAccessEvents()
+	s.wg.Go(func() {
+		for event := range s.exec.events() {
+			s.emitCommandEvent(event)
+		}
+	})
+	s.wg.Go(func() {
+		for event := range s.open.events() {
+			s.emitDiskEvent(event)
+		}
+	})
+	s.wg.Go(func() {
+		for event := range s.conn.v4Events() {
+			s.emit4NetworkEvent(event)
+		}
+	})
+	s.wg.Go(func() {
+		for event := range s.conn.v6Events() {
+			s.emit6NetworkEvent(event)
+		}
+	})
+
+	// Log the number of lost events.
+	s.wg.Go(s.logLostEvents)
 
 	return s, nil
 }
@@ -177,9 +208,10 @@ func (s *Service) Close(restarting bool) error {
 		logger.WarnContext(s.closeContext, "Failed to close cgroup", "error", err)
 	}
 
-	// Signal to the processAccessEvents pulling events off the perf buffer to shutdown.
 	s.closeFunc()
 
+	s.wg.Wait()
+
 	return nil
 }
 
@@ -270,13 +302,27 @@ func (s *Service) Enabled() bool {
 	return true
 }
 
-func sendEvents(bpfEvents chan []byte, eventBuf *ringbuf.Reader) {
+// LostEvents returns the total number of lost events for command, disk,
+// and network events since the service was started.
+func (s *Service) LostEvents() EventCount {
+	return EventCount{
+		commandEvents: s.exec.lostCounter.Count(),
+		diskEvents:    s.open.lostCounter.Count(),
+		networkEvents: s.conn.lostCounter.Count(),
+	}
+}
+
+func sendEvents(eventType string, bpfEvents chan []byte, eventBuf *ringbuf.Reader) {
 	defer eventBuf.Close()
+	defer close(bpfEvents)
+
+	timer := time.NewTimer(eventSendTimeout)
+	defer timer.Stop()
 
 	for {
 		rec, err := eventBuf.Read()
 		if err != nil {
-			if errors.Is(err, ringbuf.ErrClosed) {
+			if errors.Is(err, ringbuf.ErrClosed) || errors.Is(err, ringbuf.ErrFlushed) {
 				logger.DebugContext(context.Background(), "Received signal, exiting")
 				return
 			}
@@ -284,39 +330,40 @@ func sendEvents(bpfEvents chan []byte, eventBuf *ringbuf.Reader) {
 			return
 		}
 
-		bpfEvents <- rec.RawSample[:]
-	}
-}
-
-// processAccessEvents pulls events off the perf ring buffer, parses them, and emits them to
-// the audit log.
-// TODO(capnspacehook): combine processAccessEvents and processNetworkEvents
-func (s *Service) processAccessEvents() {
-	for {
+		// Avoid blocking on the channel if the buffer is full, this
+		// could prevent the service from shutting down.
+		timer.Reset(eventSendTimeout)
 		select {
-		// Program execution.
-		case event := <-s.exec.events():
-			s.emitCommandEvent(event)
-		// Disk access.
-		case event := <-s.open.events():
-			s.emitDiskEvent(event)
-		case <-s.closeContext.Done():
-			return
+		case bpfEvents <- rec.RawSample[:]:
+		case <-timer.C:
+			logger.WarnContext(context.Background(), "Enhanced session recording event buffer is full, dropping event", "event_type", eventType)
 		}
 	}
 }
 
-// processNetworkEvents pulls networks events of the ring buffer and emits them
-// to the audit log.
-func (s *Service) processNetworkEvents() {
+func (s *Service) logLostEvents() {
+	const interval = 5 * time.Second
+	timer := time.NewTimer(interval)
+	defer timer.Stop()
+
 	for {
 		select {
-		// Network access (IPv4).
-		case event := <-s.conn.v4Events():
-			s.emit4NetworkEvent(event)
-		// Network access (IPv6).
-		case event := <-s.conn.v6Events():
-			s.emit6NetworkEvent(event)
+		case <-timer.C:
+			le := s.LostEvents()
+			newlyLost := le.Delta(s.lostEvents)
+			if !newlyLost.Empty() {
+				logger.WarnContext(s.closeContext, "Lost some Enhanced Session Recording events in the last 5 seconds due to a full eBPF ringbuffer, consider increasing the buffer sizes; see https://goteleport.com/docs/enroll-resources/server-access/guides/bpf-session-recording/#create-a-configuration-file for more information",
+					"command_events",
+					newlyLost.commandEvents,
+					"disk_events",
+					newlyLost.diskEvents,
+					"network_events",
+					newlyLost.networkEvents,
+				)
+			}
+
+			s.lostEvents = le
+			timer.Reset(interval)
 		case <-s.closeContext.Done():
 			return
 		}

lib/bpf/command.go

@@ -47,15 +47,15 @@ type exec struct {
 	// session
 	objs commandObjects
 
-	eventBuf *ringbuf.Reader
-	lost     *ebpf.Map
-	toClose  []io.Closer
-
-	closed bool
-	mtx    sync.Mutex
-
 	bpfEvents   chan []byte
 	lostCounter *Counter
+	toClose     []io.Closer
+
+	closed   bool
+	flushBuf func() error
+
+	mtx sync.Mutex
+	wg  sync.WaitGroup
 }
 
 // startExec will load, start, and pull events off the ring buffer
@@ -119,17 +119,17 @@ func startExec(bufferSize int) (*exec, error) {
 		return nil, trace.Wrap(err)
 	}
 
-	bpfEvents := make(chan []byte, bufferSize)
-	go sendEvents(bpfEvents, eventBuf)
-
-	return &exec{
+	e := &exec{
 		objs:        objs,
-		eventBuf:    eventBuf,
-		lost:        objs.LostCounter,
-		toClose:     toClose,
-		bpfEvents:   bpfEvents,
 		lostCounter: lostCtr,
-	}, nil
+		toClose:     toClose,
+		flushBuf:    eventBuf.Flush,
+	}
+
+	e.bpfEvents = make(chan []byte, bufferSize)
+	e.wg.Go(func() { sendEvents("command", e.bpfEvents, eventBuf) })
+
+	return e, nil
 }
 
 func (e *exec) startSession(auditSessionID uint32) error {
@@ -166,14 +166,20 @@ func (e *exec) endSession(auditSessionID uint32) error {
 // program. The ring buffer is closed as part of the module being closed.
 func (e *exec) close() {
 	e.mtx.Lock()
-	defer e.mtx.Unlock()
 
 	if e.closed {
+		e.mtx.Unlock()
 		return
 	}
 
 	e.closed = true
 
+	if err := e.flushBuf(); err != nil {
+		logger.WarnContext(context.Background(), "failed to flush command ring buffer", "error", err)
+	} else {
+		logger.DebugContext(context.Background(), "Flushed command ring buffer, waiting for pending events to be processed")
+	}
+
 	for _, link := range e.toClose {
 		if err := link.Close(); err != nil {
 			logger.WarnContext(context.Background(), "failed to close link", "error", err)
@@ -188,6 +194,11 @@ func (e *exec) close() {
 		logger.WarnContext(context.Background(), "failed to close command lost counter", "error", err)
 	}
 
+	// Unlock before waiting for the goroutines to finish to avoid
+	// startSession/endSession blocking for potentially a long time.
+	e.mtx.Unlock()
+	e.wg.Wait()
+
 	logger.DebugContext(context.Background(), "Closed command BPF module")
 }
 

lib/bpf/common.go

@@ -43,6 +43,10 @@ type BPF interface {
 
 	// Enabled returns whether enhanced recording is active.
 	Enabled() bool
+
+	// LostEvents returns the total number of lost events for command, disk,
+	// and network events since the service was started.
+	LostEvents() EventCount
 }
 
 // SessionContext contains all the information needed to track and emit
@@ -95,6 +99,40 @@ type SessionContext struct {
 	AuditSessionID uint32
 }
 
+// EventCount is a simple struct to track the number of events.
+type EventCount struct {
+	commandEvents uint64
+	diskEvents    uint64
+	networkEvents uint64
+}
+
+// Delta returns the number of events that have occurred since the previous
+// EventCount.
+func (e EventCount) Delta(prev EventCount) EventCount {
+	return EventCount{
+		commandEvents: e.commandEvents - prev.commandEvents,
+		diskEvents:    e.diskEvents - prev.diskEvents,
+		networkEvents: e.networkEvents - prev.networkEvents,
+	}
+}
+
+// Empty returns true if there are no events.
+func (e EventCount) Empty() bool {
+	return max(e.commandEvents, e.diskEvents, e.networkEvents) == 0
+}
+
+func (e EventCount) CommandEvents() uint64 {
+	return e.commandEvents
+}
+
+func (e EventCount) DiskEvents() uint64 {
+	return e.diskEvents
+}
+
+func (e EventCount) NetworkEvents() uint64 {
+	return e.networkEvents
+}
+
 // NOP is used on either non-Linux systems or when BPF support is not enabled.
 type NOP struct{}
 
@@ -117,6 +155,12 @@ func (s *NOP) Enabled() bool {
 	return false
 }
 
+// LostEvents returns the number of lost events. Note this function
+// does nothing.
+func (s *NOP) LostEvents() EventCount {
+	return EventCount{}
+}
+
 // IsHostCompatible checks that BPF programs can run on this host.
 func IsHostCompatible() error {
 	version, err := utils.KernelVersion()

lib/bpf/common_test.go

@@ -32,8 +32,9 @@ import (
 // provide values for the page sizes and hard coded values (like zero or a
 // specific page size) are respected when given.
 func TestBPFConfig_CheckAndSetDefaults(t *testing.T) {
-	perfBufferPageCount := defaults.PerfBufferPageCount
+	cmdBufferPageCount := defaults.CmdPerfBufferPageCount
 	openPerfBufferPageCount := defaults.OpenPerfBufferPageCount
+	netPerfBufferPageCount := defaults.NetPerfBufferPageCount
 	zeroPageCount := 0
 
 	tests := []struct {
@@ -45,9 +46,9 @@ func TestBPFConfig_CheckAndSetDefaults(t *testing.T) {
 			name: "all defaults",
 			got:  &servicecfg.BPFConfig{},
 			want: &servicecfg.BPFConfig{
-				CommandBufferSize: &perfBufferPageCount,
+				CommandBufferSize: &cmdBufferPageCount,
 				DiskBufferSize:    &openPerfBufferPageCount,
-				NetworkBufferSize: &perfBufferPageCount,
+				NetworkBufferSize: &netPerfBufferPageCount,
 				CgroupPath:        defaults.CgroupPath,
 			},
 		},
@@ -56,13 +57,13 @@ func TestBPFConfig_CheckAndSetDefaults(t *testing.T) {
 			got: &servicecfg.BPFConfig{
 				CommandBufferSize: &zeroPageCount,
 				DiskBufferSize:    &zeroPageCount,
-				NetworkBufferSize: &perfBufferPageCount,
+				NetworkBufferSize: &netPerfBufferPageCount,
 				CgroupPath:        "/my/cgroup/",
 			},
 			want: &servicecfg.BPFConfig{
-				CommandBufferSize: &perfBufferPageCount,
+				CommandBufferSize: &cmdBufferPageCount,
 				DiskBufferSize:    &openPerfBufferPageCount,
-				NetworkBufferSize: &perfBufferPageCount,
+				NetworkBufferSize: &netPerfBufferPageCount,
 				CgroupPath:        "/my/cgroup/",
 			},
 		},

lib/bpf/disk.go

@@ -47,13 +47,15 @@ var lostDiskEvents = prometheus.NewCounter(
 type open struct {
 	objs diskObjects
 
-	eventBuf chan []byte
-	toClose  []io.Closer
-
-	closed bool
-	mtx    sync.Mutex
-
+	bpfEvents   chan []byte
 	lostCounter *Counter
+	toClose     []io.Closer
+
+	closed   bool
+	flushBuf func() error
+
+	mtx sync.Mutex
+	wg  sync.WaitGroup
 }
 
 // startOpen will compile, load, start, and pull events off the perf buffer
@@ -142,15 +144,17 @@ func startOpen(bufferSize int) (*open, error) {
 		return nil, trace.Wrap(err, "creating ring buffer reader: %v", err)
 	}
 
-	bpfEvents := make(chan []byte, bufferSize)
-	go sendEvents(bpfEvents, eventBuf)
-
-	return &open{
+	o := &open{
 		objs:        objs,
-		eventBuf:    bpfEvents,
-		toClose:     toClose,
 		lostCounter: lostCtr,
-	}, nil
+		toClose:     toClose,
+		flushBuf:    eventBuf.Flush,
+	}
+
+	o.bpfEvents = make(chan []byte, bufferSize)
+	o.wg.Go(func() { sendEvents("disk", o.bpfEvents, eventBuf) })
+
+	return o, nil
 }
 
 func (o *open) startSession(auditSessionID uint32) error {
@@ -187,14 +191,20 @@ func (o *open) endSession(auditSessionID uint32) error {
 // program. The ring buffer is closed as part of the module being closed.
 func (o *open) close() {
 	o.mtx.Lock()
-	defer o.mtx.Unlock()
 
 	if o.closed {
+		o.mtx.Unlock()
 		return
 	}
 
 	o.closed = true
 
+	if err := o.flushBuf(); err != nil {
+		logger.WarnContext(context.Background(), "failed to flush disk ring buffer", "error", err)
+	} else {
+		logger.DebugContext(context.Background(), "Flushed disk ring buffer, waiting for pending events to be processed")
+	}
+
 	for _, toClose := range o.toClose {
 		if err := toClose.Close(); err != nil {
 			logger.WarnContext(context.Background(), "failed to close link", "error", err)
@@ -209,10 +219,15 @@ func (o *open) close() {
 		logger.WarnContext(context.Background(), "failed to close disk lost counter", "error", err)
 	}
 
+	// Unlock before waiting for the goroutines to finish to avoid
+	// startSession/endSession blocking for potentially a long time.
+	o.mtx.Unlock()
+	o.wg.Wait()
+
 	logger.DebugContext(context.Background(), "Closed disk BPF module")
 }
 
 // events contains raw events off the perf buffer.
 func (o *open) events() <-chan []byte {
-	return o.eventBuf
+	return o.bpfEvents
 }

lib/bpf/helper.go

@@ -24,6 +24,7 @@ import (
 	"context"
 	"errors"
 	"sync"
+	"sync/atomic"
 
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/ringbuf"
@@ -61,7 +62,7 @@ type Counter struct {
 	// arr is a one element array containing the value
 	arr *ebpf.Map
 	// lastCnt keeps the last read counter value
-	lastCnt uint64
+	lastCnt atomic.Uint64
 
 	// wg is used to wait for the loop goroutine to finish
 	wg sync.WaitGroup
@@ -96,6 +97,10 @@ func (c *Counter) Close() error {
 	return err
 }
 
+func (c *Counter) Count() uint64 {
+	return c.lastCnt.Load()
+}
+
 func (c *Counter) loop() {
 	for {
 		_, err := c.doorbellBuf.Read()
@@ -113,9 +118,11 @@ func (c *Counter) loop() {
 			logger.ErrorContext(context.Background(), "Error reading array value at index 0", "error", err)
 			continue
 		}
-		if delta := count - c.lastCnt; delta > 0 {
+		delta := count - c.lastCnt.Load()
+		if delta > 0 {
 			c.counter.Add(float64(delta))
+			c.lastCnt.Add(delta)
+			logger.DebugContext(context.Background(), "BPF events lost", "lost_events", delta)
 		}
-		c.lastCnt = count
 	}
 }

lib/bpf/network.go

@@ -22,6 +22,7 @@ package bpf
 
 import (
 	"context"
+	"errors"
 	"io"
 	"sync"
 
@@ -44,16 +45,18 @@ var lostNetworkEvents = prometheus.NewCounter(
 )
 
 type conn struct {
-	objs *networkObjects
-
-	event4Chan chan []byte
-	event6Chan chan []byte
-	toClose    []io.Closer
-
-	closed bool
-	mtx    sync.Mutex
+	objs networkObjects
 
+	bpf4Events  chan []byte
+	bpf6Events  chan []byte
 	lostCounter *Counter
+	toClose     []io.Closer
+
+	closed   bool
+	flushBuf func() error
+
+	mtx sync.Mutex
+	wg  sync.WaitGroup
 }
 
 func startConn(bufferSize int) (*conn, error) {
@@ -133,19 +136,24 @@ func startConn(bufferSize int) (*conn, error) {
 		return nil, trace.Wrap(err)
 	}
 
-	bpfv4Events := make(chan []byte, bufferSize)
-	go sendEvents(bpfv4Events, eventBufV4)
+	flush := func() error {
+		return errors.Join(eventBufV4.Flush(), eventBufV6.Flush())
+	}
 
-	bpfv6Events := make(chan []byte, bufferSize)
-	go sendEvents(bpfv6Events, eventBufV6)
-
-	return &conn{
-		objs:        &objs,
-		event4Chan:  bpfv4Events,
-		event6Chan:  bpfv6Events,
-		toClose:     toClose,
+	c := &conn{
+		objs:        objs,
 		lostCounter: lostCtr,
-	}, nil
+		toClose:     toClose,
+		flushBuf:    flush,
+	}
+
+	c.bpf4Events = make(chan []byte, bufferSize)
+	c.wg.Go(func() { sendEvents("network", c.bpf4Events, eventBufV4) })
+
+	c.bpf6Events = make(chan []byte, bufferSize)
+	c.wg.Go(func() { sendEvents("network", c.bpf6Events, eventBufV6) })
+
+	return c, nil
 }
 
 func (c *conn) startSession(auditSessionID uint32) error {
@@ -181,16 +189,21 @@ func (c *conn) endSession(auditSessionID uint32) error {
 // close will stop reading events off the ring buffer and unload the BPF
 // program. The ring buffer is closed as part of the module being closed.
 func (c *conn) close() {
-	// c.lost.Close()
 	c.mtx.Lock()
-	defer c.mtx.Unlock()
 
 	if c.closed {
+		c.mtx.Unlock()
 		return
 	}
 
 	c.closed = true
 
+	if err := c.flushBuf(); err != nil {
+		logger.WarnContext(context.Background(), "failed to flush network ring buffer", "error", err)
+	} else {
+		logger.DebugContext(context.Background(), "Flushed network ring buffer, waiting for pending events to be processed")
+	}
+
 	for _, link := range c.toClose {
 		if err := link.Close(); err != nil {
 			logger.WarnContext(context.Background(), "failed to close link", "error", err)
@@ -205,15 +218,20 @@ func (c *conn) close() {
 		logger.WarnContext(context.Background(), "failed to close network lost counter", "error", err)
 	}
 
+	// Unlock before waiting for the goroutines to finish to avoid
+	// startSession/endSession blocking for potentially a long time.
+	c.mtx.Unlock()
+	c.wg.Wait()
+
 	logger.DebugContext(context.Background(), "Closed network BPF module")
 }
 
 // v4Events contains raw events off the perf buffer.
 func (c *conn) v4Events() <-chan []byte {
-	return c.event4Chan
+	return c.bpf4Events
 }
 
 // v6Events contains raw events off the perf buffer.
 func (c *conn) v6Events() <-chan []byte {
-	return c.event6Chan
+	return c.bpf6Events
 }

lib/defaults/defaults.go

@@ -549,13 +549,19 @@ func ReadableDatabaseProtocol(p string) string {
 }
 
 const (
-	// PerfBufferPageCount is the size of the perf ring buffer.
-	PerfBufferPageCount = 64
+	// CmdPerfBufferPageCount is the size of buffered channel that receives
+	// eBPF command event messages.
+	CmdPerfBufferPageCount = 64
 
-	// OpenPerfBufferPageCount is the page count for the perf buffer. Open
-	// events generate many events so this buffer needs to be extra large.
+	// OpenPerfBufferPageCount is the size of buffered channel that receives
+	// eBPF disk event messages. Open events generate many events so this
+	// buffer needs to be extra large.
 	OpenPerfBufferPageCount = 128
 
+	// NetPerfBufferPageCount is the size of the buffered channel that receives
+	// eBPF network event messages.
+	NetPerfBufferPageCount = 64
+
 	// CgroupPath is where the cgroupv2 hierarchy will be mounted.
 	CgroupPath = "/cgroup2"
 )
@@ -832,7 +838,7 @@ type httpClientOptions struct {
 // settings from the environment variables HTTP_PROXY, HTTPS_PROXY, and NO_PROXY.
 func UseProxyFromEnvironment() HTTPClientOption {
 	return func(opts *httpClientOptions) *httpClientOptions {
-		var enable = true
+		enable := true
 		opts.useProxyFromEnvironment = &enable
 		return opts
 	}
@@ -842,7 +848,7 @@ func UseProxyFromEnvironment() HTTPClientOption {
 // settings from the environment variables HTTP_PROXY, HTTPS_PROXY, and NO_PROXY.
 func DisableProxyFromEnvironment() HTTPClientOption {
 	return func(opts *httpClientOptions) *httpClientOptions {
-		var enable = false
+		enable := false
 		opts.useProxyFromEnvironment = &enable
 		return opts
 	}

lib/service/servicecfg/bpf.go

@@ -47,15 +47,16 @@ type BPFConfig struct {
 
 // CheckAndSetDefaults checks BPF configuration.
 func (c *BPFConfig) CheckAndSetDefaults() error {
-	perfBufferPageCount := defaults.PerfBufferPageCount
+	cmdBufferPageCount := defaults.CmdPerfBufferPageCount
 	openPerfBufferPageCount := defaults.OpenPerfBufferPageCount
+	netPerfBufferPageCount := defaults.NetPerfBufferPageCount
 
 	// Set defaults for buffer sizes if they are unset or zero.
 	// A zero value was accepted before but is undesirable now as it
 	// will result in blocking event channels, so we set it to a sane
 	// default to maintain backwards compatibility.
 	if c.CommandBufferSize == nil || *c.CommandBufferSize == 0 {
-		c.CommandBufferSize = &perfBufferPageCount
+		c.CommandBufferSize = &cmdBufferPageCount
 	} else if *c.CommandBufferSize < 0 {
 		return trace.BadParameter("CommandBufferSize must not be negative")
 	}
@@ -65,7 +66,7 @@ func (c *BPFConfig) CheckAndSetDefaults() error {
 		return trace.BadParameter("DiskBufferSize must not be negative")
 	}
 	if c.NetworkBufferSize == nil || *c.NetworkBufferSize == 0 {
-		c.NetworkBufferSize = &perfBufferPageCount
+		c.NetworkBufferSize = &netPerfBufferPageCount
 	} else if *c.NetworkBufferSize < 0 {
 		return trace.BadParameter("NetworkBufferSize must not be negative")
 	}

lib/srv/bpf_test.go

@@ -63,9 +63,15 @@ const (
 
 	longArgBase = "averylongargument"
 
+	// increased buffer sizes to accommodate the onslaught of events
+	// during the stress test
+	stressCommandBufSize = 256
+	stressDiskBufSize    = 512
+	stressNetworkBufSize = 256
 	// number of commands that will be run in parallel during the
 	// stress test test case
-	stressTestRunCount = 10
+	// This was the max that would pass consistently
+	stressTestRunCount = 35
 )
 
 var (
@@ -115,10 +121,16 @@ type addrInfo struct {
 	port int
 }
 
+type bpfTestCase struct {
+	name       string
+	command    string
+	eventInfos []expectedEvents
+}
+
 func TestBPFRecording(t *testing.T) {
 	skipIfNoBPF(t)
 
-	srv, bpfSrv := newServices(t)
+	srv, bpfSrv := newServices(t, nil)
 
 	testBPFRecording(t, srv, bpfSrv)
 }
@@ -130,7 +142,7 @@ func TestBPFRecordingWithPAM(t *testing.T) {
 	skipIfNoBPF(t)
 	skipIfNoPAM(t)
 
-	srv, bpfSrv := newServices(t)
+	srv, bpfSrv := newServices(t, nil)
 	srv.pamCfg = &pamcfg.PAMConfig{
 		Enabled:     true,
 		ServiceName: "sshd",
@@ -150,6 +162,8 @@ func TestBPFRecordingWithPAM(t *testing.T) {
 // Filesystem Hierarchy Standard (FHS) or are specified in their
 // respective man pages.
 func testBPFRecording(t *testing.T, srv Server, bpfSrv bpf.BPF) {
+	var lostEvents bpf.EventCount
+
 	// Create a temp dir and files for commands to use.
 	cmdDir := t.TempDir()
 	tempFilePath := filepath.Join(cmdDir, "file")
@@ -165,7 +179,7 @@ func testBPFRecording(t *testing.T, srv Server, bpfSrv bpf.BPF) {
 	// Create a TCP listener for each IP family. Register the listeners
 	// to be closed in case the test fails before the connection
 	// handling goroutines are started, the listeners will be closed
-	// before waitinf for the goroutines to finish otherwise.
+	// before waiting for the goroutines to finish otherwise.
 	lis4, err := net.Listen("tcp4", "127.0.0.1:0")
 	require.NoError(t, err)
 	t.Cleanup(func() {
@@ -223,11 +237,7 @@ eval $(echo %s | base64 --decode)`,
 	}
 
 	// Define the test cases.
-	tests := []struct {
-		name       string
-		command    string
-		eventInfos []expectedEvents
-	}{
+	tests := []bpfTestCase{
 		{
 			name:    "no commands",
 			command: "true",
@@ -509,162 +519,251 @@ eval $(echo %s | base64 --decode)`,
 				},
 			},
 		},
-		{
-			name:    "stress test",
-			command: fmt.Sprintf("for i in $(seq 1 %d); do { curl %s; ls -lahiZ; } & done; wait", stressTestRunCount, netAddr4),
-			eventInfos: []expectedEvents{
-				{
-					cmdInfo: commandInfo{
-						program: "ls",
-						args:    []string{"-lahiZ"},
-					},
-					paths: []string{
-						"/proc/filesystems",
-						".",
-						"/etc/nsswitch.conf",
-						"/etc/passwd",
-						"/etc/group",
-						"/etc/localtime",
-					},
-					count: stressTestRunCount,
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			runBPFTestCase(t, srv, bpfSrv, tt, false)
+
+			le := bpfSrv.LostEvents()
+			newlyLost := le.Delta(lostEvents)
+			if newlyLost.CommandEvents() > 0 {
+				// TODO(capnspacehook): make error once new trace points are in place
+				t.Logf("error: %d command events were lost", newlyLost.CommandEvents())
+			}
+			if newlyLost.DiskEvents() > 0 {
+				t.Errorf("error: %d disk events were lost", newlyLost.DiskEvents())
+			}
+			if newlyLost.NetworkEvents() > 0 {
+				t.Errorf("error: %d network events were lost", newlyLost.NetworkEvents())
+			}
+
+			lostEvents = le
+		})
+	}
+}
+
+func runBPFTestCase(t *testing.T, srv Server, bpfSrv bpf.BPF, tt bpfTestCase, closeBPFSrv bool) {
+	expectedCmdFail := slices.ContainsFunc(tt.eventInfos, func(info expectedEvents) bool {
+		return info.cmdInfo.expectedFail
+	})
+
+	// Run the command and capture the events.
+	recordedEvents := runCommand(t, srv, bpfSrv, tt.command, expectedCmdFail, recordAllEvents, closeBPFSrv)
+
+	commandArgs := make(map[string]int)
+	programPaths := make(map[string]string)
+	programLibs := make(map[string][]countedValue[string])
+	commandOpens := make(map[string][]countedValue[string])
+	commandDstAddrs := make(map[string][]countedValue[addrInfo])
+
+	for _, info := range tt.eventInfos {
+		cmdInfo := info.cmdInfo
+		count := max(info.count, 1)
+
+		// Build a map of expected command arguments.
+		cmdArgKey := commandKey(cmdInfo.program, cmdInfo.args)
+		commandArgs[cmdArgKey] = count
+
+		// Build a map of program paths on disk. Also build a
+		// map of dynamic libraries that should be opened.
+		// We conservatively only expect to see libraries loaded
+		// in the DT_NEEDED section of the program binary.
+		// This gives us more disk events to check against for
+		// free, which can be useful to catch rare bugs in the
+		// BPF disk tracing program.
+		if info.cmdInfo.scriptPath == "" {
+			programPath, err := exec.LookPath(cmdInfo.program)
+			require.NoError(t, err, "%s command is required for these tests but was not found", cmdInfo.program)
+			programPaths[cmdInfo.program] = programPath
+
+			_, ok := programLibs[cmdInfo.program]
+			if !ok {
+				programLibs[cmdInfo.program] = getProgramLibs(t, programPath, count)
+			}
+		} else {
+			// If a script is being run, the command event will
+			// show the script as the program, but it will act
+			// like the interpreter.
+			programPaths[cmdInfo.program] = info.cmdInfo.scriptPath
+			interpPath, err := exec.LookPath(cmdInfo.interpreter)
+			require.NoError(t, err)
+
+			_, ok := programLibs[cmdInfo.program]
+			if !ok {
+				programLibs[cmdInfo.program] = getProgramLibs(t, interpPath, count)
+			}
+		}
+
+		expectedPaths := commandOpens[cmdInfo.program]
+		commandOpens[cmdInfo.program] = append(expectedPaths, makeCounted(info.paths, count)...)
+
+		// Build a map of program destination addresses.
+		if info.dstAddr != nil {
+			addrs := commandDstAddrs[cmdInfo.program]
+			commandDstAddrs[cmdInfo.program] = append(addrs, countedValue[addrInfo]{value: *info.dstAddr, count: count})
+		}
+	}
+
+	// Check that the emitted events have expected contents.
+	for _, event := range recordedEvents {
+		switch e := event.(type) {
+		case *apievents.SessionCommand:
+			t.Logf("  command event: Command=%s Path=%s Args=[%s] NArgs=%d", e.BPFMetadata.Program, e.Path, quoteStrings(e.Argv), len(e.Argv))
+
+			checkCommandEvent(t, e, programPaths, commandArgs)
+		case *apievents.SessionDisk:
+			t.Logf("  disk event: Command=%s Path=%s", e.BPFMetadata.Program, e.Path)
+
+			if checkDiskEvent(t, e, programLibs, true) {
+				continue
+			}
+			checkDiskEvent(t, e, commandOpens, false)
+		case *apievents.SessionNetwork:
+			t.Logf("  network event: Command=%s SrcAddr=%s DstAddr=%s DstPort=%d", e.BPFMetadata.Program, e.SrcAddr, e.DstAddr, e.DstPort)
+
+			checkNetworkEvent(t, e, commandDstAddrs)
+		}
+	}
+
+	// Check that expected events were found the expected number of times.
+	for cmd, count := range commandArgs {
+		if count > 0 {
+			t.Errorf("error: command event for %s was expected %d more times", cmd, count)
+		}
+	}
+	for cmd, paths := range commandOpens {
+		for _, path := range paths {
+			if path.count > 0 {
+				t.Errorf("error: disk event for program %q opening %q was expected %d more times", cmd, path.value, path.count)
+			}
+		}
+	}
+	for cmd, addrs := range commandDstAddrs {
+		for _, addr := range addrs {
+			if addr.count > 0 {
+				t.Errorf("error: network event for program %q with destination address %q was expected %d more times", cmd, addr.value, addr.count)
+			}
+		}
+	}
+	for cmd, libs := range programLibs {
+		for _, lib := range libs {
+			if lib.count > 0 {
+				t.Errorf("error: disk event for program %q opening library %q was expected %d more times", cmd, lib.value, lib.count)
+			}
+		}
+	}
+}
+
+func TestBPFStress(t *testing.T) {
+	skipIfNoBPF(t)
+
+	cmdBufSize := stressCommandBufSize
+	diskBufSize := stressDiskBufSize
+	networkBufSize := stressNetworkBufSize
+	cfg := &servicecfg.BPFConfig{
+		Enabled:           true,
+		CommandBufferSize: &cmdBufSize,
+		DiskBufferSize:    &diskBufSize,
+		NetworkBufferSize: &networkBufSize,
+	}
+
+	srv, bpfSrv := newServices(t, cfg)
+
+	testBPFStress(t, srv, bpfSrv)
+}
+
+// TODO(capnspacehook): test with PAM auth enabled, and with a different
+// login user once https://github.com/gravitational/teleport/issues/61692
+// is fixed.
+func TestBPFStressWithPAM(t *testing.T) {
+	skipIfNoBPF(t)
+	skipIfNoPAM(t)
+
+	cmdBufSize := stressCommandBufSize
+	diskBufSize := stressDiskBufSize
+	networkBufSize := stressNetworkBufSize
+	cfg := &servicecfg.BPFConfig{
+		Enabled:           true,
+		CommandBufferSize: &cmdBufSize,
+		DiskBufferSize:    &diskBufSize,
+		NetworkBufferSize: &networkBufSize,
+	}
+
+	srv, bpfSrv := newServices(t, cfg)
+	srv.pamCfg = &pamcfg.PAMConfig{
+		Enabled:     true,
+		ServiceName: "sshd",
+	}
+
+	testBPFStress(t, srv, bpfSrv)
+}
+
+func testBPFStress(t *testing.T, srv Server, bpfSrv bpf.BPF) {
+	// Create a TCP listener. Register the listeners to be closed in
+	// case the test fails before the connection handling goroutines
+	// are started, the listeners will be closed before waiting for
+	// the goroutines to finish otherwise.
+	lis, err := net.Listen("tcp4", "127.0.0.1:0")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		lis.Close()
+	})
+
+	netAddr := lis.Addr().String()
+	addr, portStr, err := net.SplitHostPort(netAddr)
+	require.NoError(t, err)
+	port, err := strconv.Atoi(portStr)
+	require.NoError(t, err)
+
+	var wg sync.WaitGroup
+	wg.Go(func() { handleConnections(lis) })
+	t.Cleanup(func() {
+		lis.Close()
+		wg.Wait()
+	})
+
+	tt := bpfTestCase{
+		name:    "stress test",
+		command: fmt.Sprintf("for i in $(seq 1 %d); do { curl %s; ls -lahiZ; } & done; wait", stressTestRunCount, netAddr),
+		eventInfos: []expectedEvents{
+			{
+				cmdInfo: commandInfo{
+					program: "ls",
+					args:    []string{"-lahiZ"},
 				},
-				{
-					cmdInfo: commandInfo{
-						program: "curl",
-						args:    []string{netAddr4},
-					},
-					dstAddr: &addrInfo{addr: addr4, port: port4},
-					count:   stressTestRunCount,
+				paths: []string{
+					"/proc/filesystems",
+					"/etc/nsswitch.conf",
+					"/etc/passwd",
+					"/etc/group",
 				},
-				{
-					cmdInfo: commandInfo{
-						program: "seq",
-						args:    []string{"1", strconv.Itoa(stressTestRunCount)},
-					},
+				count: stressTestRunCount,
+			},
+			{
+				cmdInfo: commandInfo{
+					program: "curl",
+					args:    []string{netAddr},
+				},
+				dstAddr: &addrInfo{addr: addr, port: port},
+				count:   stressTestRunCount,
+			},
+			{
+				cmdInfo: commandInfo{
+					program: "seq",
+					args:    []string{"1", strconv.Itoa(stressTestRunCount)},
 				},
 			},
 		},
 	}
 
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			expectedCmdFail := slices.ContainsFunc(tt.eventInfos, func(info expectedEvents) bool {
-				return info.cmdInfo.expectedFail
-			})
-
-			// Run the command and capture the events.
-			recordedEvents := runCommand(t, srv, bpfSrv, tt.command, expectedCmdFail, recordAllEvents)
-
-			commandArgs := make(map[string]int)
-			programPaths := make(map[string]string)
-			programLibs := make(map[string][]countedValue[string])
-			commandOpens := make(map[string][]countedValue[string])
-			commandDstAddrs := make(map[string][]countedValue[addrInfo])
-
-			for _, info := range tt.eventInfos {
-				cmdInfo := info.cmdInfo
-				count := max(info.count, 1)
-
-				// Build a map of expected command arguments.
-				cmdArgKey := commandKey(cmdInfo.program, cmdInfo.args)
-				commandArgs[cmdArgKey] = count
-
-				// Build a map of program paths on disk. Also build a
-				// map of dynamic libraries that should be opened.
-				// We conservatively only expect to see libraries loaded
-				// in the DT_NEEDED section of the program binary.
-				// This gives us more disk events to check against for
-				// free, which can be useful to catch rare bugs in the
-				// BPF disk tracing program.
-				if info.cmdInfo.scriptPath == "" {
-					programPath, err := exec.LookPath(cmdInfo.program)
-					require.NoError(t, err, "%s command is required for these tests but was not found", cmdInfo.program)
-					programPaths[cmdInfo.program] = programPath
-
-					_, ok := programLibs[cmdInfo.program]
-					if !ok {
-						programLibs[cmdInfo.program] = getProgramLibs(t, programPath, count)
-					}
-				} else {
-					// If a script is being run, the command event will
-					// show the script as the program, but it will act
-					// like the interpreter.
-					programPaths[cmdInfo.program] = info.cmdInfo.scriptPath
-					interpPath, err := exec.LookPath(cmdInfo.interpreter)
-					require.NoError(t, err)
-
-					_, ok := programLibs[cmdInfo.program]
-					if !ok {
-						programLibs[cmdInfo.program] = getProgramLibs(t, interpPath, count)
-					}
-				}
-
-				expectedPaths := commandOpens[cmdInfo.program]
-				commandOpens[cmdInfo.program] = append(expectedPaths, makeCounted(info.paths, count)...)
-
-				// Build a map of program destination addresses.
-				if info.dstAddr != nil {
-					addrs := commandDstAddrs[cmdInfo.program]
-					commandDstAddrs[cmdInfo.program] = append(addrs, countedValue[addrInfo]{value: *info.dstAddr, count: count})
-				}
-			}
-
-			// Check that the emitted events have expected contents.
-			for _, event := range recordedEvents {
-				switch e := event.(type) {
-				case *apievents.SessionCommand:
-					t.Logf("  command event: Command=%s Path=%s Args=[%s] NArgs=%d", e.BPFMetadata.Program, e.Path, quoteStrings(e.Argv), len(e.Argv))
-
-					checkCommandEvent(t, e, programPaths, commandArgs)
-				case *apievents.SessionDisk:
-					t.Logf("  disk event: Command=%s Path=%s", e.BPFMetadata.Program, e.Path)
-
-					if checkDiskEvent(t, e, programLibs, true) {
-						continue
-					}
-					checkDiskEvent(t, e, commandOpens, false)
-				case *apievents.SessionNetwork:
-					t.Logf("  network event: Command=%s SrcAddr=%s DstAddr=%s DstPort=%d", e.BPFMetadata.Program, e.SrcAddr, e.DstAddr, e.DstPort)
-
-					checkNetworkEvent(t, e, commandDstAddrs)
-				}
-			}
-
-			// Check that expected events were found the expected number of times.
-			for cmd, count := range commandArgs {
-				if count > 0 {
-					t.Errorf("error: command event for %s was expected %d more times", cmd, count)
-				}
-			}
-			for cmd, paths := range commandOpens {
-				for _, path := range paths {
-					if path.count > 0 {
-						t.Errorf("error: disk event for program %q opening %q was expected %d more times", cmd, path.value, path.count)
-					}
-				}
-			}
-			for cmd, addrs := range commandDstAddrs {
-				for _, addr := range addrs {
-					if addr.count > 0 {
-						t.Errorf("error: network event for program %q with destination address %q was expected %d more times", cmd, addr.value, addr.count)
-					}
-				}
-			}
-			for cmd, libs := range programLibs {
-				for _, lib := range libs {
-					if lib.count > 0 {
-						t.Errorf("error: disk event for program %q opening library %q was expected %d more times", cmd, lib.value, lib.count)
-					}
-				}
-			}
-		})
-	}
+	runBPFTestCase(t, srv, bpfSrv, tt, true)
 }
 
 func TestBPFMonitoring(t *testing.T) {
 	skipIfNoBPF(t)
 
-	srv, bpfSrv := newServices(t)
+	srv, bpfSrv := newServices(t, nil)
 
 	testBPFMonitoring(t, srv, bpfSrv)
 }
@@ -676,7 +775,7 @@ func TestBPFMonitoringWithPAM(t *testing.T) {
 	skipIfNoBPF(t)
 	skipIfNoPAM(t)
 
-	srv, bpfSrv := newServices(t)
+	srv, bpfSrv := newServices(t, nil)
 	srv.pamCfg = &pamcfg.PAMConfig{
 		Enabled:     true,
 		ServiceName: "sshd",
@@ -702,7 +801,7 @@ func testBPFMonitoring(t *testing.T, srv Server, bpfSrv bpf.BPF) {
 	// commands will.
 	eventsCh := make(chan []apievents.AuditEvent)
 	wg.Go(func() {
-		eventsCh <- runCommand(t, srv, bpfSrv, "sleep 3", false, recordAllEvents)
+		eventsCh <- runCommand(t, srv, bpfSrv, "sleep 3", false, recordAllEvents, false)
 	})
 
 	// Run curl commands that the bpf programs should ignore.
@@ -757,7 +856,7 @@ func testBPFMonitoring(t *testing.T, srv Server, bpfSrv bpf.BPF) {
 func TestBPFRoleOptions(t *testing.T) {
 	skipIfNoBPF(t)
 
-	srv, bpfSrv := newServices(t)
+	srv, bpfSrv := newServices(t, nil)
 
 	lis, err := net.Listen("tcp4", "127.0.0.1:0")
 	require.NoError(t, err)
@@ -822,7 +921,7 @@ func TestBPFRoleOptions(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			// Run the command and capture the events.
-			recordedEvents := runCommand(t, srv, bpfSrv, command, true, tt.events)
+			recordedEvents := runCommand(t, srv, bpfSrv, command, true, tt.events, false)
 
 			// Check that only configured events were recorded.
 			if len(tt.events) == 0 {
@@ -882,10 +981,14 @@ func commandKey(program string, args []string) string {
 	return fmt.Sprintf("%s [%s]", program, quoteStrings(args))
 }
 
-func newServices(t *testing.T) (*mockServer, bpf.BPF) {
+func newServices(t *testing.T, cfg *servicecfg.BPFConfig) (*mockServer, bpf.BPF) {
 	t.Helper()
 
-	bpfSrv, err := bpf.New(&servicecfg.BPFConfig{Enabled: true})
+	if cfg == nil {
+		cfg = &servicecfg.BPFConfig{}
+	}
+	cfg.Enabled = true
+	bpfSrv, err := bpf.New(cfg)
 	require.NoError(t, err)
 
 	t.Cleanup(func() {
@@ -913,7 +1016,7 @@ func handleConnections(l net.Listener) {
 
 // runCommand runs the given command with Enhanced Session Recording
 // enabled and returns the recorded events.
-func runCommand(t *testing.T, srv Server, bpfSrv bpf.BPF, command string, expectedCmdFail bool, recordEvents map[string]struct{}) []apievents.AuditEvent {
+func runCommand(t *testing.T, srv Server, bpfSrv bpf.BPF, command string, expectedCmdFail bool, recordEvents map[string]struct{}, closeBPFSrv bool) []apievents.AuditEvent {
 	ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
 	t.Cleanup(cancel)
 
@@ -1018,6 +1121,12 @@ func runCommand(t *testing.T, srv Server, bpfSrv bpf.BPF, command string, expect
 		}
 	}
 
+	// Close the BPF service if requested, this will flush the ebpf ring
+	// buffers that hold events and ensure all pending events are processed.
+	if closeBPFSrv {
+		require.NoError(t, bpfSrv.Close(true))
+	}
+
 	return emitter.Events()
 }
 

lib/srv/mock_test.go

@@ -479,3 +479,7 @@ func (f fakeBPF) Close(restarting bool) error {
 func (f fakeBPF) Enabled() bool {
 	return true
 }
+
+func (f fakeBPF) LostEvents() bpf.EventCount {
+	return bpf.EventCount{}
+}