reduce number of protobuf serializations per program

Author: hl4x

Sources/Fuzzilli/Database/DatabaseUtils.swift

@@ -4,9 +4,19 @@ import Crypto
 
 public class DatabaseUtils {
 
-    public static func encodeProgramToBase64(program: Program) throws -> String {
-        let data = try encodeProgramToProtobuf(program: program)
-        return data.base64EncodedString()
+    public static func prepareProgram(program: Program) throws -> (hash: String, base64: String) {
+        if program.code.contains(where: { $0.op is JsInternalOperation }) {
+            throw DatabaseUtilsError.programContainsInternalOperations
+        }
+
+        var proto = program.asProtobuf()
+        proto.clearParent()
+        let data = try proto.serializedData()
+
+        let hash = SHA256.hash(data: data).map { String(format: "%02x", $0) }.joined()
+        let base64 = data.base64EncodedString()
+
+        return (hash, base64)
     }
 
     public static func decodeProgramFromBase64(base64: String) throws -> Program {
@@ -16,6 +26,7 @@ public class DatabaseUtils {
         return try decodeProgramFromProtobuf(data: data)
     }
 
+    /*
     public static func encodeProgramToProtobuf(program: Program) throws -> Data {
         // Make sure the program does not contain internal operations
         if program.code.contains(where: { $0.op is JsInternalOperation }) {
@@ -29,12 +40,14 @@ public class DatabaseUtils {
 
         return try proto.serializedData()
     }
+    */
 
     public static func decodeProgramFromProtobuf(data: Data) throws -> Program {
         let proto = try Fuzzilli_Protobuf_Program(serializedBytes: data)
         return try Program(from: proto)
     }
 
+    /*
     public static func calculateProgramHash(program: Program) throws -> String {
         // Make sure the program does not contain internal operations
         if program.code.contains(where: { $0.op is JsInternalOperation }) {
@@ -57,6 +70,7 @@ public class DatabaseUtils {
         // Convert to hex string (64 characters for SHA-256)
         return digest.map { String(format: "%02x", $0) }.joined()
     }
+    */
 
     public static func mapExecutionOutcome(outcome: ExecutionOutcome) -> Int {
         switch outcome {

Sources/Fuzzilli/Database/PostgresSQLStorage.swift

@@ -67,20 +67,29 @@ public actor PostgresSQLStorage {
 
     public struct ProgramInput {
         public let program: Program
+        public let programHash: String
+        public let programBase64: String
         public let fuzzerId: Int
         public let mutatorNames: [String]
         public let contributorNames: [String]
+        public let parentHash: String?
 
         public init(
             program: Program, 
+            programHash: String, 
+            programBase64: String, 
             fuzzerId: Int, 
             mutatorNames: [String], 
-            contributorNames: [String]
+            contributorNames: [String],
+            parentHash: String?
         ) {
             self.program = program
+            self.programHash = programHash
+            self.programBase64 = programBase64
             self.fuzzerId = fuzzerId
             self.mutatorNames = mutatorNames
             self.contributorNames = contributorNames
+            self.parentHash = parentHash
         }
     }
 
@@ -369,6 +378,7 @@ public actor PostgresSQLStorage {
     }
 
     public func addProgramToBatch(_ programInput: ProgramInput) {
+        /*
         guard let programHash = try? DatabaseUtils.calculateProgramHash(program: programInput.program) else {
             if self.enableLogging {
                 self.logger.warning("Failed to calculate program hash, skipping program")
@@ -385,6 +395,7 @@ public actor PostgresSQLStorage {
         }
 
         seenProgramHashes.insert(programHash)
+        */
         pendingPrograms.append(programInput)
     }
 
@@ -421,6 +432,7 @@ public actor PostgresSQLStorage {
 
         if !programsToStore.isEmpty {
             // Group by fuzzerId to use storeProgramsBatch
+            // TODO Aleksi: Is this even necessary?
             let groupedPrograms = Dictionary(grouping: programsToStore, by: { $0.fuzzerId})
             for (fuzzerId, programInputs) in groupedPrograms {
                 _ = try await storeProgramsBatch(programInputs: programInputs, fuzzerId: fuzzerId)
@@ -441,66 +453,8 @@ public actor PostgresSQLStorage {
     }
 
     private func _storeProgramsBatchImpl(programInputs: [ProgramInput], fuzzerId: Int) async throws -> [String] {
-        // Pre-calculate all hashes and sort to ensure consistent lock ordering
-        struct PreparedProgram {
-            let hash: String
-            let input: ProgramInput
-            let parentHash: String?
-            let programData: String
-        }
-        
-        var preparedPrograms: [PreparedProgram] = []
-        
-        for programInput in programInputs {
-            let program = programInput.program
-            
-            // Calculate program hash
-            let programHash: String
-            do {
-                programHash = try DatabaseUtils.calculateProgramHash(program: program)
-            } catch {
-                if self.enableLogging {
-                    self.logger.warning("Failed to calculate hash for program, skipping: \(error)")
-                }
-                continue
-            }
-            
-            // Calculate parent hash if exists
-            let parentHash: String?
-            if let parentProgram = program.parent {
-                do {
-                    parentHash = try DatabaseUtils.calculateProgramHash(program: parentProgram)
-                } catch {
-                    if self.enableLogging {
-                        self.logger.warning("Failed to calculate parent hash, using nil: \(error)")
-                    }
-                    parentHash = nil
-                }
-            } else {
-                parentHash = nil
-            }
-            
-            // Encode program (must be done after parent hash calculation)
-            let programData: String
-            do {
-                programData = try DatabaseUtils.encodeProgramToBase64(program: program)
-            } catch {
-                if self.enableLogging {
-                    self.logger.warning("Failed to encode program with hash \(programHash), skipping: \(error)")
-                }
-                continue
-            }
-            
-            preparedPrograms.append(PreparedProgram(
-                hash: programHash,
-                input: programInput,
-                parentHash: parentHash,
-                programData: programData
-            ))
-        }
-        
         // Sort by hash - ensures all workers acquire locks in the same order
-        preparedPrograms.sort { $0.hash < $1.hash }
+        let sortedProgramInputs = programInputs.sorted { $0.programHash < $1.programHash }
 
         return try await databasePool.withConnection { connection in
             var programHashes: [String] = []
@@ -511,11 +465,10 @@ public actor PostgresSQLStorage {
                 var insertedCount = 0
                 var skippedCount = 0
 
-                for prepared in preparedPrograms {
-                    let programHash = prepared.hash
-                    let programInput = prepared.input
-                    let mutatorNames = programInput.mutatorNames
-                    let contributorNames = programInput.contributorNames
+                for input in sortedProgramInputs {
+                    let programHash = input.programHash
+                    let mutatorNames = input.mutatorNames
+                    let contributorNames = input.contributorNames
 
                     // Format as PostgreSQL arrays: ARRAY['name1', 'name2', ...]
                     let mutatorsArray: String
@@ -539,7 +492,7 @@ public actor PostgresSQLStorage {
                     // Single unified INSERT into the program table (previously split between fuzzer and program tables)
                     let programQuery = PostgresQuery(stringLiteral: """
                         INSERT INTO program (program_hash, fuzzer_id, inserted_at, program_base64, created_at, source_mutators, contributors, parent_program_hash) 
-                        VALUES ('\(programHash)', \(fuzzerId), NOW(), '\(prepared.programData)', NOW(), \(mutatorsArray), \(contributorsArray), \(prepared.parentHash != nil ? "'\(prepared.parentHash!)'" : "NULL"))
+                        VALUES ('\(programHash)', \(fuzzerId), NOW(), '\(input.programBase64)', NOW(), \(mutatorsArray), \(contributorsArray), \(input.parentHash != nil ? "'\(input.parentHash!)'" : "NULL"))
                         ON CONFLICT (program_hash) DO NOTHING
                         RETURNING program_hash
                     """)
@@ -683,7 +636,7 @@ public actor PostgresSQLStorage {
         }
 
         var allPrograms: [Program] = []
-        var seenHashes = Set<String>()
+        //var seenHashes = Set<String>()
         let batchSize = 5000  // Fetch 5k programs at a time
         var offset = 0
         var totalFetched = 0
@@ -696,6 +649,7 @@ public actor PostgresSQLStorage {
             }
 
             // Deduplicate and add to results
+            /*
             for program in batch {
                 do {
                     let hash = try DatabaseUtils.calculateProgramHash(program: program)
@@ -710,6 +664,7 @@ public actor PostgresSQLStorage {
                     }
                 }
             }
+            */
 
             totalFetched += batch.count
 

Sources/Fuzzilli/Modules/PostgreSQLSync.swift

@@ -111,7 +111,7 @@ public class PostgreSQLSync: Module {
             self.cleanupCache(&self.executionCache)
         }
 
-        // Cache all mutator names and contributor names from ProgramGenerated event (before minimization)
+        // Cache all mutator names and contributor names from ProgramGenerated event (before minddimization)
         // This works around the fact that contributors don't survive protobuf serialization
         fuzzer.registerEventListener(for: fuzzer.events.ProgramGenerated) { program in
             let programId = program.id.uuidString
@@ -149,15 +149,7 @@ public class PostgreSQLSync: Module {
                     return
                 }
 
-                let mutatorNames =  self.mutatorCache[programId] ?? [] 
-                let contributorNames = self.contributorCache[programId] ?? []
-                
-                let programInput = PostgresSQLStorage.ProgramInput(
-                    program: program,
-                    fuzzerId: fuzzerId,
-                    mutatorNames: mutatorNames,
-                    contributorNames: contributorNames
-                )
+                let programInput = try self.prepareProgramInput(program: program, fuzzerId: fuzzerId, programId: programId)
                 await self.storage.addProgramToBatch(programInput)
 
                 if let execution = execution {
@@ -200,18 +192,8 @@ public class PostgreSQLSync: Module {
                         feedbackNexusCount = Int(evaluator.getFeedbackNexusCount())
                     }
 
-                    let programHash: String
-                    do {
-                        programHash = try DatabaseUtils.calculateProgramHash(program: program)
-                    } catch {
-                        if self.enableLogging {
-                            self.logger.warning("Failed to calculate program hash, skipping execution record: \(error)")
-                        }
-                        return
-                    }
-                    
                     let executionInput = PostgresSQLStorage.ExecutionInput(
-                        programHash: programHash,
+                        programHash: programInput.programHash,
                         executionOutcomeId: outcomeId,
                         coverageTotal: coverageTotal,
                         edgesFound: edgesFound,
@@ -254,15 +236,7 @@ public class PostgreSQLSync: Module {
                     return
                 }
 
-                let mutatorNames =  self.mutatorCache[programId] ?? [] 
-                let contributorNames = self.contributorCache[programId] ?? []
-                
-                let programInput = PostgresSQLStorage.ProgramInput(
-                    program: program,
-                    fuzzerId: fuzzerId,
-                    mutatorNames: mutatorNames,
-                    contributorNames: contributorNames
-                )
+                let programInput = try self.prepareProgramInput(program: program, fuzzerId: fuzzerId, programId: programId)
                 await self.storage.addProgramToBatch(programInput)
 
                 // Get coverage metrics if available (crashes may still have coverage)
@@ -289,19 +263,9 @@ public class PostgreSQLSync: Module {
                     feedbackNexusCount = Int(evaluator.getFeedbackNexusCount())
                 }
 
-                let programHash: String
-                do {
-                    programHash = try DatabaseUtils.calculateProgramHash(program: program)
-                } catch {
-                    if self.enableLogging {
-                        self.logger.warning("Failed to calculate program hash for crash, skipping execution record: \(error)")
-                    }
-                    return
-                }
-                
                 // Create execution record with outcome_id = 1 (Crashed)
                 let executionInput = PostgresSQLStorage.ExecutionInput(
-                    programHash: programHash,
+                    programHash: programInput.programHash,
                     executionOutcomeId: 1,  // Crashed
                     coverageTotal: coverageTotal,
                     edgesFound: edgesFound,
@@ -473,4 +437,34 @@ public class PostgreSQLSync: Module {
             }
         }
     }
+
+    private func prepareProgramInput(program: Program, fuzzerId: Int, programId: String) throws -> PostgresSQLStorage.ProgramInput {
+        let mutatorNames =  self.mutatorCache[programId] ?? [] 
+        self.mutatorCache.removeValue(forKey: programId)
+
+        let contributorNames = self.contributorCache[programId] ?? []
+        self.contributorCache.removeValue(forKey: programId)
+                
+        let (programHash, programBase64) = try DatabaseUtils.prepareProgram(program: program)
+        let parentHash: String?
+        if let parentProgram = program.parent {
+            parentHash = try? DatabaseUtils.prepareProgram(program: parentProgram).hash
+        } else {
+            parentHash = nil
+        }
+
+        return PostgresSQLStorage.ProgramInput(
+            program: program,
+            programHash: programHash,
+            programBase64: programBase64,
+            fuzzerId: fuzzerId,
+            mutatorNames: mutatorNames,
+            contributorNames: contributorNames,
+            parentHash: parentHash
+        )
+    }
+
+    private func prepareExecutionInput() {
+
+    }
 }