From 3df12bf32b995ab4ef2c34ffe95d547f474d7573 Mon Sep 17 00:00:00 2001 From: Victor Lyuboslavsky <2685025+getvictor@users.noreply.github.com> Date: Mon, 8 Sep 2025 10:24:22 -0500 Subject: [PATCH] Improved performance when modifying config with a large number of yara rules (#32696) Fixes #29909 - Do not update DB if rules haven't changed - Cache Yara rules when retrieved by hosts. This should reduce DB accesses with large number of hosts retrieving large numbers of rules I manually QA'd using OpenTelemetry (APM would also work) and monitoring the DB accesses when updating or retrieving yara rules. # Checklist for submitter - [x] Changes file added for user-visible changes in `changes/`, `orbit/changes/` or `ee/fleetd-chrome/changes`. See [Changes files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/guides/committing-changes.md#changes-files) for more information. ## Testing - [x] Added/updated automated tests - [x] QA'd all new/changed functionality manually ## Summary by CodeRabbit * **Performance Improvements** * Faster config saves when many YARA rules are present (incremental updates, reduced work). * Lower latency and load when many hosts fetch YARA rules (caching and smarter retrieval). * More efficient handling of unchanged, added, modified, and removed YARA rules. * **Documentation** * Changelog entry noting YARA rules performance and fetch improvements. --- changes/29909-yara-rules-performance | 1 + server/datastore/cached_mysql/cached_mysql.go | 49 +++++++ .../cached_mysql/cached_mysql_test.go | 133 ++++++++++++++++++ server/datastore/mysql/app_configs.go | 74 ++++++++-- server/datastore/mysql/app_configs_test.go | 94 +++++++++++++ server/fleet/app.go | 7 + .../cloner-check/generated_files/yararule.txt | 2 + tools/cloner-check/main.go | 1 + 8 files changed, 353 insertions(+), 8 deletions(-) create mode 100644 changes/29909-yara-rules-performance create mode 100644 tools/cloner-check/generated_files/yararule.txt diff --git a/changes/29909-yara-rules-performance b/changes/29909-yara-rules-performance new file mode 100644 index 0000000000..8f0d02b736 --- /dev/null +++ b/changes/29909-yara-rules-performance @@ -0,0 +1 @@ +* Improved performance for YARA rules: when modifying config (PATCH /api/latest/fleet/config) with a large number of yara rules and when large numbers of hosts fetch rules via /api/osquery/yara/{name} endpoint. diff --git a/server/datastore/cached_mysql/cached_mysql.go b/server/datastore/cached_mysql/cached_mysql.go index cb2ed3d2bd..6b84e7e212 100644 --- a/server/datastore/cached_mysql/cached_mysql.go +++ b/server/datastore/cached_mysql/cached_mysql.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "strings" "time" "github.com/fleetdm/fleet/v4/server/contexts/ctxdb" @@ -53,6 +54,9 @@ const ( defaultQueryByNameExpiration = 1 * time.Second queryResultsCountKey = "QueryResultsCount:%d" defaultQueryResultsCountExpiration = 1 * time.Second + yaraRuleCachePrefix = "YaraRuleByName:" + yaraRuleByNameKey = yaraRuleCachePrefix + "%s" + defaultYaraRuleByNameExpiration = 1 * time.Minute // NOTE: MDM assets are cached using their checksum as well, as it's // important for them to always be fresh if they changed (see cachedi // mplementation below for details) @@ -119,6 +123,7 @@ type cachedMysql struct { defaultTeamConfigExp time.Duration queryByNameExp time.Duration queryResultsCountExp time.Duration + yaraRuleByNameExp time.Duration mdmConfigAssetExp time.Duration } @@ -172,6 +177,12 @@ func WithQueryResultsCountExpiration(d time.Duration) Option { } } +func WithYaraRuleByNameExpiration(d time.Duration) Option { + return func(o *cachedMysql) { + o.yaraRuleByNameExp = d + } +} + func WithMDMConfigAssetExpiration(d time.Duration) Option { return func(o *cachedMysql) { o.mdmConfigAssetExp = d @@ -197,6 +208,7 @@ func New(ds fleet.Datastore, opts ...Option) fleet.Datastore { defaultTeamConfigExp: defaultDefaultTeamConfigExpiration, queryByNameExp: defaultQueryByNameExpiration, queryResultsCountExp: defaultQueryResultsCountExpiration, + yaraRuleByNameExp: defaultYaraRuleByNameExpiration, mdmConfigAssetExp: defaultMDMConfigAssetExpiration, } for _, fn := range opts { @@ -487,3 +499,40 @@ func (ds *cachedMysql) SaveDefaultTeamConfig(ctx context.Context, config *fleet. return nil } + +func (ds *cachedMysql) YaraRuleByName(ctx context.Context, name string) (*fleet.YaraRule, error) { + key := fmt.Sprintf(yaraRuleByNameKey, name) + + if x, found := ds.c.Get(ctx, key); found { + if rule, ok := x.(*fleet.YaraRule); ok { + return rule, nil + } + } + + rule, err := ds.Datastore.YaraRuleByName(ctx, name) + if err != nil { + return nil, err + } + + ds.c.Set(ctx, key, rule, ds.yaraRuleByNameExp) + + return rule, nil +} + +func (ds *cachedMysql) ApplyYaraRules(ctx context.Context, rules []fleet.YaraRule) error { + err := ds.Datastore.ApplyYaraRules(ctx, rules) + if err != nil { + return err + } + + // Invalidate all cached YARA rules + // We need to flush all because we don't know which rules were added/removed/modified + items := ds.c.Items() + for k := range items { + if strings.HasPrefix(k, yaraRuleCachePrefix) { + ds.c.Delete(k) + } + } + + return nil +} diff --git a/server/datastore/cached_mysql/cached_mysql_test.go b/server/datastore/cached_mysql/cached_mysql_test.go index 81c5f23fe0..5ab88b4a28 100644 --- a/server/datastore/cached_mysql/cached_mysql_test.go +++ b/server/datastore/cached_mysql/cached_mysql_test.go @@ -60,6 +60,17 @@ func TestClone(t *testing.T) { }, }, }, + { + name: "yara rule", + src: &fleet.YaraRule{ + Name: "test_rule.yar", + Contents: "rule TestRule { condition: true }", + }, + want: &fleet.YaraRule{ + Name: "test_rule.yar", + Contents: "rule TestRule { condition: true }", + }, + }, } for _, tc := range tests { @@ -848,3 +859,125 @@ func TestGetAllMDMConfigAssetsByName(t *testing.T) { require.Error(t, err) require.Equal(t, "error fetching hashes", err.Error()) } + +func TestCachedYaraRules(t *testing.T) { + t.Parallel() + + mockedDS := new(mock.Store) + ds := New(mockedDS, WithYaraRuleByNameExpiration(100*time.Millisecond)) + + testRule1 := &fleet.YaraRule{ + Name: "rule1.yar", + Contents: "rule Rule1 { condition: true }", + } + testRule2 := &fleet.YaraRule{ + Name: "rule2.yar", + Contents: "rule Rule2 { condition: true }", + } + + // Setup mock functions + mockedDS.YaraRuleByNameFunc = func(_ context.Context, name string) (*fleet.YaraRule, error) { + switch name { + case "rule1.yar": + return testRule1, nil + case "rule2.yar": + return testRule2, nil + default: + return nil, errors.New("rule not found") + } + } + + mockedDS.ApplyYaraRulesFunc = func(_ context.Context, _ []fleet.YaraRule) error { + return nil + } + + // Test 1: First call gets the result from the DB + rule1, err := ds.YaraRuleByName(context.Background(), "rule1.yar") + require.NoError(t, err) + require.Equal(t, testRule1, rule1) + require.Same(t, testRule1, rule1) + require.True(t, mockedDS.YaraRuleByNameFuncInvoked) + mockedDS.YaraRuleByNameFuncInvoked = false + + // Test 2: Cached call returns cloned value + rule1Cached, err := ds.YaraRuleByName(context.Background(), "rule1.yar") + require.NoError(t, err) + require.Equal(t, testRule1, rule1Cached) // returns the cached value + require.NotSame(t, testRule1, rule1Cached) // have been cloned + require.False(t, mockedDS.YaraRuleByNameFuncInvoked) // from cache + + // Test 3: Deep change doesn't alter the stored value + rule1Cached.Contents = "modified content" + require.NotEqual(t, rule1, rule1Cached) + + // Verify original cached value is unchanged + rule1Again, err := ds.YaraRuleByName(context.Background(), "rule1.yar") + require.NoError(t, err) + require.Equal(t, testRule1, rule1Again) + require.False(t, mockedDS.YaraRuleByNameFuncInvoked) // still from cache + + // Test 4: Cache multiple rules + rule2, err := ds.YaraRuleByName(context.Background(), "rule2.yar") + require.NoError(t, err) + require.Equal(t, testRule2, rule2) + require.True(t, mockedDS.YaraRuleByNameFuncInvoked) + mockedDS.YaraRuleByNameFuncInvoked = false + + // Both rules should now be cached + rule2Cached, err := ds.YaraRuleByName(context.Background(), "rule2.yar") + require.NoError(t, err) + require.Equal(t, testRule2, rule2Cached) + require.False(t, mockedDS.YaraRuleByNameFuncInvoked) // from cache + + // Test 5: ApplyYaraRules invalidates all cached rules + newRules := []fleet.YaraRule{ + {Name: "rule1.yar", Contents: "rule Rule1Modified { condition: false }"}, + {Name: "rule3.yar", Contents: "rule Rule3 { condition: true }"}, + } + err = ds.ApplyYaraRules(context.Background(), newRules) + require.NoError(t, err) + require.True(t, mockedDS.ApplyYaraRulesFuncInvoked) + + // Update mock to return modified rule + testRule1 = &fleet.YaraRule{ + Name: "rule1.yar", + Contents: "rule Rule1Modified { condition: false }", + } + + // After ApplyYaraRules, cache should be invalidated, so next call hits the DB + rule1AfterApply, err := ds.YaraRuleByName(context.Background(), "rule1.yar") + require.NoError(t, err) + require.Equal(t, testRule1, rule1AfterApply) // updated rule + require.True(t, mockedDS.YaraRuleByNameFuncInvoked) // from DB, not cache + mockedDS.YaraRuleByNameFuncInvoked = false + + // rule2 should also be invalidated even though it wasn't changed + rule2AfterApply, err := ds.YaraRuleByName(context.Background(), "rule2.yar") + require.NoError(t, err) + require.Equal(t, testRule2, rule2AfterApply) + require.True(t, mockedDS.YaraRuleByNameFuncInvoked) // from DB, not cache + mockedDS.YaraRuleByNameFuncInvoked = false + + // Test 6: Cache expiration + // The previous call (rule1AfterApply) already cached rule1, so verify it's cached + rule1Cached2, err := ds.YaraRuleByName(context.Background(), "rule1.yar") + require.NoError(t, err) + require.Equal(t, testRule1, rule1Cached2) + require.False(t, mockedDS.YaraRuleByNameFuncInvoked) // should be from cache + + // Wait for cache to expire + time.Sleep(200 * time.Millisecond) + + // Update mock to return a different value + testRule1 = &fleet.YaraRule{ + Name: "rule1.yar", + Contents: "rule Rule1Expired { condition: true }", + } + + // This call should get from DB again since cache expired + rule1Expired, err := ds.YaraRuleByName(context.Background(), "rule1.yar") + require.NoError(t, err) + require.Equal(t, testRule1, rule1Expired) // new value from DB + require.Same(t, testRule1, rule1Expired) + require.True(t, mockedDS.YaraRuleByNameFuncInvoked) // from DB after expiration +} diff --git a/server/datastore/mysql/app_configs.go b/server/datastore/mysql/app_configs.go index 6f2e9d8eda..52337159ed 100644 --- a/server/datastore/mysql/app_configs.go +++ b/server/datastore/mysql/app_configs.go @@ -310,16 +310,69 @@ func (ds *Datastore) ApplyYaraRules(ctx context.Context, rules []fleet.YaraRule) } func applyYaraRulesDB(ctx context.Context, q sqlx.ExtContext, rules []fleet.YaraRule) error { - const delStmt = "DELETE FROM yara_rules" - if _, err := q.ExecContext(ctx, delStmt); err != nil { - return ctxerr.Wrap(ctx, err, "clear before insert") + // First, load existing rules to check if there are any changes + existingRules, err := getYaraRulesDB(ctx, q) + if err != nil { + return ctxerr.Wrap(ctx, err, "get existing yara rules") } - if len(rules) > 0 { + // Create maps for efficient comparison + existingMap := make(map[string]string, len(existingRules)) + for _, rule := range existingRules { + existingMap[rule.Name] = rule.Contents + } + + newMap := make(map[string]string, len(rules)) + for _, rule := range rules { + if _, exists := newMap[rule.Name]; exists { + return ctxerr.Wrap(ctx, &fleet.BadRequestError{Message: fmt.Sprintf("duplicate YARA rule name: %s", rule.Name)}, "duplicate rule name") + } + newMap[rule.Name] = rule.Contents + } + + // Determine which rules to delete (removed or need updating) + var toDelete []string + + // Rules that exist in DB but not in new rules (removed) + for name := range existingMap { + if _, exists := newMap[name]; !exists { + toDelete = append(toDelete, name) + } + } + + // Determine which rules to insert (new or updated) + var toInsert []fleet.YaraRule + for _, rule := range rules { + existingContent, exists := existingMap[rule.Name] + if !exists || existingContent != rule.Contents { + // Rule is new or has been modified + toInsert = append(toInsert, rule) + + // If it exists but content changed, we need to delete it first + if exists { + toDelete = append(toDelete, rule.Name) + } + } + } + + // Single DELETE for both removed rules and rules that need updating + if len(toDelete) > 0 { + stmt := fmt.Sprintf("DELETE FROM yara_rules WHERE name IN (%s)", strings.TrimSuffix(strings.Repeat("?,", len(toDelete)), ",")) + args := make([]any, len(toDelete)) + for i, name := range toDelete { + args[i] = name + } + if _, err := q.ExecContext(ctx, stmt, args...); err != nil { + return ctxerr.Wrap(ctx, err, "delete yara rules") + } + } + + // Insert new and updated rules + if len(toInsert) > 0 { const insStmt = `INSERT INTO yara_rules (name, contents) VALUES %s` - var args []interface{} - sql := fmt.Sprintf(insStmt, strings.TrimSuffix(strings.Repeat(`(?, ?),`, len(rules)), ",")) - for _, r := range rules { + args := make([]any, 0, len(toInsert)*2) + sql := fmt.Sprintf(insStmt, strings.TrimSuffix(strings.Repeat(`(?, ?),`, len(toInsert)), ",")) + for _, r := range toInsert { args = append(args, r.Name, r.Contents) } @@ -332,9 +385,14 @@ func applyYaraRulesDB(ctx context.Context, q sqlx.ExtContext, rules []fleet.Yara } func (ds *Datastore) GetYaraRules(ctx context.Context) ([]fleet.YaraRule, error) { + return getYaraRulesDB(ctx, ds.reader(ctx)) +} + +// getYaraRulesDB is a helper to get YARA rules using a specific database connection/transaction +func getYaraRulesDB(ctx context.Context, q sqlx.QueryerContext) ([]fleet.YaraRule, error) { sql := "SELECT name, contents FROM yara_rules" rules := []fleet.YaraRule{} - if err := sqlx.SelectContext(ctx, ds.reader(ctx), &rules, sql); err != nil { + if err := sqlx.SelectContext(ctx, q, &rules, sql); err != nil { return nil, ctxerr.Wrap(ctx, err, "get yara rules") } return rules, nil diff --git a/server/datastore/mysql/app_configs_test.go b/server/datastore/mysql/app_configs_test.go index edf6c80220..2bc198e5b7 100644 --- a/server/datastore/mysql/app_configs_test.go +++ b/server/datastore/mysql/app_configs_test.go @@ -627,6 +627,100 @@ func testYaraRulesRoundtrip(t *testing.T, ds *Datastore) { require.NoError(t, err) assert.Equal(t, &expectedRules[1], rule) + // Apply the same rules again - this should be a no-op due to optimization + // (rules haven't changed, so no DELETE/INSERT should occur) + err = ds.ApplyYaraRules(ctx, expectedRules) + require.NoError(t, err) + rules, err = ds.GetYaraRules(ctx) + require.NoError(t, err) + assert.Equal(t, expectedRules, rules) + + // Test edge case: Apply same rules but in different order + reorderedRules := []fleet.YaraRule{expectedRules[1], expectedRules[0]} + err = ds.ApplyYaraRules(ctx, reorderedRules) + require.NoError(t, err) + // Verify both rules still exist + rule, err = ds.YaraRuleByName(ctx, expectedRules[0].Name) + require.NoError(t, err) + assert.Equal(t, &expectedRules[0], rule) + rule, err = ds.YaraRuleByName(ctx, expectedRules[1].Name) + require.NoError(t, err) + assert.Equal(t, &expectedRules[1], rule) + + // Test: Modify only the content of one rule (same name, different content) + modifiedContentRules := []fleet.YaraRule{ + { + Name: "wildcard.yar", + Contents: `rule WildcardExampleModified +{ + strings: + $hex_string = { E2 34 ?? C8 A? FB FF } + + condition: + $hex_string +}`, + }, + { + Name: "jump-modified.yar", + Contents: `rule JumpExample +{ + strings: + $hex_string = true + + condition: + $hex_string +}`, + }, + } + err = ds.ApplyYaraRules(ctx, modifiedContentRules) + require.NoError(t, err) + // Verify the content was actually updated + rule, err = ds.YaraRuleByName(ctx, "wildcard.yar") + require.NoError(t, err) + assert.Contains(t, rule.Contents, "WildcardExampleModified") + assert.Contains(t, rule.Contents, "E2 34 ?? C8 A? FB FF") + + // Test: Mixed operations - add new, keep one unchanged, delete one + mixedRules := []fleet.YaraRule{ + { + Name: "wildcard.yar", + Contents: `rule WildcardExampleModified +{ + strings: + $hex_string = { E2 34 ?? C8 A? FB FF } + + condition: + $hex_string +}`, + }, // unchanged from previous + // jump-modified.yar is deleted + { + Name: "new-rule.yar", + Contents: `rule NewRule { condition: true }`, + }, // new rule + } + err = ds.ApplyYaraRules(ctx, mixedRules) + require.NoError(t, err) + + // Verify mixed operation results + rules, err = ds.GetYaraRules(ctx) + require.NoError(t, err) + require.Len(t, rules, 2) + + // Check wildcard.yar is unchanged + rule, err = ds.YaraRuleByName(ctx, "wildcard.yar") + require.NoError(t, err) + assert.Contains(t, rule.Contents, "WildcardExampleModified") + + // Check jump-modified.yar is deleted + _, err = ds.YaraRuleByName(ctx, "jump-modified.yar") + require.Error(t, err) + + // Check new-rule.yar is added + rule, err = ds.YaraRuleByName(ctx, "new-rule.yar") + require.NoError(t, err) + assert.Equal(t, `rule NewRule { condition: true }`, rule.Contents) + // Clear rules expectedRules = []fleet.YaraRule{} err = ds.ApplyYaraRules(ctx, expectedRules) diff --git a/server/fleet/app.go b/server/fleet/app.go index ebd9d33055..2f015c3207 100644 --- a/server/fleet/app.go +++ b/server/fleet/app.go @@ -1540,3 +1540,10 @@ type YaraRule struct { Name string `json:"name"` Contents string `json:"contents"` } + +func (r *YaraRule) Clone() (Cloner, error) { + return &YaraRule{ + Name: r.Name, + Contents: r.Contents, + }, nil +} diff --git a/tools/cloner-check/generated_files/yararule.txt b/tools/cloner-check/generated_files/yararule.txt new file mode 100644 index 0000000000..7c5f78f3ab --- /dev/null +++ b/tools/cloner-check/generated_files/yararule.txt @@ -0,0 +1,2 @@ +github.com/fleetdm/fleet/v4/server/fleet/YaraRule Name string +github.com/fleetdm/fleet/v4/server/fleet/YaraRule Contents string diff --git a/tools/cloner-check/main.go b/tools/cloner-check/main.go index 83eb3efcb6..d474e44c9d 100644 --- a/tools/cloner-check/main.go +++ b/tools/cloner-check/main.go @@ -50,6 +50,7 @@ var cacheableItems = []fleet.Cloner{ &fleet.Query{}, &fleet.MDMProfileSpec{}, &fleet.MDMConfigAsset{}, + &fleet.YaraRule{}, // TeamAgentOptions is not in the list because it is a json.RawMessage, no fields can change. // Same for ResultCountForQuery, it's just an int. }