Merge pull request #11 from BlackVectorOps/perf/optimize-indexing-loop-9245333493878548011

⚡ Optimize indexing loop by batching signature additions

Author: xkilldash9x

internal/cli/index.go

@@ -166,10 +166,13 @@ func RunIndexJSON(target string, results []diff.FingerprintResult, name, severit
 			Created: time.Now().Format("2006-01-02"),
 		}
 
-		scanner.AddSignature(&sig)
 		indexed = append(indexed, sig)
 	}
 
+	if err := scanner.AddSignatures(indexed); err != nil {
+		return nil, 0, err
+	}
+
 	if err := scanner.SaveDatabase(dbPath); err != nil {
 		return nil, 0, err
 	}

pkg/storage/jsondb/json_store.go

@@ -177,6 +177,18 @@ func (s *Scanner) SaveDatabase(path string) error {
 	return nil
 }
 
+// Generates a secure random ID if the signature doesn't have one.
+func (s *Scanner) ensureID(sig *detection.Signature) error {
+	if sig.ID == "" {
+		b := make([]byte, 8)
+		if _, err := rand.Read(b); err != nil {
+			return fmt.Errorf("failed to generate secure random ID: %w", err)
+		}
+		sig.ID = fmt.Sprintf("SFW-AUTO-%s", hex.EncodeToString(b))
+	}
+	return nil
+}
+
 // Adds a new signature to the database.
 // We use crypto/rand for ID generation because math/rand is deterministic
 // and we don't want ID collisions if the seed isn't set properly.
@@ -195,20 +207,44 @@ func (s *Scanner) AddSignature(sig *detection.Signature) error {
 		}
 	}
 
-	// Generate ID if not provided using secure entropy.
-	if sig.ID == "" {
-		b := make([]byte, 8)
-		if _, err := rand.Read(b); err != nil {
-			return fmt.Errorf("failed to generate secure random ID: %w", err)
-		}
-		sig.ID = fmt.Sprintf("SFW-AUTO-%s", hex.EncodeToString(b))
+	if err := s.ensureID(sig); err != nil {
+		return err
 	}
 
 	// Append copies the struct value.
 	s.db.Signatures = append(s.db.Signatures, *sig)
 	return nil
 }
 
+// Adds multiple signatures to the database in a batch.
+// This reduces locking overhead when indexing many files.
+func (s *Scanner) AddSignatures(sigs []detection.Signature) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if len(sigs) == 0 {
+		return nil
+	}
+
+	if s.db == nil {
+		s.db = &detection.SignatureDatabase{
+			Version:     "1.0",
+			Description: "Semantic Firewall Malware Signature Database",
+		}
+	}
+
+	for i := range sigs {
+		sig := &sigs[i]
+
+		if err := s.ensureID(sig); err != nil {
+			return err
+		}
+
+		s.db.Signatures = append(s.db.Signatures, *sig)
+	}
+	return nil
+}
+
 // Retrieves a signature by ID.
 // Returns a deep copy to prevent the caller from modifying the internal
 // database state without a lock. Shared mutable state is the root of all evil.

pkg/storage/jsondb/json_store_test.go

@@ -0,0 +1,69 @@
+package jsondb
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/BlackVectorOps/semantic_firewall/v3/pkg/detection"
+)
+
+func TestAddSignatures(t *testing.T) {
+	s := NewScanner()
+	sigs := []detection.Signature{
+		{Name: "Sig1"},
+		{Name: "Sig2"},
+	}
+	if err := s.AddSignatures(sigs); err != nil {
+		t.Fatalf("AddSignatures failed: %v", err)
+	}
+
+	if len(s.db.Signatures) != 2 {
+		t.Errorf("Expected 2 signatures, got %d", len(s.db.Signatures))
+	}
+	if s.db.Signatures[0].Name != "Sig1" {
+		t.Errorf("Expected Sig1, got %s", s.db.Signatures[0].Name)
+	}
+    // Verify IDs were generated
+    if s.db.Signatures[0].ID == "" {
+        t.Error("ID should have been generated for Sig1")
+    }
+}
+
+const count = 10000
+
+func BenchmarkAddSignatureLoop(b *testing.B) {
+	sigs := make([]detection.Signature, count)
+	for i := 0; i < count; i++ {
+		sigs[i] = detection.Signature{
+			ID:   fmt.Sprintf("SIG-%d", i),
+			Name: fmt.Sprintf("Signature %d", i),
+		}
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		s := NewScanner()
+		s.db.Signatures = make([]detection.Signature, 0, count)
+
+		for j := 0; j < count; j++ {
+			_ = s.AddSignature(&sigs[j])
+		}
+	}
+}
+
+func BenchmarkAddSignaturesBatch(b *testing.B) {
+	sigs := make([]detection.Signature, count)
+	for i := 0; i < count; i++ {
+		sigs[i] = detection.Signature{
+			ID:   fmt.Sprintf("SIG-%d", i),
+			Name: fmt.Sprintf("Signature %d", i),
+		}
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		s := NewScanner()
+		s.db.Signatures = make([]detection.Signature, 0, count)
+		_ = s.AddSignatures(sigs)
+	}
+}