Implement rate limit for RSA operations. (Velocidex#3957)

When the server first starts up it needs to populate the RSA cache to mitigate high CPU load. This PR implements a rate limiter to ensure that RSA operations are made at a fixed limit. This delays populating the cache in such a way that the server is not immediately loaded at start up. Also bugfixes: * Export STACK tables works correctly now * GetHunt API removes the CompiledCollectorArgs field to avoid sending very large grpc responses. * Fixed crash in server_frontend_cert() VQL function when run on the client. --------- Co-authored-by: snyk-bot <snyk-bot@snyk.io>
bvhuntcontent · Dec 9, 2024 · ff467f9 · ff467f9
1 parent 3196f4e
commit ff467f9
Show file tree

Hide file tree

Showing 26 changed files with 343 additions and 289 deletions.
diff --git a/api/download.go b/api/download.go
@@ -399,6 +399,24 @@ func getRows(
 
 		return rs_reader.Rows(ctx), rs_reader.Close, log_path, err
 
+	} else if request.Type == "STACK" {
+		log_path = path_specs.NewUnsafeFilestorePath(
+			utils.FilterSlice(request.StackPath, "")...).
+			SetType(api.PATH_TYPE_FILESTORE_JSON)
+
+		options, err := tables.GetTableOptions(request)
+		if err != nil {
+			return nil, nil, nil, err
+		}
+
+		rs_reader, err := result_sets.NewResultSetReaderWithOptions(
+			ctx, config_obj, file_store_factory, log_path, options)
+		if err != nil {
+			return nil, nil, nil, err
+		}
+
+		return rs_reader.Rows(ctx), rs_reader.Close, log_path, err
+
 	} else {
 		log_path, err := tables.GetPathSpec(
 			ctx, config_obj, request, principal)
@@ -408,7 +426,9 @@ func getRows(
 
 		rs_reader, err := result_sets.NewResultSetReader(
 			file_store_factory, log_path)
-
+		if err != nil {
+			return nil, nil, nil, err
+		}
 		return rs_reader.Rows(ctx), rs_reader.Close, log_path, err
 	}
 }

diff --git a/api/hunts.go b/api/hunts.go
@@ -439,6 +439,10 @@ func (self *ApiServer) GetHunt(
 			fmt.Errorf("%w: %v", services.HuntNotFoundError, in.HuntId))
 	}
 
+	if !in.IncludeRequest && result.StartRequest != nil {
+		result.StartRequest.CompiledCollectorArgs = nil
+	}
+
 	return result, nil
 }
 

