Fix Issue #148 - Allow mend to fix open Zstandard files

cmd/warc/mend/mend.go

@@ -238,6 +238,33 @@
 	return header[0] == 0x1f && header[1] == 0x8b
 }
 
+func isCompressedFile(filepath string) bool {
+    file, err := os.Open(filepath)
+    if err != nil {
+        return false
+    }
+    defer file.Close()
+
+    header := make([]byte, 4)
+    n, err := file.Read(header)
+    if err != nil || n < 4 {
+        return false
+    }
+
+    // gzip magic
+    if header[0] == 0x1f && header[1] == 0x8b {
+        return true
+    }
+
+    // zstd magic 28 B5 2F FD
+    if header[0] == 0x28 && header[1] == 0xB5 && header[2] == 0x2F && header[3] == 0xFD {
+        return true
+    }
+
+    return false
+}
+
+
 func analyzeWARCFile(filepath string, verbose bool, force bool) mendResult {
 	result := mendResult{
 		filepath: filepath,
@@ -253,26 +280,27 @@
 			return result
 		}
 
-		// Only support gzip-compressed WARC files (.gz.open)
-		if !strings.HasSuffix(strings.ToLower(filepath), ".gz.open") {
-			slog.Error("only gzip-compressed WARC files (.gz.open) are supported", "file", filepath)
+		// Support gzip (.gz.open) and zstd (.zst.open)
+		lower := strings.ToLower(filepath)
+		if !(strings.HasSuffix(lower, ".gz.open") || strings.HasSuffix(lower, ".zst.open")) {
+			slog.Error("only gzip or zstd WARC files (.gz.open or .zst.open) are supported", "file", filepath)
 			return result
-		}
-	} else {
+		}else {
 		// With --force, process any gzip WARC file (.gz or .gz.open)
-		if !strings.HasSuffix(strings.ToLower(filepath), ".gz") &&
-			!strings.HasSuffix(strings.ToLower(filepath), ".gz.open") {
-			slog.Error("only gzip-compressed WARC files (.gz or .gz.open) are supported", "file", filepath)
-			return result
-		}
+		if !(strings.HasSuffix(lower, ".gz") || strings.HasSuffix(lower, ".gz.open") || strings.HasSuffix(lower, ".zst") || strings.HasSuffix(lower, ".zst.open")) {
+		slog.Error("only gzip or zstd WARC files (.gz/.zst/.gz.open/.zst.open) are supported", "file", filepath)
+		return result
+}
+
 	}
 
 	// Verify the file is actually gzip compressed by checking magic bytes
-	if !isGzipFile(filepath) {
-		slog.Error("file is not gzip compressed (must be gzip (.gz) format)", "file", filepath)
-		return result
+	if !isCompressedFile(filepath) {
+    slog.Error("file is not gzip or zstd compressed", "file", filepath)
+    return result
 	}
 
+
 	// Get file size
 	fileInfo, err := os.Stat(filepath)
 	if err != nil {
@@ -377,7 +405,7 @@
 	return result
 }
 
 func truncateFile(filepath string, position int64) error {
 	file, err := os.OpenFile(filepath, os.O_RDWR, 0)
 	if err != nil {
 		return fmt.Errorf("failed to open file for truncation: %w", err)
@@ -391,7 +419,7 @@
 	return nil
 }
 
 func confirmAction(prompt string, autoYes bool) bool {
 	if autoYes {
 		fmt.Println("y")
 		return true
@@ -409,7 +437,7 @@
 }
 
 // displaySummary shows final statistics for the mend operation
 func displaySummary(stats mendStats) {
 	elapsed := time.Since(stats.startTime)
 
 	// Build summary parts
@@ -476,7 +504,7 @@
 }
 
 // formatBytes formats a byte count as a human-readable string
 func formatBytes(bytes int64) string {
 	const unit = 1024
 	if bytes < unit {
 		return fmt.Sprintf("%d B", bytes)

cmd/warc/utils/utils.go

@@ -1,14 +1,18 @@
 package utils
 
 import (
-	"log/slog"
-	"os"
-	"strings"
+    "bytes"
+    "compress/gzip"
+    "log/slog"
+    "os"
+    "strings"
 
-	warc "github.com/internetarchive/gowarc"
-	"github.com/spf13/cobra"
+    warc "github.com/internetarchive/gowarc"
+    "github.com/klauspost/compress/zstd"
+    "github.com/spf13/cobra"
 )
 
+
 // GetThreadsFlag extracts the threads flag value from a cobra command
 // Cobra already validates that it's a valid integer, but we still check for errors
 func GetThreadsFlag(cmd *cobra.Command) int {
@@ -21,24 +25,60 @@ func GetThreadsFlag(cmd *cobra.Command) int {
 	return threads
 }
 
-// OpenWARCFile opens a WARC file and returns a reader and file handle
+// OpenWARCFile opens a WARC file (supports: gzip, zstd, uncompressed)
 func OpenWARCFile(filepath string) (*warc.Reader, *os.File, error) {
-	f, err := os.Open(filepath)
-	if err != nil {
-		slog.Error("unable to open file", "err", err.Error(), "file", filepath)
-		return nil, nil, err
-	}
+    f, err := os.Open(filepath)
+    if err != nil {
+        slog.Error("unable to open file", "err", err.Error(), "file", filepath)
+        return nil, nil, err
+    }
 
-	reader, err := warc.NewReader(f)
-	if err != nil {
-		slog.Error("warc.NewReader failed", "err", err.Error(), "file", filepath)
-		f.Close()
-		return nil, nil, err
-	}
+    // Read magic bytes
+    magic := make([]byte, 4)
+    n, err := f.Read(magic)
+    if err != nil || n < 4 {
+        slog.Error("failed to read magic bytes", "file", filepath)
+        f.Close()
+        return nil, nil, err
+    }
+    f.Seek(0, 0)
 
-	return reader, f, nil
+    // GZIP magic bytes: 1F 8B
+    if magic[0] == 0x1F && magic[1] == 0x8B {
+        gz, err := gzip.NewReader(f)
+        if err != nil {
+            slog.Error("gzip reader failed", "err", err.Error(), "file", filepath)
+            f.Close()
+            return nil, nil, err
+        }
+        r, err := warc.NewReader(gz)
+        return r, f, err
+    }
+
+    // ZSTD magic bytes: 28 B5 2F FD
+    if bytes.Equal(magic, []byte{0x28, 0xB5, 0x2F, 0xFD}) {
+        dec, err := zstd.NewReader(f)
+        if err != nil {
+            slog.Error("zstd reader failed", "err", err.Error(), "file", filepath)
+            f.Close()
+            return nil, nil, err
+        }
+        r, err := warc.NewReader(dec)
+        return r, f, err
+    }
+
+    // UNCOMPRESSED fallback
+    reader, err := warc.NewReader(f)
+    if err != nil {
+        slog.Error("warc.NewReader failed", "err", err.Error(), "file", filepath)
+        f.Close()
+        return nil, nil, err
+    }
+
+    return reader, f, nil
 }
 
+
 // ShouldSkipRecord determines if a WARC record should be skipped during processing
 func ShouldSkipRecord(record *warc.Record) bool {
 	// Skip revisit records