Make vulnerability processing less RAM hungry (#2739)

* Make preview work when run from scratch (no orbit running)

* Make vulnerability processing less RAM hungry

* Add changes file

* Only get the cpe list once

* Remove cache

* Try killing osquery as well as orbit and clear their db dir

Co-authored-by: Lucas Rodriguez <lucas@fleetdm.com>

changes/make-vuln-processing-less-intense

@@ -0,0 +1 @@
+* Make vulnerability processing consume less RAM

cmd/fleetctl/preview.go

@@ -576,17 +576,17 @@ func storePidFile(destDir string, pid int) error {
 	return nil
 }
 
-func readPidFromFile(destDir string) (int, error) {
-	pidFilePath := path.Join(destDir, "orbit.pid")
+func readPidFromFile(destDir string, what string) (int, error) {
+	pidFilePath := path.Join(destDir, what)
 	data, err := os.ReadFile(pidFilePath)
 	if err != nil {
-		return -1, fmt.Errorf("error reading pidfile %s: %s", pidFilePath, err)
+		return -1, fmt.Errorf("error reading pidfile %s: %w", pidFilePath, err)
 	}
 	return strconv.Atoi(string(data))
 }
 
 func isOrbitAlreadyRunning(destDir string) bool {
-	pid, err := readPidFromFile(destDir)
+	pid, err := readPidFromFile(destDir, "orbit.pid")
 	if err != nil {
 		// if any error occurs reading the pid file, we assume orbit is not running
 		return false
@@ -610,6 +610,15 @@ func downloadOrbitAndStart(destDir string, enrollSecret string, address string)
 		fmt.Println("Orbit is already running.")
 		return nil
 	}
+
+	fmt.Println("Trying to clear orbit and osquery directories...")
+	if err := os.RemoveAll(path.Join(destDir, "osquery.db")); err != nil {
+		fmt.Println("Warning: clearing osquery db dir:", err)
+	}
+	if err := os.RemoveAll(path.Join(destDir, "orbit.db")); err != nil {
+		fmt.Println("Warning: clearing orbit db dir:", err)
+	}
+
 	updateOpt := update.DefaultOptions
 	switch runtime.GOOS {
 	case "linux":
@@ -633,6 +642,7 @@ func downloadOrbitAndStart(destDir string, enrollSecret string, address string)
 		"--root-dir", destDir,
 		"--fleet-url", address,
 		"--insecure",
+		"--debug",
 		"--enroll-secret", enrollSecret,
 		"--log-file", path.Join(destDir, "orbit.log"),
 	)
@@ -647,13 +657,33 @@ func downloadOrbitAndStart(destDir string, enrollSecret string, address string)
 }
 
 func stopOrbit(destDir string) error {
-	pid, err := readPidFromFile(destDir)
+	err := killFromPIDFile(destDir, "osquery.pid")
+	if err != nil {
+		return err
+	}
+	err = killFromPIDFile(destDir, "orbit.pid")
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func killFromPIDFile(destDir string, w string) error {
+	pid, err := readPidFromFile(destDir, w)
 	if err != nil {
 		return errors.Wrap(err, "reading pid")
 	}
+	switch {
+	case err == nil:
+		// OK
+	case errors.Is(err, os.ErrNotExist):
+		return nil // we assume it's not running
+	default:
+		return errors.Wrapf(err, "reading pid from: %s", destDir)
+	}
 	err = killPID(pid)
 	if err != nil {
-		return errors.Wrapf(err, "killing orbit %d", pid)
+		return errors.Wrapf(err, "killing %d", pid)
 	}
 	return nil
 }

orbit/cmd/orbit/orbit.go

@@ -115,7 +115,7 @@ func main() {
 
 		log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: time.RFC3339Nano, NoColor: true})
 		if logfile := c.String("log-file"); logfile != "" {
-			f, err := secure.OpenFile(logfile, os.O_CREATE|os.O_APPEND, 0o600)
+			f, err := secure.OpenFile(logfile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
 			if err != nil {
 				return errors.Wrap(err, "open logfile")
 			}

server/vulnerabilities/cve.go

@@ -62,6 +62,18 @@ func TranslateCPEToCVE(
 		return err
 	}
 
+	var files []string
+	err = filepath.Walk(vulnPath, func(path string, info os.FileInfo, err error) error {
+		if match, err := regexp.MatchString("nvdcve.*\\.gz$", path); !match || err != nil {
+			return nil
+		}
+		files = append(files, path)
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
 	cpeList, err := ds.AllCPEs(ctx)
 	if err != nil {
 		return err
@@ -80,24 +92,24 @@ func TranslateCPEToCVE(
 		return nil
 	}
 
-	var files []string
-	err = filepath.Walk(vulnPath, func(path string, info os.FileInfo, err error) error {
-		if match, err := regexp.MatchString("nvdcve.*\\.gz$", path); !match || err != nil {
-			return nil
+	for _, file := range files {
+		err := checkCVEs(ctx, ds, logger, cpes, file)
+		if err != nil {
+			return err
 		}
-		files = append(files, path)
-		return nil
-	})
-	if err != nil {
-		return err
 	}
 
+	return nil
+}
+
+func checkCVEs(ctx context.Context, ds fleet.Datastore, logger kitlog.Logger, cpes []*wfn.Attributes, files ...string) error {
 	dict, err := cvefeed.LoadJSONDictionary(files...)
 	if err != nil {
 		return err
 	}
-	cache := cvefeed.NewCache(dict).SetRequireVersion(true).SetMaxSize(0)
-	cache.Idx = cvefeed.NewIndex(dict)
+	cache := cvefeed.NewCache(dict).SetRequireVersion(true).SetMaxSize(-1)
+	// This index consumes too much RAM
+	//cache.Idx = cvefeed.NewIndex(dict)
 
 	cpeCh := make(chan *wfn.Attributes)
 

server/vulnerabilities/cve_test.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"runtime"
 	"strings"
 	"sync"
 	"testing"
@@ -25,6 +26,20 @@ var cvetests = []struct {
 	{"cpe:2.3:a:1password:1password:3.9.9:*:*:*:*:*:*:*", "CVE-2012-6369"},
 }
 
+func PrintMemUsage() {
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	// For info on each, see: https://golang.org/pkg/runtime/#MemStats
+	fmt.Printf("Alloc = %v MiB", bToMb(m.Alloc))
+	fmt.Printf("\tTotalAlloc = %v MiB", bToMb(m.TotalAlloc))
+	fmt.Printf("\tSys = %v MiB", bToMb(m.Sys))
+	fmt.Printf("\tNumGC = %v\n", m.NumGC)
+}
+
+func bToMb(b uint64) uint64 {
+	return b / 1024 / 1024
+}
+
 func TestTranslateCPEToCVE(t *testing.T) {
 	if os.Getenv("NETWORK_TEST") == "" {
 		t.Skip("set environment variable NETWORK_TEST=1 to run")
@@ -55,6 +70,8 @@ func TestTranslateCPEToCVE(t *testing.T) {
 			err := TranslateCPEToCVE(ctx, ds, tempDir, kitlog.NewLogfmtLogger(os.Stdout), config.FleetConfig{})
 			require.NoError(t, err)
 
+			PrintMemUsage()
+
 			require.Equal(t, []string{tt.cve}, cvesFound)
 			require.Equal(t, []string{tt.cpe}, cveToCPEs[tt.cve])
 		})