From 0e53dc040a230719c452ea1678b4b584f3a21042 Mon Sep 17 00:00:00 2001
From: Thomas Waldmann <tw@waldmann-edv.de>
Date: Tue, 3 May 2022 04:15:01 +0200
Subject: [PATCH] api: enable giving already compressed data

the api already offered getting compressed data,
but not giving compressed data.

thus: cache.add_chunk and key.encrypt improved.
---
 src/borg/cache.py      | 18 ++++++++++++------
 src/borg/crypto/key.py | 22 +++++++++++++---------
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/src/borg/cache.py b/src/borg/cache.py
index 46848c8da..6fa74e692 100644
--- a/src/borg/cache.py
+++ b/src/borg/cache.py
@@ -939,14 +939,17 @@ class LocalCache(CacheStatsMixin):
         self.cache_config.ignored_features.update(repo_features - my_features)
         self.cache_config.mandatory_features.update(repo_features & my_features)
 
-    def add_chunk(self, id, chunk, stats, overwrite=False, wait=True):
+    def add_chunk(self, id, chunk, stats, *, overwrite=False, wait=True, compress=True, size=None):
         if not self.txn_active:
             self.begin_txn()
-        size = len(chunk)
+        if size is None and compress:
+            size = len(chunk)  # chunk is still uncompressed
         refcount = self.seen_chunk(id, size)
         if refcount and not overwrite:
             return self.chunk_incref(id, stats)
-        data = self.key.encrypt(id, chunk)
+        if size is None:
+            raise ValueError("when giving compressed data for a new chunk, the uncompressed size must be given also")
+        data = self.key.encrypt(id, chunk, compress=compress)
         csize = len(data)
         self.repository.put(id, data, wait=wait)
         self.chunks.add(id, 1, size, csize)
@@ -1103,15 +1106,18 @@ Chunk index:    {0.total_unique_chunks:20d}             unknown"""
     def memorize_file(self, hashed_path, path_hash, st, ids):
         pass
 
-    def add_chunk(self, id, chunk, stats, overwrite=False, wait=True):
+    def add_chunk(self, id, chunk, stats, *, overwrite=False, wait=True, compress=True, size=None):
         assert not overwrite, 'AdHocCache does not permit overwrites — trying to use it for recreate?'
         if not self._txn_active:
             self.begin_txn()
-        size = len(chunk)
+        if size is None and compress:
+            size = len(chunk)  # chunk is still uncompressed
+        if size is None:
+            raise ValueError("when giving compressed data for a chunk, the uncompressed size must be given also")
         refcount = self.seen_chunk(id, size)
         if refcount:
             return self.chunk_incref(id, stats, size=size)
-        data = self.key.encrypt(id, chunk)
+        data = self.key.encrypt(id, chunk, compress=compress)
         csize = len(data)
         self.repository.put(id, data, wait=wait)
         self.chunks.add(id, 1, size, csize)
diff --git a/src/borg/crypto/key.py b/src/borg/crypto/key.py
index 8a561d51f..ed4c375d2 100644
--- a/src/borg/crypto/key.py
+++ b/src/borg/crypto/key.py
@@ -167,7 +167,7 @@ class KeyBase:
         """
         raise NotImplementedError
 
-    def encrypt(self, id, data):
+    def encrypt(self, id, data, compress=True):
         pass
 
     def decrypt(self, id, data, decompress=True):
@@ -273,8 +273,9 @@ class PlaintextKey(KeyBase):
     def id_hash(self, data):
         return sha256(data).digest()
 
-    def encrypt(self, id, data):
-        data = self.compressor.compress(data)
+    def encrypt(self, id, data, compress=True):
+        if compress:
+            data = self.compressor.compress(data)
         return b''.join([self.TYPE_STR, data])
 
     def decrypt(self, id, data, decompress=True):
@@ -349,8 +350,9 @@ class AESKeyBase(KeyBase):
 
     logically_encrypted = True
 
-    def encrypt(self, id, data):
-        data = self.compressor.compress(data)
+    def encrypt(self, id, data, compress=True):
+        if compress:
+            data = self.compressor.compress(data)
         next_iv = self.nonce_manager.ensure_reservation(self.cipher.next_iv(),
                                                         self.cipher.block_count(len(data)))
         return self.cipher.encrypt(data, header=self.TYPE_STR, iv=next_iv)
@@ -809,8 +811,9 @@ class AuthenticatedKeyBase(AESKeyBase, FlexiKey):
         if manifest_data is not None:
             self.assert_type(manifest_data[0])
 
-    def encrypt(self, id, data):
-        data = self.compressor.compress(data)
+    def encrypt(self, id, data, compress=True):
+        if compress:
+            data = self.compressor.compress(data)
         return b''.join([self.TYPE_STR, data])
 
     def decrypt(self, id, data, decompress=True):
@@ -865,9 +868,10 @@ class AEADKeyBase(KeyBase):
 
     MAX_IV = 2 ** 48 - 1
 
-    def encrypt(self, id, data):
+    def encrypt(self, id, data, compress=True):
         # to encrypt new data in this session we use always self.cipher and self.sessionid
-        data = self.compressor.compress(data)
+        if compress:
+            data = self.compressor.compress(data)
         reserved = b'\0'
         iv = self.cipher.next_iv()
         if iv > self.MAX_IV:  # see the data-structures docs about why the IV range is enough