diff --git a/api/proto/hunts.pb.go b/api/proto/hunts.pb.go
diff --git a/api/proto/hunts.proto b/api/proto/hunts.proto
@@ -173,6 +173,7 @@ message ListHuntsResponse {
 
 message GetHuntRequest {
     string hunt_id = 1;
+    bool include_request = 2;
 }
 
 message GetHuntResultsRequest {

diff --git a/artifacts/definitions/Windows/Search/FileFinder.yaml b/artifacts/definitions/Windows/Search/FileFinder.yaml
@@ -34,16 +34,12 @@ precondition:
   SELECT * FROM info() where OS = 'windows'
 
 parameters:
-  - name: SearchFilesGlob
-    default: C:\Users\*
-    description: Use a glob to define the files that will be searched.
-
   - name: SearchFilesGlobTable
     type: csv
     default: |
       Glob
       C:/Users/SomeUser/*
-    description: Alternative specify multiple globs in a table
+    description: Specify multiple globs to search for.
 
   - name: Accessor
     default: auto
@@ -93,7 +89,7 @@ sources:
                Btime AS BTime,
                Ctime AS CTime, "" AS Keywords,
                IsDir, Data
-        FROM glob(globs=SearchFilesGlobTable.Glob + SearchFilesGlob,
+        FROM glob(globs=SearchFilesGlobTable.Glob,
                   accessor=Accessor)
 
       LET more_recent = SELECT * FROM if(

diff --git a/artifacts/testdata/server/testcases/file_finder.in.yaml b/artifacts/testdata/server/testcases/file_finder.in.yaml
@@ -44,34 +44,33 @@ Queries:
 # Just find a zip file.
   - SELECT basename(path=OSPath) AS File, Hash, Size, Upload, Keywords
     FROM Artifact.Windows.Search.FileFinder(
-        SearchFilesGlob=srcDir + "/artifacts/testdata/files/*.zip")
+        SearchFilesGlobTable="Glob\n" + srcDir + "/artifacts/testdata/files/*.zip")
     WHERE File =~ "test"
 
 # Calculate the hash.
   - SELECT basename(path=OSPath) AS File, Hash, Size, Upload, Keywords
     FROM Artifact.Windows.Search.FileFinder(
         Calculate_Hash="Y",
-        SearchFilesGlob=srcDir + "/artifacts/testdata/files/*.zip")
+        SearchFilesGlobTable="Glob\n" + srcDir + "/artifacts/testdata/files/*.zip")
     WHERE File =~ "test"
 
 # Upload it.
   - SELECT basename(path=OSPath) AS File, Hash, Size, Upload.md5, Keywords
     FROM Artifact.Windows.Search.FileFinder(
         Upload_File="Y",
-        SearchFilesGlob=srcDir + "/artifacts/testdata/files/*.zip")
+        SearchFilesGlobTable="Glob\n" + srcDir + "/artifacts/testdata/files/*.zip")
     WHERE File =~ "test"
 
 # Search for files that match a keyword
   - SELECT basename(path=OSPath) AS File, Keywords
     FROM Artifact.Windows.Search.FileFinder(
         YaraRule="wide nocase ascii:supercalifragilisticexpialidocious",
-        SearchFilesGlob=srcDir + "/artifacts/testdata/**/*.in.yaml")
+        SearchFilesGlobTable="Glob\n"+srcDir + "/artifacts/testdata/**/*.in.yaml")
 
 # Search globs provided in a csv table
   - |
     SELECT basename(path=OSPath) AS File, Hash, Size, Upload, Keywords
     FROM Artifact.Windows.Search.FileFinder(
-        SearchFilesGlob="",
         SearchFilesGlobTable="Glob\n"+ srcDir + "/artifacts/testdata/files/*.zip\n")
     WHERE File =~ "test"
 
@@ -85,13 +84,16 @@ Queries:
   - LET _ <= mock(plugin="glob", results=MockData)
   - SELECT OSPath, MTime
     FROM Artifact.Windows.Search.FileFinder(
+        SearchFilesGlobTable="Glob\nC:/Users/*",
         MoreRecentThan=timestamp(string="2020-04-29T13:39:00Z"))
 
   - SELECT OSPath, MTime
     FROM Artifact.Windows.Search.FileFinder(
+        SearchFilesGlobTable="Glob\nC:/Users/*",
         ModifiedBefore=timestamp(string="2020-04-29T13:39:00Z"))
 
   - SELECT OSPath, MTime
     FROM Artifact.Windows.Search.FileFinder(
+        SearchFilesGlobTable="Glob\nC:/Users/*",
         MoreRecentThan=timestamp(string="2020-04-29T13:39:00Z"),
         ModifiedBefore=timestamp(string="2020-04-29T13:41:00Z"))
diff --git a/artifacts/testdata/server/testcases/file_finder.out.yaml b/artifacts/testdata/server/testcases/file_finder.out.yaml
@@ -2,7 +2,7 @@ SELECT mock(plugin='info', results=[dict(OS='windows'), dict(OS='windows')] ) FR
  {
   "mock(plugin='info', results=[dict(OS='windows'), dict(OS='windows')])": null
  }
-]SELECT basename(path=OSPath) AS File, Hash, Size, Upload, Keywords FROM Artifact.Windows.Search.FileFinder( SearchFilesGlob=srcDir + "/artifacts/testdata/files/*.zip") WHERE File =~ "test"[
+]SELECT basename(path=OSPath) AS File, Hash, Size, Upload, Keywords FROM Artifact.Windows.Search.FileFinder( SearchFilesGlobTable="Glob\n" + srcDir + "/artifacts/testdata/files/*.zip") WHERE File =~ "test"[
  {
   "File": "ext4_tests.zip",
   "Hash": null,
@@ -24,7 +24,7 @@ SELECT mock(plugin='info', results=[dict(OS='windows'), dict(OS='windows')] ) FR
   "Upload": null,
   "Keywords": ""
  }
-]SELECT basename(path=OSPath) AS File, Hash, Size, Upload, Keywords FROM Artifact.Windows.Search.FileFinder( Calculate_Hash="Y", SearchFilesGlob=srcDir + "/artifacts/testdata/files/*.zip") WHERE File =~ "test"[
+]SELECT basename(path=OSPath) AS File, Hash, Size, Upload, Keywords FROM Artifact.Windows.Search.FileFinder( Calculate_Hash="Y", SearchFilesGlobTable="Glob\n" + srcDir + "/artifacts/testdata/files/*.zip") WHERE File =~ "test"[
  {
   "File": "ext4_tests.zip",
   "Hash": {
@@ -58,7 +58,7 @@ SELECT mock(plugin='info', results=[dict(OS='windows'), dict(OS='windows')] ) FR
   "Upload": null,
   "Keywords": ""
  }
-]SELECT basename(path=OSPath) AS File, Hash, Size, Upload.md5, Keywords FROM Artifact.Windows.Search.FileFinder( Upload_File="Y", SearchFilesGlob=srcDir + "/artifacts/testdata/files/*.zip") WHERE File =~ "test"[
+]SELECT basename(path=OSPath) AS File, Hash, Size, Upload.md5, Keywords FROM Artifact.Windows.Search.FileFinder( Upload_File="Y", SearchFilesGlobTable="Glob\n" + srcDir + "/artifacts/testdata/files/*.zip") WHERE File =~ "test"[
  {
   "File": "ext4_tests.zip",
   "Hash": null,
@@ -80,14 +80,13 @@ SELECT mock(plugin='info', results=[dict(OS='windows'), dict(OS='windows')] ) FR
   "Upload.md5": "0f72ff84405a6f2def7d2446d9d5e836",
   "Keywords": ""
  }
-]SELECT basename(path=OSPath) AS File, Keywords FROM Artifact.Windows.Search.FileFinder( YaraRule="wide nocase ascii:supercalifragilisticexpialidocious", SearchFilesGlob=srcDir + "/artifacts/testdata/**/*.in.yaml")[
+]SELECT basename(path=OSPath) AS File, Keywords FROM Artifact.Windows.Search.FileFinder( YaraRule="wide nocase ascii:supercalifragilisticexpialidocious", SearchFilesGlobTable="Glob\n"+srcDir + "/artifacts/testdata/**/*.in.yaml")[
  {
   "File": "file_finder.in.yaml",
   "Keywords": "supercalifragilisticexpialidocious"
  }
 ]SELECT basename(path=OSPath) AS File, Hash, Size, Upload, Keywords
 FROM Artifact.Windows.Search.FileFinder(
-    SearchFilesGlob="",
     SearchFilesGlobTable="Glob\n"+ srcDir + "/artifacts/testdata/files/*.zip\n")
 WHERE File =~ "test"
 [
@@ -112,7 +111,7 @@ WHERE File =~ "test"
   "Upload": null,
   "Keywords": ""
  }
-]LET MockData <= SELECT *, timestamp(epoch=Mtime) AS Mtime, timestamp(epoch=Atime) AS Atime, timestamp(epoch=Ctime) AS Ctime, timestamp(epoch=Btime) AS Btime FROM parse_json_array(data=GlobMock)[]LET _ <= mock(plugin="glob", results=MockData)[]SELECT OSPath, MTime FROM Artifact.Windows.Search.FileFinder( MoreRecentThan=timestamp(string="2020-04-29T13:39:00Z"))[
+]LET MockData <= SELECT *, timestamp(epoch=Mtime) AS Mtime, timestamp(epoch=Atime) AS Atime, timestamp(epoch=Ctime) AS Ctime, timestamp(epoch=Btime) AS Btime FROM parse_json_array(data=GlobMock)[]LET _ <= mock(plugin="glob", results=MockData)[]SELECT OSPath, MTime FROM Artifact.Windows.Search.FileFinder( SearchFilesGlobTable="Glob\nC:/Users/*", MoreRecentThan=timestamp(string="2020-04-29T13:39:00Z"))[
  {
   "OSPath": "Users/yolo/Desktop/file2.txt",
   "MTime": "2020-04-29T13:40:10Z"
@@ -121,12 +120,12 @@ WHERE File =~ "test"
   "OSPath": "Users/yolo/Desktop/file3.txt",
   "MTime": "2020-04-29T13:41:50Z"
  }
-]SELECT OSPath, MTime FROM Artifact.Windows.Search.FileFinder( ModifiedBefore=timestamp(string="2020-04-29T13:39:00Z"))[
+]SELECT OSPath, MTime FROM Artifact.Windows.Search.FileFinder( SearchFilesGlobTable="Glob\nC:/Users/*", ModifiedBefore=timestamp(string="2020-04-29T13:39:00Z"))[
  {
   "OSPath": "Users/yolo/Desktop/file.txt",
   "MTime": "2020-04-29T13:38:27Z"
  }
-]SELECT OSPath, MTime FROM Artifact.Windows.Search.FileFinder( MoreRecentThan=timestamp(string="2020-04-29T13:39:00Z"), ModifiedBefore=timestamp(string="2020-04-29T13:41:00Z"))[
+]SELECT OSPath, MTime FROM Artifact.Windows.Search.FileFinder( SearchFilesGlobTable="Glob\nC:/Users/*", MoreRecentThan=timestamp(string="2020-04-29T13:39:00Z"), ModifiedBefore=timestamp(string="2020-04-29T13:41:00Z"))[
  {
   "OSPath": "Users/yolo/Desktop/file2.txt",
   "MTime": "2020-04-29T13:40:10Z"

diff --git a/artifacts/testdata/server/testcases/kapefiles_remap.in.yaml b/artifacts/testdata/server/testcases/kapefiles_remap.in.yaml
@@ -13,7 +13,7 @@ Queries:
     FROM collect(artifacts=["Windows.Forensics.Usn", "Windows.Search.FileFinder"],
        args=dict(`Windows.Forensics.Usn`=dict(),
                  `Windows.Search.FileFinder`=dict(
-                      SearchFilesGlob="HKEY_USERS/*/*",
+                      SearchFilesGlobTable="Glob\nHKEY_USERS/*/*",
                       Accessor="registry")),
        output=tmp,
        remapping=GetRemapping(FlowId=FlowId, ClientId=ClientId))

diff --git a/artifacts/testdata/server/testcases/kapefiles_remap.out.yaml b/artifacts/testdata/server/testcases/kapefiles_remap.out.yaml
@@ -2,7 +2,7 @@ LET ClientId <= client_create(hostname="TestVirtualClient").client_id[]LET FlowI
 FROM collect(artifacts=["Windows.Forensics.Usn", "Windows.Search.FileFinder"],
    args=dict(`Windows.Forensics.Usn`=dict(),
              `Windows.Search.FileFinder`=dict(
-                  SearchFilesGlob="HKEY_USERS/*/*",
+                  SearchFilesGlobTable="Glob\nHKEY_USERS/*/*",
                   Accessor="registry")),
    output=tmp,
    remapping=GetRemapping(FlowId=FlowId, ClientId=ClientId))

diff --git a/crypto/client/client.go b/crypto/client/client.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/Velocidex/ttlcache/v2"
 	"github.com/go-errors/errors"
+	"golang.org/x/time/rate"
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
 	crypto_utils "www.velocidex.com/golang/velociraptor/crypto/utils"
 	"www.velocidex.com/golang/velociraptor/logging"
@@ -80,6 +81,12 @@ func NewClientCryptoManager(
 		lru_size = config_obj.Frontend.Resources.ExpectedClients
 	}
 
+	limit_rate := int64(100)
+	if config_obj.Frontend != nil &&
+		config_obj.Frontend.Resources.EnrollmentsPerSecond > 0 {
+		limit_rate = config_obj.Frontend.Resources.EnrollmentsPerSecond
+	}
+
 	result := &ClientCryptoManager{CryptoManager{
 		client_id:           client_id,
 		private_key:         private_key,
@@ -88,6 +95,7 @@ func NewClientCryptoManager(
 		unauthenticated_lru: ttlcache.NewCache(),
 		caPool:              roots,
 		logger:              logger,
+		limiter:             rate.NewLimiter(rate.Limit(limit_rate), 100),
 	}}
 
 	go func() {