aquasecurity · geyslan · Mar 16, 2023 · Mar 1, 2023
diff --git a/cmd/tracee-rules/main.go b/cmd/tracee-rules/main.go
@@ -159,13 +159,13 @@ func main() {
 				return err
 			}
 
-			if httpServer != nil {
-				go httpServer.Start()
-			}
-
 			ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
 			defer stop()
 
+			if httpServer != nil {
+				go httpServer.Start(ctx)
+			}
+
 			e.Start(ctx)
 
 			return nil

diff --git a/pkg/cmd/tracee.go b/pkg/cmd/tracee.go
@@ -37,7 +37,7 @@ func (r Runner) Run(ctx context.Context) error {
 			}
 		}
 
-		go r.Server.Start()
+		go r.Server.Start(ctx)
 	}
 
 	// Print statistics at the end

diff --git a/pkg/ebpf/bpf_log.go b/pkg/ebpf/bpf_log.go
@@ -1,6 +1,7 @@
 package ebpf
 
 import (
+	"context"
 	"encoding/binary"
 	"fmt"
 	"strings"
@@ -158,7 +159,10 @@ func (b *BPFLog) Decode(rawBuffer []byte) error {
 	return nil
 }
 
-func (t *Tracee) processBPFLogs() {
+func (t *Tracee) processBPFLogs(ctx context.Context) {
+	logger.Debugw("Starting processBPFLogs go routine")
+	defer logger.Debugw("Stopped processBPFLogs go routine")
+
 	for {
 		select {
 		case rawData := <-t.bpfLogsChannel:
@@ -193,6 +197,9 @@ func (t *Tracee) processBPFLogs() {
 				t.stats.LostBPFLogsCount.Increment(lost)
 				logger.Warnw(fmt.Sprintf("Lost %d ebpf logs events", lost))
 			}
+
+		case <-ctx.Done():
+			return
 		}
 	}
 }
diff --git a/pkg/ebpf/events_processor.go b/pkg/ebpf/events_processor.go
@@ -2,6 +2,7 @@ package ebpf
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"path/filepath"
 	"strconv"
@@ -26,15 +27,22 @@ func init() {
 	initKernelReadFileTypes()
 }
 
-func (t *Tracee) processLostEvents() {
+func (t *Tracee) processLostEvents(ctx context.Context) {
+	logger.Debugw("Starting processLostEvents go routine")
+	defer logger.Debugw("Stopped processLostEvents go routine")
+
 	for {
-		lost := <-t.lostEvChannel
-		// When terminating tracee-ebpf the lost channel receives multiple "0 lost events" events.
-		// This check prevents those 0 lost events messages to be written to stderr until the bug is fixed:
-		// https://github.com/aquasecurity/libbpfgo/issues/122
-		if lost > 0 {
-			t.stats.LostEvCount.Increment(lost)
-			logger.Warnw(fmt.Sprintf("Lost %d events", lost))
+		select {
+		case lost := <-t.lostEvChannel:
+			// When terminating tracee-ebpf the lost channel receives multiple "0 lost events" events.
+			// This check prevents those 0 lost events messages to be written to stderr until the bug is fixed:
+			// https://github.com/aquasecurity/libbpfgo/issues/122
+			if lost > 0 {
+				_ = t.stats.LostEvCount.Increment(lost)
+				logger.Warnw(fmt.Sprintf("Lost %d events", lost))
+			}
+		case <-ctx.Done():
+			return
 		}
 	}
 }

diff --git a/pkg/ebpf/net_capture.go b/pkg/ebpf/net_capture.go
@@ -1,7 +1,7 @@
 package ebpf
 
 import (
-	gocontext "context"
+	"context"
 	"encoding/binary"
 	"fmt"
 
@@ -29,7 +29,10 @@ const (
 	familyIpv6
 )
 
-func (t *Tracee) processNetCaptureEvents(ctx gocontext.Context) {
+func (t *Tracee) processNetCaptureEvents(ctx context.Context) {
+	logger.Debugw("Starting processNetCaptureEvents go routine")
+	defer logger.Debugw("Stopped processNetCaptureEvents go routine")
+
 	var errChanList []<-chan error
 
 	// source pipeline stage (re-used from regular pipeline)
@@ -44,7 +47,7 @@ func (t *Tracee) processNetCaptureEvents(ctx gocontext.Context) {
 	t.WaitForPipeline(errChanList...)
 }
 
-func (t *Tracee) processNetCapEvents(ctx gocontext.Context, in <-chan *trace.Event) <-chan error {
+func (t *Tracee) processNetCapEvents(ctx context.Context, in <-chan *trace.Event) <-chan error {
 	errc := make(chan error, 1)
 
 	go func() {

diff --git a/pkg/ebpf/tracee.go b/pkg/ebpf/tracee.go
@@ -1359,18 +1359,18 @@ func (t *Tracee) Run(ctx gocontext.Context) error {
 		logger.Warnw("Memory dump", "error", err)
 	}
 	t.eventsPerfMap.Start()
-	go t.processLostEvents()
+	go t.processLostEvents(ctx)
 	go t.handleEvents(ctx)
 	if t.config.BlobPerfBufferSize > 0 {
 		t.fileWrPerfMap.Start()
-		go t.processFileWrites()
+		go t.processFileWrites(ctx)
 	}
 	if pcaps.PcapsEnabled(t.config.Capture.Net) {
 		t.netCapPerfMap.Start()
 		go t.processNetCaptureEvents(ctx)
 	}
 	t.bpfLogsPerfMap.Start()
-	go t.processBPFLogs()
+	go t.processBPFLogs(ctx)
 
 	// write pid file
 	err = t.writePid()

diff --git a/pkg/ebpf/write_capture.go b/pkg/ebpf/write_capture.go
@@ -1,6 +1,7 @@
 package ebpf
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -12,7 +13,9 @@ import (
 	"github.com/aquasecurity/tracee/pkg/utils"
 )
 
-func (t *Tracee) processFileWrites() {
+func (t *Tracee) processFileWrites(ctx context.Context) {
+	logger.Debugw("Starting processFileWrites go routine")
+	defer logger.Debugw("Stopped processFileWrites go routine")
 
 	const (
 		//S_IFMT uint32 = 0170000 // bit mask for the file type bit field
@@ -159,6 +162,9 @@ func (t *Tracee) processFileWrites() {
 				t.stats.LostWrCount.Increment(lost)
 				logger.Warnw(fmt.Sprintf("Lost %d write events", lost))
 			}
+
+		case <-ctx.Done():
+			return
 		}
 	}
 }
diff --git a/pkg/server/server.go b/pkg/server/server.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/http/pprof"
@@ -12,16 +13,21 @@ import (
 
 // Server represents a http server
 type Server struct {
-	mux            *http.ServeMux
-	listenAddr     string
+	hs             *http.Server
+	mux            *http.ServeMux // just an exposed copy of hs.Handler
 	metricsEnabled bool
 }
 
 // New creates a new server
 func New(listenAddr string) *Server {
+	mux := http.NewServeMux()
+
 	return &Server{
-		mux:        http.NewServeMux(),
-		listenAddr: listenAddr,
+		hs: &http.Server{
+			Addr:    listenAddr,
+			Handler: mux,
+		},
+		mux: mux,
 	}
 }
 
@@ -38,15 +44,35 @@ func (s *Server) EnableHealthzEndpoint() {
 	})
 }
 
-// Start starts the http server on the listen addr
-func (s *Server) Start() {
-	logger.Debugw("Serving metrics endpoint", "address", s.listenAddr)
+// Start starts the http server on the listen address
+func (s *Server) Start(ctx context.Context) {
+	srvCtx, srvCancel := context.WithCancel(ctx)
+	defer srvCancel()
+
+	go func() {
+		logger.Debugw("Starting serving metrics endpoint go routine", "address", s.hs.Addr)
+		defer logger.Debugw("Stopped serving metrics endpoint")
+
+		if err := s.hs.ListenAndServe(); err != http.ErrServerClosed {
+			logger.Errorw("Serving metrics endpoint", "error", err)
+		}
+
+		srvCancel()
+	}()
+
+	select {
+	case <-ctx.Done():
+		logger.Debugw("Context cancelled, shutting down metrics endpoint server")
+		if err := s.hs.Shutdown(ctx); err != nil {
+			logger.Errorw("Stopping serving metrics endpoint", "error", err)
+		}
 
-	if err := http.ListenAndServe(s.listenAddr, s.mux); err != http.ErrServerClosed {
-		logger.Errorw("Serving metrics endpoint", "error", err)
+	// if server error occurred while base ctx is not done, we should exit via this case
+	case <-srvCtx.Done():
 	}
 }
 
+// EnablePProfEndpoint enables pprof endpoint for debugging
 func (s *Server) EnablePProfEndpoint() {
 	s.mux.HandleFunc("/debug/pprof/", pprof.Index)
 	s.mux.Handle("/debug/pprof/allocs", pprof.Handler("allocs"))
@@ -59,6 +85,7 @@ func (s *Server) EnablePProfEndpoint() {
 	s.mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
 }
 
+// MetricsEndpointEnabled returns true if metrics endpoint is enabled
 func (s *Server) MetricsEndpointEnabled() bool {
 	return s.metricsEnabled
 }
diff --git a/tests/integration/integration_test.go b/tests/integration/integration_test.go
@@ -397,12 +397,21 @@ func Test_EventFilters(t *testing.T) {
 			eventOutput := &eventOutput{}
 
 			go func() {
-				for evt := range eventChan {
-					eventOutput.addEvent(evt)
+				for {
+					select {
+					case evt, ok := <-eventChan:
+						if !ok {
+							return
+						}
+						eventOutput.addEvent(evt)
+
+					case <-ctx.Done():
+						return
+					}
 				}
 			}()
 
-			trc := startTracee(t, config, nil, nil, ctx)
+			trc := startTracee(t, ctx, config, nil, nil)
 
 			waitforTraceeStart(t, trc, time.Now())
 

diff --git a/tests/integration/tracee.go b/tests/integration/tracee.go
@@ -17,7 +17,7 @@ import (
 )
 
 // load tracee into memory with args
-func startTracee(t *testing.T, config tracee.Config, output *tracee.OutputConfig, capture *tracee.CaptureConfig, ctx context.Context) *tracee.Tracee {
+func startTracee(t *testing.T, ctx context.Context, config tracee.Config, output *tracee.OutputConfig, capture *tracee.CaptureConfig) *tracee.Tracee {
 	initialize.SetLibbpfgoCallbacks()
 
 	kernelConfig, err := initialize.KernelConfig()
@@ -43,8 +43,16 @@ func startTracee(t *testing.T, config tracee.Config, output *tracee.OutputConfig
 	errChan := make(chan error)
 
 	go func() {
-		for err := range errChan {
-			t.Logf("received error while testing: %s\n", err)
+		for {
+			select {
+			case err, ok := <-errChan:
+				if !ok {
+					return
+				}
+				t.Logf("received error while testing: %s\n", err)
+			case <-ctx.Done():
+				return
+			}
 		}
 	}()
-Original file line number
+Diff line change
@@ Expand Up / @@ -37,7 +37,7 @@ func (r Runner) Run(ctx context.Context) error { @@
     			}
     		}
-    		go r.Server.Start()
+    		go r.Server.Start(ctx)
     	}
     	// Print statistics at the end
@@ Expand Down @@