From 2e067a7ae8120e822a97561f02b1db1228ad62ed Mon Sep 17 00:00:00 2001
From: Marian Beermann <public@enkore.de>
Date: Fri, 2 Jun 2017 11:36:58 +0200
Subject: [PATCH] repository: add refcount corruption test

---
 src/borg/repository.py           | 17 ++++-----
 src/borg/testsuite/repository.py | 59 ++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 8 deletions(-)

diff --git a/src/borg/repository.py b/src/borg/repository.py
index 43183223b..a34600313 100644
--- a/src/borg/repository.py
+++ b/src/borg/repository.py
@@ -373,17 +373,18 @@ def commit(self, save_space=False):
         self.write_index()
         self.rollback()
 
-    def _read_integrity(self, transaction_id, key=None):
-        integrity_path = os.path.join(self.path, 'integrity.%d' % transaction_id)
+    def _read_integrity(self, transaction_id, key):
+        integrity_file = 'integrity.%d' % transaction_id
+        integrity_path = os.path.join(self.path, integrity_file)
         try:
             with open(integrity_path, 'rb') as fd:
                 integrity = msgpack.unpack(fd)
         except FileNotFoundError:
             return
-        if key:
-            return integrity[key].decode()
-        else:
-            return integrity
+        if integrity.get(b'version') != 2:
+            logger.warning('Unknown integrity data version %r in %s', integrity.get(b'version'), integrity_file)
+            return
+        return integrity[key].decode()
 
     def open_index(self, transaction_id, auto_recover=True):
         if transaction_id is None:
@@ -617,7 +618,7 @@ def complete_xfer(intermediate=True):
             # get rid of the old, sparse, unused segments. free space.
             for segment in unused:
                 logger.debug('complete_xfer: deleting unused segment %d', segment)
-                assert self.segments.pop(segment) == 0
+                assert self.segments.pop(segment) == 0, 'Corrupted segment reference count - corrupted index or hints'
                 self.io.delete_segment(segment)
                 del self.compact[segment]
             unused = []
@@ -711,7 +712,7 @@ def complete_xfer(intermediate=True):
                             new_segment, size = self.io.write_delete(key)
                         self.compact[new_segment] += size
                         segments.setdefault(new_segment, 0)
-            assert segments[segment] == 0
+            assert segments[segment] == 0, 'Corrupted segment reference count - corrupted index or hints'
             unused.append(segment)
             pi.show()
         pi.finish()
diff --git a/src/borg/testsuite/repository.py b/src/borg/testsuite/repository.py
index f507504e6..cfbb3283d 100644
--- a/src/borg/testsuite/repository.py
+++ b/src/borg/testsuite/repository.py
@@ -6,6 +6,8 @@
 import tempfile
 from unittest.mock import patch
 
+import msgpack
+
 import pytest
 
 from ..hashindex import NSIndex
@@ -555,6 +557,63 @@ def test_unreadable_index(self):
         with self.assert_raises(OSError):
             self.do_commit()
 
+    def test_unknown_integrity_version(self):
+        integrity_path = os.path.join(self.repository.path, 'integrity.1')
+        with open(integrity_path, 'r+b') as fd:
+            msgpack.pack({
+                b'version': 4.7,
+            }, fd)
+            fd.truncate()
+        with self.repository:
+            assert len(self.repository) == 1
+            assert self.repository.get(H(0)) == b'foo'
+
+    def _subtly_corrupted_hints_setup(self):
+        with self.repository:
+            self.repository.append_only = True
+            assert len(self.repository) == 1
+            assert self.repository.get(H(0)) == b'foo'
+            self.repository.put(H(1), b'bar')
+            self.repository.put(H(2), b'baz')
+            self.repository.commit()
+            self.repository.put(H(2), b'bazz')
+            self.repository.commit()
+
+        hints_path = os.path.join(self.repository.path, 'hints.5')
+        with open(hints_path, 'r+b') as fd:
+            hints = msgpack.unpack(fd)
+            fd.seek(0)
+            # Corrupt segment refcount
+            assert hints[b'segments'][2] == 1
+            hints[b'segments'][2] = 0
+            msgpack.pack(hints, fd)
+            fd.truncate()
+
+    def test_subtly_corrupted_hints(self):
+        self._subtly_corrupted_hints_setup()
+        with self.repository:
+            self.repository.append_only = False
+            self.repository.put(H(3), b'1234')
+            # Do a compaction run. Succeeds, since the failed checksum prompted a rebuild of the index+hints.
+            self.repository.commit()
+
+            assert len(self.repository) == 4
+            assert self.repository.get(H(0)) == b'foo'
+            assert self.repository.get(H(1)) == b'bar'
+            assert self.repository.get(H(2)) == b'bazz'
+
+    def test_subtly_corrupted_hints_without_integrity(self):
+        self._subtly_corrupted_hints_setup()
+        integrity_path = os.path.join(self.repository.path, 'integrity.5')
+        os.unlink(integrity_path)
+        with self.repository:
+            self.repository.append_only = False
+            self.repository.put(H(3), b'1234')
+            # Do a compaction run. Fails, since the corrupted refcount was not detected and leads to an assertion failure.
+            with pytest.raises(AssertionError) as exc_info:
+                self.repository.commit()
+            assert 'Corrupted segment reference count' in str(exc_info.value)
+
 
 class RepositoryCheckTestCase(RepositoryTestCaseBase):