fix: address CodeRabbitAI review comments for database output

- Fix SQL injection in postgres.go using pq.QuoteIdentifier for table/index names
- Fix SQL injection in mysql.go using custom quoteIdentifier function
- Fix race condition in writer.go by checking closed state before channel send
- Add missing idx_url index in mysql.go for parity with PostgreSQL schema
- Include RawHeaders in omitRaw check for consistency

Author: dogancanbakir

internal/db/mysql.go

@@ -5,11 +5,16 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
+	"strings"
 
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/projectdiscovery/httpx/runner"
 )
 
+func quoteIdentifier(name string) string {
+	return "`" + strings.ReplaceAll(name, "`", "``") + "`"
+}
+
 func init() {
 	Register(MySQL, newMySQLDatabase)
 }
@@ -45,6 +50,7 @@ func (m *mysqlDatabase) Close() error {
 }
 
 func (m *mysqlDatabase) EnsureSchema(ctx context.Context) error {
+	tableName := quoteIdentifier(m.cfg.TableName)
 	schema := fmt.Sprintf(`
 		CREATE TABLE IF NOT EXISTS %s (
 			id BIGINT AUTO_INCREMENT PRIMARY KEY,
@@ -118,7 +124,7 @@ func (m *mysqlDatabase) EnsureSchema(ctx context.Context) error {
 
 			-- Metrics
 			words INT,
-			` + "`lines`" + ` INT,
+			`+"`lines`"+` INT,
 
 			-- Headers and extracts
 			header JSON,
@@ -146,10 +152,11 @@ func (m *mysqlDatabase) EnsureSchema(ctx context.Context) error {
 			trace JSON,
 
 			INDEX idx_timestamp (timestamp),
+			INDEX idx_url (url(255)),
 			INDEX idx_host (host),
 			INDEX idx_status_code (status_code)
 		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
-	`, m.cfg.TableName)
+	`, tableName)
 
 	_, err := m.db.ExecContext(ctx, schema)
 	if err != nil {
@@ -172,6 +179,8 @@ func (m *mysqlDatabase) InsertBatch(ctx context.Context, results []runner.Result
 		_ = tx.Rollback()
 	}()
 
+	// Use quoteIdentifier to safely quote table name to prevent SQL injection
+	tableName := quoteIdentifier(m.cfg.TableName)
 	query := fmt.Sprintf(`
 		INSERT INTO %s (
 			timestamp, url, input, host, port, scheme, path, method, final_url,
@@ -197,7 +206,7 @@ func (m *mysqlDatabase) InsertBatch(ctx context.Context, results []runner.Result
 			?, ?,
 			?, ?, ?, ?, ?,
 			?, ?, ?
-		)`, m.cfg.TableName)
+		)`, tableName)
 
 	stmt, err := tx.PrepareContext(ctx, query)
 	if err != nil {

internal/db/postgres.go

@@ -45,6 +45,13 @@ func (p *postgresDatabase) Close() error {
 }
 
 func (p *postgresDatabase) EnsureSchema(ctx context.Context) error {
+	tableName := pq.QuoteIdentifier(p.cfg.TableName)
+	idxTimestamp := pq.QuoteIdentifier("idx_" + p.cfg.TableName + "_timestamp")
+	idxURL := pq.QuoteIdentifier("idx_" + p.cfg.TableName + "_url")
+	idxHost := pq.QuoteIdentifier("idx_" + p.cfg.TableName + "_host")
+	idxStatusCode := pq.QuoteIdentifier("idx_" + p.cfg.TableName + "_status_code")
+	idxTech := pq.QuoteIdentifier("idx_" + p.cfg.TableName + "_tech")
+
 	schema := fmt.Sprintf(`
 		CREATE TABLE IF NOT EXISTS %s (
 			id BIGSERIAL PRIMARY KEY,
@@ -146,18 +153,18 @@ func (p *postgresDatabase) EnsureSchema(ctx context.Context) error {
 			trace JSONB
 		);
 
-		CREATE INDEX IF NOT EXISTS idx_%s_timestamp ON %s(timestamp DESC);
-		CREATE INDEX IF NOT EXISTS idx_%s_url ON %s(url);
-		CREATE INDEX IF NOT EXISTS idx_%s_host ON %s(host);
-		CREATE INDEX IF NOT EXISTS idx_%s_status_code ON %s(status_code);
-		CREATE INDEX IF NOT EXISTS idx_%s_tech ON %s USING GIN(tech);
+		CREATE INDEX IF NOT EXISTS %s ON %s(timestamp DESC);
+		CREATE INDEX IF NOT EXISTS %s ON %s(url);
+		CREATE INDEX IF NOT EXISTS %s ON %s(host);
+		CREATE INDEX IF NOT EXISTS %s ON %s(status_code);
+		CREATE INDEX IF NOT EXISTS %s ON %s USING GIN(tech);
 	`,
-		p.cfg.TableName,
-		p.cfg.TableName, p.cfg.TableName,
-		p.cfg.TableName, p.cfg.TableName,
-		p.cfg.TableName, p.cfg.TableName,
-		p.cfg.TableName, p.cfg.TableName,
-		p.cfg.TableName, p.cfg.TableName,
+		tableName,
+		idxTimestamp, tableName,
+		idxURL, tableName,
+		idxHost, tableName,
+		idxStatusCode, tableName,
+		idxTech, tableName,
 	)
 
 	_, err := p.db.ExecContext(ctx, schema)
@@ -181,6 +188,7 @@ func (p *postgresDatabase) InsertBatch(ctx context.Context, results []runner.Res
 		_ = tx.Rollback()
 	}()
 
+	tableName := pq.QuoteIdentifier(p.cfg.TableName)
 	query := fmt.Sprintf(`
 		INSERT INTO %s (
 			timestamp, url, input, host, port, scheme, path, method, final_url,
@@ -206,7 +214,7 @@ func (p *postgresDatabase) InsertBatch(ctx context.Context, results []runner.Res
 			$53, $54,
 			$55, $56, $57, $58, $59,
 			$60, $61, $62
-		)`, p.cfg.TableName)
+		)`, tableName)
 
 	stmt, err := tx.PrepareContext(ctx, query)
 	if err != nil {

internal/db/writer.go

@@ -60,6 +60,10 @@ func NewWriter(ctx context.Context, cfg *Config) (*Writer, error) {
 
 func (w *Writer) GetWriterCallback() runner.OnResultCallback {
 	return func(r runner.Result) {
+		if w.closed.Load() {
+			return
+		}
+
 		if r.Err != nil {
 			return
 		}
@@ -68,6 +72,7 @@ func (w *Writer) GetWriterCallback() runner.OnResultCallback {
 			r.Raw = ""
 			r.Request = ""
 			r.ResponseBody = ""
+			r.RawHeaders = ""
 		}
 
 		select {