Batch JSONL output and add trailing commas (#5705)

projectdiscovery · Nov 15, 2024 · ba6a050 · ba6a050
1 parent 3e4ec90
commit ba6a050
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 26 deletions.
diff --git a/internal/runner/options.go b/internal/runner/options.go
@@ -304,10 +304,17 @@ func createReportingOptions(options *types.Options) (*reporting.Options, error)
 			OmitRaw: options.OmitRawRequests,
 		}
 	}
+	// Combine options.
 	if options.JSONLExport != "" {
-		reportingOptions.JSONLExporter = &jsonl.Options{
-			File:    options.JSONLExport,
-			OmitRaw: options.OmitRawRequests,
+		// Combine the CLI options with the config file options with the CLI options taking precedence
+		if reportingOptions.JSONLExporter != nil {
+			reportingOptions.JSONLExporter.File = options.JSONLExport
+			reportingOptions.JSONLExporter.OmitRaw = options.OmitRawRequests
+		} else {
+			reportingOptions.JSONLExporter = &jsonl.Options{
+				File:    options.JSONLExport,
+				OmitRaw: options.OmitRawRequests,
+			}
 		}
 	}
 

diff --git a/pkg/reporting/exporters/jsonl/jsonl.go b/pkg/reporting/exporters/jsonl/jsonl.go
@@ -10,16 +10,21 @@ import (
 )
 
 type Exporter struct {
-	options *Options
-	mutex   *sync.Mutex
-	rows    []output.ResultEvent
+	options    *Options
+	mutex      *sync.Mutex
+	rows       []output.ResultEvent
+	outputFile *os.File
 }
 
 // Options contains the configuration options for JSONL exporter client
 type Options struct {
 	// File is the file to export found JSONL result to
-	File    string `yaml:"file"`
-	OmitRaw bool   `yaml:"omit-raw"`
+	File string `yaml:"file"`
+	// OmitRaw whether to exclude the raw request and response from the output
+	OmitRaw bool `yaml:"omit-raw"`
+	// BatchSize the number of records to keep in memory before writing them out to the JSONL file or 0 to disable
+	// batching (default)
+	BatchSize int `yaml:"batch-size"`
 }
 
 // New creates a new JSONL exporter integration client based on options.
@@ -32,8 +37,7 @@ func New(options *Options) (*Exporter, error) {
 	return exporter, nil
 }
 
-// Export appends the passed result event to the list of objects to be exported to
-// the resulting JSONL file
+// Export appends the passed result event to the list of objects to be exported to the resulting JSONL file
 func (exporter *Exporter) Export(event *output.ResultEvent) error {
 	exporter.mutex.Lock()
 	defer exporter.mutex.Unlock()
@@ -46,40 +50,73 @@ func (exporter *Exporter) Export(event *output.ResultEvent) error {
 	// Add the event to the rows
 	exporter.rows = append(exporter.rows, *event)
 
+	// If the batch size is greater than 0 and the number of rows has reached the batch, flush it to the database
+	if exporter.options.BatchSize > 0 && len(exporter.rows) >= exporter.options.BatchSize {
+		err := exporter.WriteRows()
+		if err != nil {
+			// The error is already logged, return it to bubble up to the caller
+			return err
+		}
+	}
+
 	return nil
 }
 
-// Close writes the in-memory data to the JSONL file specified by options.JSONLExport
-// and closes the exporter after operation
-func (exporter *Exporter) Close() error {
-	exporter.mutex.Lock()
-	defer exporter.mutex.Unlock()
-
-	// Open the JSONL file for writing and create it if it doesn't exist
-	f, err := os.OpenFile(exporter.options.File, os.O_WRONLY|os.O_CREATE, 0644)
-	if err != nil {
-		return errors.Wrap(err, "failed to create JSONL file")
+// WriteRows writes all rows from the rows list to JSONL file and removes them from the list
+func (exporter *Exporter) WriteRows() error {
+	// Open the file for writing if it's not already.
+	// This will recreate the file if it exists, but keep the file handle so that batched writes within the same
+	// execution are appended to the same file.
+	var err error
+	if exporter.outputFile == nil {
+		// Open the JSONL file for writing and create it if it doesn't exist
+		exporter.outputFile, err = os.OpenFile(exporter.options.File, os.O_WRONLY|os.O_CREATE, 0644)
+		if err != nil {
+			return errors.Wrap(err, "failed to create JSONL file")
+		}
 	}
 
-	// Loop through the rows and convert each to a JSON byte array and write to file
-	for _, row := range exporter.rows {
+	// Loop through the rows and write them, removing them as they're entered
+	for len(exporter.rows) > 0 {
+		row := exporter.rows[0]
+
 		// Convert the row to JSON byte array and append a trailing newline. This is treated as a single line in JSONL
 		obj, err := json.Marshal(row)
 		if err != nil {
 			return errors.Wrap(err, "failed to generate row for JSONL report")
 		}
 
 		// Add a trailing newline to the JSON byte array to confirm with the JSONL format
-		obj = append(obj, '\n')
+		obj = append(obj, ',', '\n')
 
 		// Attempt to append the JSON line to file specified in options.JSONLExport
-		if _, err = f.Write(obj); err != nil {
+		if _, err = exporter.outputFile.Write(obj); err != nil {
 			return errors.Wrap(err, "failed to append JSONL line")
 		}
+
+		// Remove the item from the list
+		exporter.rows = exporter.rows[1:]
+	}
+
+	return nil
+}
+
+// Close writes the in-memory data to the JSONL file specified by options.JSONLExport and closes the exporter after
+// operation
+func (exporter *Exporter) Close() error {
+	exporter.mutex.Lock()
+	defer exporter.mutex.Unlock()
+
+	// Write any remaining rows to the file
+	// Write all pending rows
+	err := exporter.WriteRows()
+	if err != nil {
+		// The error is already logged, return it to bubble up to the caller
+		return err
 	}
 
 	// Close the file
-	if err := f.Close(); err != nil {
+	if err := exporter.outputFile.Close(); err != nil {
 		return errors.Wrap(err, "failed to close JSONL file")
 	}
 

diff --git a/pkg/reporting/reporting.go b/pkg/reporting/reporting.go
@@ -197,7 +197,7 @@ func New(options *Options, db string, doNotDedupe bool) (Client, error) {
 	return client, nil
 }
 
-// CreateConfigIfNotExists creates report-config if it doesn't exists
+// CreateConfigIfNotExists creates report-config if it doesn't exist
 func CreateConfigIfNotExists() error {
 	reportingConfig := config.DefaultConfig.GetReportingConfigFilePath()