From 9990889ccef1b6b8fdb7085858ca1c5910a296bb Mon Sep 17 00:00:00 2001
From: Victor Shyba <victor1984@riseup.net>
Date: Tue, 9 Jul 2019 01:07:07 -0300
Subject: [PATCH] add repair ability to baseheaders

---
 torba/tests/client_tests/unit/test_headers.py | 21 +++++++++++++++++++
 torba/torba/client/baseheader.py              | 16 +++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/torba/tests/client_tests/unit/test_headers.py b/torba/tests/client_tests/unit/test_headers.py
index 819169cf1..9d3b65cc6 100644
--- a/torba/tests/client_tests/unit/test_headers.py
+++ b/torba/tests/client_tests/unit/test_headers.py
@@ -94,6 +94,27 @@ class BasicHeadersTests(BitcoinHeadersTestCase):
         await headers.connect(len(headers), remainder)
         self.assertEqual(headers.height, 32259)
 
+    async def test_repair(self):
+        headers = MainHeaders(':memory:')
+        await headers.connect(0, self.get_bytes(block_bytes(3001)))
+        self.assertEqual(headers.height, 3000)
+        headers.repair()
+        self.assertEqual(headers.height, 3000)
+        # corrupt the middle of it
+        headers.io.seek(block_bytes(1500))
+        headers.io.write(b"wtf")
+        headers.repair()
+        self.assertEqual(headers.height, 1499)
+        self.assertEqual(len(headers), 1500)
+        # corrupt by appending
+        headers.io.seek(block_bytes(len(headers)))
+        headers.io.write(b"appending")
+        headers._size = None
+        headers.repair()
+        self.assertEqual(headers.height, 1499)
+        await headers.connect(len(headers), self.get_bytes(block_bytes(3001 - 1500), after=block_bytes(1500)))
+        self.assertEqual(headers.height, 3000)
+
     async def test_concurrency(self):
         BLOCKS = 30
         headers_temporary_file = tempfile.mktemp()
diff --git a/torba/torba/client/baseheader.py b/torba/torba/client/baseheader.py
index c1cc62ce9..6130f2fea 100644
--- a/torba/torba/client/baseheader.py
+++ b/torba/torba/client/baseheader.py
@@ -164,12 +164,26 @@ class BaseHeaders:
                         proof_of_work.value, target.value)
                 )
 
+    def repair(self):
+        for height in range(self.height):
+            chunk = self.get_raw_header(height)
+            try:
+                # validate_chunk() is CPU bound and reads previous chunks from file system
+                self.validate_chunk(height, chunk)
+            except InvalidHeader as e:
+                log.warning("Header file corrupted at height %s, truncating it.", e.height)
+                self.io.seek((e.height) * self.header_size, os.SEEK_SET)
+                self.io.truncate()
+                self.io.flush()
+                self._size = None
+                return
+
     @staticmethod
     def get_proof_of_work(header_hash: bytes) -> ArithUint256:
         return ArithUint256(int(b'0x' + header_hash, 16))
 
     def _iterate_chunks(self, height: int, headers: bytes) -> Iterator[Tuple[int, bytes]]:
-        assert len(headers) % self.header_size == 0
+        assert len(headers) % self.header_size == 0, f"{len(headers)} {len(headers)%self.header_size}"
         start = 0
         end = (self.chunk_size - height % self.chunk_size) * self.header_size
         while start < end: