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


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## 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.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

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.

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
+}

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
+}

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

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)

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
+}

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

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.
 }