diff --git a/cmd/warc/mend/mend.go b/cmd/warc/mend/mend.go index 06065ae..d29162c 100644 --- a/cmd/warc/mend/mend.go +++ b/cmd/warc/mend/mend.go @@ -238,6 +238,33 @@ func isGzipFile(filepath string) bool { 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 @@ func analyzeWARCFile(filepath string, verbose bool, force bool) mendResult { 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 { diff --git a/cmd/warc/utils/utils.go b/cmd/warc/utils/utils.go index b7b8674..3eed22b 100644 --- a/cmd/warc/utils/utils.go +++ b/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