From 4b957e7373db43f79ccd612d12eba76af8a41e12 Mon Sep 17 00:00:00 2001
From: Michael Eischer <michael.eischer@fau.de>
Date: Sun, 13 Feb 2022 00:12:40 +0100
Subject: [PATCH] repository: Implement index/snapshot/lock compression

The config file is not compressed as it should remain readable by older
restic versions such that these can return a proper error.

As the old format for unpacked data does not include a version header,
make use of a trick: The old data is always encoded as JSON. Thus it can
only start with '{' or '['. For any other value the first byte indicates
a versioned format. The version is set to 2 for now. Then the zstd
compressed data follows.
---
 go.mod                            |  2 +-
 internal/repository/repository.go | 58 +++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index b5d72874f..a20ae44a1 100644
--- a/go.mod
+++ b/go.mod
@@ -21,7 +21,7 @@ require (
 	github.com/hashicorp/golang-lru v0.5.4
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/juju/ratelimit v1.0.1
-	github.com/klauspost/compress v1.15.1 // indirect
+	github.com/klauspost/compress v1.15.1
 	github.com/klauspost/cpuid/v2 v2.0.12 // indirect
 	github.com/kurin/blazer v0.5.4-0.20211030221322-ba894c124ac6
 	github.com/minio/md5-simd v1.1.2 // indirect
diff --git a/internal/repository/repository.go b/internal/repository/repository.go
index d74868895..eb2d0a109 100644
--- a/internal/repository/repository.go
+++ b/internal/repository/repository.go
@@ -12,6 +12,7 @@ import (
 	"sync"
 
 	"github.com/cenkalti/backoff/v4"
+	"github.com/klauspost/compress/zstd"
 	"github.com/restic/chunker"
 	"github.com/restic/restic/internal/backend/dryrun"
 	"github.com/restic/restic/internal/cache"
@@ -40,6 +41,9 @@ type Repository struct {
 
 	treePM *packerManager
 	dataPM *packerManager
+
+	enc *zstd.Encoder
+	dec *zstd.Decoder
 }
 
 // New returns a new repository with backend be.
@@ -51,6 +55,16 @@ func New(be restic.Backend) *Repository {
 		treePM: newPackerManager(be, nil),
 	}
 
+	enc, err := zstd.NewWriter(nil)
+	if err != nil {
+		panic(err)
+	}
+	repo.enc = enc
+	dec, err := zstd.NewReader(nil)
+	if err != nil {
+		panic(err)
+	}
+	repo.dec = dec
 	return repo
 }
 
@@ -125,6 +139,9 @@ func (r *Repository) LoadUnpacked(ctx context.Context, buf []byte, t restic.File
 	if err != nil {
 		return nil, err
 	}
+	if t != restic.ConfigFile {
+		return r.decompressUnpacked(plaintext)
+	}
 
 	return plaintext, nil
 }
@@ -312,9 +329,50 @@ func (r *Repository) SaveJSONUnpacked(ctx context.Context, t restic.FileType, it
 	return r.SaveUnpacked(ctx, t, plaintext)
 }
 
+func (r *Repository) compressUnpacked(p []byte) ([]byte, error) {
+	// compression is only available starting from version 2
+	if r.cfg.Version < 2 {
+		return p, nil
+	}
+
+	// version byte
+	out := []byte{2}
+	out = r.enc.EncodeAll(p, out)
+	return out, nil
+}
+
+func (r *Repository) decompressUnpacked(p []byte) ([]byte, error) {
+	// compression is only available starting from version 2
+	if r.cfg.Version < 2 {
+		return p, nil
+	}
+
+	if len(p) < 1 {
+		// too short for version header
+		return p, nil
+	}
+	if p[0] == '[' || p[0] == '{' {
+		// probably raw JSON
+		return p, nil
+	}
+	// version
+	if p[0] != 2 {
+		return nil, errors.New("not supported encoding format")
+	}
+
+	return r.dec.DecodeAll(p[1:], nil)
+}
+
 // SaveUnpacked encrypts data and stores it in the backend. Returned is the
 // storage hash.
 func (r *Repository) SaveUnpacked(ctx context.Context, t restic.FileType, p []byte) (id restic.ID, err error) {
+	if t != restic.ConfigFile {
+		p, err = r.compressUnpacked(p)
+		if err != nil {
+			return restic.ID{}, err
+		}
+	}
+
 	ciphertext := restic.NewBlobBuffer(len(p))
 	ciphertext = ciphertext[:0]
 	nonce := crypto.NewRandomNonce()