From b3a45ae9b260fa4e1d1db43ed86cd70c282d4ed9 Mon Sep 17 00:00:00 2001
From: Felicitas Pojtinger <felicitas@pojtinger.com>
Date: Tue, 30 Nov 2021 18:44:27 +0100
Subject: [PATCH] feat: Add ZStandard compression support

---
 cmd/stbak/cmd/archive.go        | 74 +++++++++++++++++++++++++++++++++
 cmd/stbak/cmd/recovery_fetch.go | 10 +++++
 cmd/stbak/cmd/recovery_index.go | 22 +++++-----
 cmd/stbak/cmd/root.go           |  5 ++-
 4 files changed, 99 insertions(+), 12 deletions(-)

diff --git a/cmd/stbak/cmd/archive.go b/cmd/stbak/cmd/archive.go
index b93ee68..a142d68 100644
--- a/cmd/stbak/cmd/archive.go
+++ b/cmd/stbak/cmd/archive.go
@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"strconv"
 
+	"github.com/klauspost/compress/zstd"
 	"github.com/klauspost/pgzip"
 	"github.com/pierrec/lz4/v4"
 	"github.com/pojntfx/stfs/pkg/adapters"
@@ -253,6 +254,43 @@ func archive(
 				hdr.Size = int64(fileSizeCounter.BytesRead)
 
 				hdr.Name += compressionFormatLZ4Suffix
+			case compressionFormatZStandardKey:
+				// Get the compressed size for the header
+				file, err := os.Open(path)
+				if err != nil {
+					return err
+				}
+
+				fileSizeCounter := counters.CounterWriter{
+					Writer: io.Discard,
+				}
+
+				zz, err := zstd.NewWriter(&fileSizeCounter)
+				if err != nil {
+					return err
+				}
+
+				if _, err := io.Copy(zz, file); err != nil {
+					return err
+				}
+
+				if err := zz.Flush(); err != nil {
+					return err
+				}
+				if err := zz.Close(); err != nil {
+					return err
+				}
+				if err := file.Close(); err != nil {
+					return err
+				}
+
+				if hdr.PAXRecords == nil {
+					hdr.PAXRecords = map[string]string{}
+				}
+				hdr.PAXRecords[pax.STFSRecordUncompressedSize] = strconv.Itoa(int(hdr.Size))
+				hdr.Size = int64(fileSizeCounter.BytesRead)
+
+				hdr.Name += compressionFormatZStandardSuffix
 			case compressionFormatNoneKey:
 			default:
 				return errUnsupportedCompressionFormat
@@ -352,6 +390,42 @@ func archive(
 			if err := file.Close(); err != nil {
 				return err
 			}
+		case compressionFormatZStandardKey:
+			// Compress and write the file
+			file, err := os.Open(path)
+			if err != nil {
+				return err
+			}
+
+			zz, err := zstd.NewWriter(tw)
+			if err != nil {
+				return err
+			}
+
+			if _, err := io.Copy(zz, file); err != nil {
+				return err
+			}
+
+			if isRegular {
+				if _, err := io.Copy(zz, file); err != nil {
+					return err
+				}
+			} else {
+				buf := make([]byte, controllers.BlockSize*recordSize)
+				if _, err := io.CopyBuffer(zz, file, buf); err != nil {
+					return err
+				}
+			}
+
+			if err := zz.Flush(); err != nil {
+				return err
+			}
+			if err := zz.Close(); err != nil {
+				return err
+			}
+			if err := file.Close(); err != nil {
+				return err
+			}
 		case compressionFormatNoneKey:
 			// Write the file
 			file, err := os.Open(path)
diff --git a/cmd/stbak/cmd/recovery_fetch.go b/cmd/stbak/cmd/recovery_fetch.go
index 1976b0d..8a601cd 100644
--- a/cmd/stbak/cmd/recovery_fetch.go
+++ b/cmd/stbak/cmd/recovery_fetch.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/klauspost/compress/zstd"
 	"github.com/klauspost/pgzip"
 	"github.com/pierrec/lz4/v4"
 	"github.com/pojntfx/stfs/pkg/controllers"
@@ -160,6 +161,15 @@ func restoreFromRecordAndBlock(
 			if _, err := io.Copy(dstFile, lz); err != nil {
 				return err
 			}
+		case compressionFormatZStandardKey:
+			zz, err := zstd.NewReader(tr)
+			if err != nil {
+				return err
+			}
+
+			if _, err := io.Copy(dstFile, zz); err != nil {
+				return err
+			}
 		case compressionFormatNoneKey:
 			if _, err := io.Copy(dstFile, tr); err != nil {
 				return err
diff --git a/cmd/stbak/cmd/recovery_index.go b/cmd/stbak/cmd/recovery_index.go
index 030ffca..bd17d22 100644
--- a/cmd/stbak/cmd/recovery_index.go
+++ b/cmd/stbak/cmd/recovery_index.go
@@ -267,20 +267,20 @@ func indexHeader(
 		hdr.Size = int64(size)
 	}
 
-	switch compressionFormat {
-	case compressionFormatGZipKey:
-		fallthrough
-	case compressionFormatParallelGZipKey:
-		if hdr.FileInfo().Mode().IsRegular() {
+	if hdr.FileInfo().Mode().IsRegular() {
+		switch compressionFormat {
+		case compressionFormatGZipKey:
+			fallthrough
+		case compressionFormatParallelGZipKey:
 			hdr.Name = strings.TrimSuffix(hdr.Name, compressionFormatGZipSuffix)
-		}
-	case compressionFormatLZ4Key:
-		if hdr.FileInfo().Mode().IsRegular() {
+		case compressionFormatLZ4Key:
 			hdr.Name = strings.TrimSuffix(hdr.Name, compressionFormatLZ4Suffix)
+		case compressionFormatZStandardKey:
+			hdr.Name = strings.TrimSuffix(hdr.Name, compressionFormatZStandardSuffix)
+		case compressionFormatNoneKey:
+		default:
+			return errUnsupportedCompressionFormat
 		}
-	case compressionFormatNoneKey:
-	default:
-		return errUnsupportedCompressionFormat
 	}
 
 	if err := formatting.PrintCSV(formatting.GetTARHeaderAsCSV(record, block, hdr)); err != nil {
diff --git a/cmd/stbak/cmd/root.go b/cmd/stbak/cmd/root.go
index 8fd7d78..b80fa60 100644
--- a/cmd/stbak/cmd/root.go
+++ b/cmd/stbak/cmd/root.go
@@ -26,10 +26,13 @@ const (
 
 	compressionFormatLZ4Key    = "lz4"
 	compressionFormatLZ4Suffix = ".lz4"
+
+	compressionFormatZStandardKey    = "zstandard"
+	compressionFormatZStandardSuffix = ".zst"
 )
 
 var (
-	knownCompressionFormats = []string{compressionFormatNoneKey, compressionFormatGZipKey, compressionFormatParallelGZipKey, compressionFormatLZ4Key}
+	knownCompressionFormats = []string{compressionFormatNoneKey, compressionFormatGZipKey, compressionFormatParallelGZipKey, compressionFormatLZ4Key, compressionFormatZStandardKey}
 
 	errUnknownCompressionFormat     = errors.New("unknown compression format")
 	errUnsupportedCompressionFormat = errors.New("unsupported compression format")