[agent-smith] Improve cpu and memory utilization

Author: Furisto

components/ee/agent-smith/cmd/signature-matches.go

@@ -28,6 +28,10 @@ var signatureMatchesCmd = &cobra.Command{
 		}
 		defer f.Close()
 
+		sfc := classifier.SignatureReadCache{
+			Reader: f,
+		}
+
 		if cfgFile == "" {
 			log.Info("no config present - reading signature from STDIN")
 			var sig classifier.Signature
@@ -36,7 +40,7 @@ var signatureMatchesCmd = &cobra.Command{
 				log.Fatal(err)
 			}
 
-			match, err := sig.Matches(f)
+			match, err := sig.Matches(&sfc)
 			if err != nil {
 				log.Fatal(err)
 			}
@@ -60,7 +64,7 @@ var signatureMatchesCmd = &cobra.Command{
 		var res []*classifier.Signature
 		for _, bl := range cfg.Blocklists.Levels() {
 			for _, s := range bl.Signatures {
-				m, err := s.Matches(f)
+				m, err := s.Matches(&sfc)
 				if err != nil {
 					log.WithError(err).WithField("signature", s.Name).Warn("cannot match signature")
 					continue

components/ee/agent-smith/pkg/classifier/classifier.go

@@ -7,6 +7,7 @@ package classifier
 import (
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"regexp"
 	"strings"
@@ -196,8 +197,12 @@ func (sigcl *SignatureMatchClassifier) Matches(executable string, cmdline []stri
 	defer r.Close()
 
 	var serr error
+
+	src := SignatureReadCache{
+		Reader: r,
+	}
 	for _, sig := range sigcl.Signatures {
-		match, err := sig.Matches(r)
+		match, err := sig.Matches(&src)
 		if match {
 			sigcl.signatureHitTotal.Inc()
 			return &Classification{
@@ -217,6 +222,13 @@ func (sigcl *SignatureMatchClassifier) Matches(executable string, cmdline []stri
 	return sigNoMatch, nil
 }
 
+type SignatureReadCache struct {
+	Reader  io.ReaderAt
+	header  []byte
+	symbols []string
+	rodata  []byte
+}
+
 func (sigcl *SignatureMatchClassifier) Describe(d chan<- *prometheus.Desc) {
 	sigcl.processMissTotal.Describe(d)
 	sigcl.signatureHitTotal.Describe(d)

components/ee/agent-smith/pkg/classifier/signature_test.go

@@ -22,6 +22,10 @@ func TestMatchELF(t *testing.T) {
 	}
 	defer input.Close()
 
+	sfc := SignatureReadCache{
+		Reader: input,
+	}
+
 	sig := Signature{
 		Kind:    ObjectELFSymbols,
 		Pattern: []byte("bash_groupname_completion_function"),
@@ -32,7 +36,7 @@ func TestMatchELF(t *testing.T) {
 		return
 	}
 
-	matches, err := sig.Matches(input)
+	matches, err := sig.Matches(&sfc)
 	if err != nil {
 		t.Errorf("cannot match signature: %v", err)
 		return
@@ -60,7 +64,11 @@ func TestMatchAny(t *testing.T) {
 			return
 		}
 
-		matches, err := test.Signature.matchAny(bytes.NewReader(test.Input))
+		sfc := SignatureReadCache{
+			Reader: bytes.NewReader(test.Input),
+		}
+
+		matches, err := test.Signature.matchAny(&sfc)
 		if err != nil {
 			t.Errorf("[%03d] cannot match signature: %v", i, err)
 			return

components/ee/agent-smith/pkg/classifier/sinature.go

@@ -103,16 +103,16 @@ func (s *Signature) Validate() error {
 }
 
 // Matches checks if the signature applies to the stream
-func (s *Signature) Matches(in io.ReaderAt) (bool, error) {
+func (s *Signature) Matches(in *SignatureReadCache) (bool, error) {
 	if s.Slice.Start > 0 {
-		_, err := in.ReadAt([]byte{}, s.Slice.Start)
+		_, err := in.Reader.ReadAt([]byte{}, s.Slice.Start)
 		// slice start exceeds what we can read - this signature cannot match
 		if err != nil {
 			return false, nil
 		}
 	}
 	if s.Slice.End > 0 {
-		_, err := in.ReadAt([]byte{}, s.Slice.End)
+		_, err := in.Reader.ReadAt([]byte{}, s.Slice.End)
 		// slice start exceeds what we can read - this signature cannot match
 		if err != nil {
 			return false, nil
@@ -121,21 +121,27 @@ func (s *Signature) Matches(in io.ReaderAt) (bool, error) {
 
 	// check the object kind
 	if s.Kind != ObjectAny {
-		head := make([]byte, 261)
-		_, err := in.ReadAt(head, 0)
-		if err == io.EOF {
-			// cannot read header which means that only Any rules would apply
-			return false, nil
-		}
-		if err != nil {
-			return false, xerrors.Errorf("cannot read stream head: %w", err)
+		var head []byte
+		if len(in.header) > 0 {
+			head = in.header
+		} else {
+			head = make([]byte, 261)
+			_, err := in.Reader.ReadAt(head, 0)
+			if err == io.EOF {
+				// cannot read header which means that only Any rules would apply
+				return false, nil
+			}
+			if err != nil {
+				return false, xerrors.Errorf("cannot read stream head: %w", err)
+			}
+			in.header = head
 		}
 
 		matches := false
 		switch s.Kind {
 		case ObjectELFSymbols, ObjectELFRodata:
 			matches = isELF(head)
-		case ObjectAny:
+		default:
 			matches = true
 		}
 		if !matches {
@@ -172,16 +178,23 @@ func isELF(head []byte) bool {
 }
 
 // matchELF matches a signature against an ELF file
-func (s *Signature) matchELFRodata(in io.ReaderAt) (bool, error) {
-	executable, err := elf.NewFile(in)
-	if err != nil {
-		return false, xerrors.Errorf("cannot anaylse ELF file: %w", err)
-	}
+func (s *Signature) matchELFRodata(in *SignatureReadCache) (bool, error) {
+	var rodata []byte
+	if len(in.rodata) > 0 {
+		rodata = in.rodata
+	} else {
+		executable, err := elf.NewFile(in.Reader)
+		if err != nil {
+			return false, xerrors.Errorf("cannot anaylse ELF file: %w", err)
+		}
 
-	rodata, err := ExtractELFRodata(executable)
-	if err != nil {
-		return false, err
+		rodata, err = ExtractELFRodata(executable)
+		if err != nil {
+			return false, err
+		}
+		in.rodata = rodata
 	}
+
 	matches, err := s.matches(rodata)
 	if matches || err != nil {
 		return matches, err
@@ -191,16 +204,23 @@ func (s *Signature) matchELFRodata(in io.ReaderAt) (bool, error) {
 }
 
 // matchELF matches a signature against an ELF file
-func (s *Signature) matchELF(in io.ReaderAt) (bool, error) {
-	executable, err := elf.NewFile(in)
-	if err != nil {
-		return false, xerrors.Errorf("cannot anaylse ELF file: %w", err)
-	}
+func (s *Signature) matchELF(in *SignatureReadCache) (bool, error) {
+	var symbols []string
+	if len(in.symbols) > 0 {
+		symbols = in.symbols
+	} else {
+		executable, err := elf.NewFile(in.Reader)
+		if err != nil {
+			return false, xerrors.Errorf("cannot anaylse ELF file: %w", err)
+		}
 
-	symbols, err := ExtractELFSymbols(executable)
-	if err != nil {
-		return false, err
+		symbols, err = ExtractELFSymbols(executable)
+		if err != nil {
+			return false, err
+		}
+		in.symbols = symbols
 	}
+
 	for _, sym := range symbols {
 		matches, err := s.matches([]byte(sym))
 		if matches || err != nil {
@@ -213,22 +233,28 @@ func (s *Signature) matchELF(in io.ReaderAt) (bool, error) {
 
 // ExtractELFSymbols extracts all ELF symbol names from an ELF binary
 func ExtractELFSymbols(executable *elf.File) ([]string, error) {
-	var symbols []string
 	syms, err := executable.Symbols()
 	if err != nil && err != elf.ErrNoSymbols {
 		return nil, xerrors.Errorf("cannot get dynsym section: %w", err)
 	}
-	for _, s := range syms {
-		symbols = append(symbols, s.Name)
-	}
 
 	dynsyms, err := executable.DynamicSymbols()
 	if err != nil && err != elf.ErrNoSymbols {
 		return nil, xerrors.Errorf("cannot get dynsym section: %w", err)
 	}
+
+	symbols := make([]string, len(syms)+len(dynsyms))
+	i := 0
+	for _, s := range syms {
+		symbols[i] = s.Name
+		i += 1
+	}
+
 	for _, s := range dynsyms {
-		symbols = append(symbols, s.Name)
+		symbols[i] = s.Name
+		i += 1
 	}
+
 	return symbols, nil
 }
 
@@ -247,11 +273,11 @@ func ExtractELFRodata(executable *elf.File) ([]byte, error) {
 }
 
 // matchAny matches a signature against a binary file
-func (s *Signature) matchAny(in io.ReaderAt) (bool, error) {
+func (s *Signature) matchAny(in *SignatureReadCache) (bool, error) {
 	buffer := make([]byte, 8096)
 	pos := s.Slice.Start
 	for {
-		n, err := in.ReadAt(buffer, pos)
+		n, err := in.Reader.ReadAt(buffer, pos)
 		sub := buffer[0:n]
 		pos += int64(n)