From b0bc06ae000ff620f2b7a0c1b24f8aeec6803508 Mon Sep 17 00:00:00 2001
From: Victor Shyba <victor1984@riseup.net>
Date: Tue, 25 Sep 2018 22:06:22 -0300
Subject: [PATCH] improve reorgs

---
 .../test_blockchain_reorganization.py         | 27 ++++++++++++++++---
 torba/baseheader.py                           |  2 +-
 tox.ini                                       |  3 ++-
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/tests/integration/test_blockchain_reorganization.py b/tests/integration/test_blockchain_reorganization.py
index f61d31f7c..0d6cd4068 100644
--- a/tests/integration/test_blockchain_reorganization.py
+++ b/tests/integration/test_blockchain_reorganization.py
@@ -1,17 +1,36 @@
 from orchstr8.testcase import IntegrationTestCase
+from asyncio import sleep
 
 
 class BlockchainReorganizationTests(IntegrationTestCase):
 
     VERBOSE = True
 
-    async def test(self):
+    async def test_reorg(self):
         self.assertEqual(self.ledger.headers.height, 200)
 
         await self.blockchain.generate(1)
         await self.on_header(201)
         self.assertEqual(self.ledger.headers.height, 201)
+        height = 201
 
-        await self.blockchain.invalidateblock(self.ledger.headers.hash(201).decode())
-        await self.blockchain.generate(2)
-        await self.on_header(203)
+        #simple fork (rewind+advance to immediate best)
+        height = await self._simulate_reorg(height, 1, 1, 2)
+        height = await self._simulate_reorg(height, 2, 1, 10)
+        height = await self._simulate_reorg(height, 4, 1, 3)
+        #lagged fork (rewind+batch catch up with immediate best)
+        height = await self._simulate_reorg(height, 4, 2, 3)
+        await self._simulate_reorg(height, 4, 4, 3)
+
+    async def _simulate_reorg(self, height, rewind, winners, progress):
+        for i in range(rewind):
+            await self.blockchain.invalidateblock(self.ledger.headers.hash(height - i).decode())
+        await self.blockchain.generate(rewind + winners)
+        height = height + winners
+        await self.on_header(height)
+        for i in range(progress):
+            await self.blockchain.generate(1)
+            height += 1
+            await self.on_header(height)
+        self.assertEquals(height, self.ledger.headers.height)
+        return height
diff --git a/torba/baseheader.py b/torba/baseheader.py
index a4b451919..9fca794f0 100644
--- a/torba/baseheader.py
+++ b/torba/baseheader.py
@@ -107,7 +107,7 @@ class BaseHeaders:
                     yield threads.deferToThread(self.validate_chunk, height, chunk)
                 except InvalidHeader as e:
                     bail = True
-                    chunk = chunk[:(height-e.height)*self.header_size]
+                    chunk = chunk[:(height-e.height+1)*self.header_size]
                 written = 0
                 if chunk:
                     self.io.seek(height * self.header_size, os.SEEK_SET)
diff --git a/tox.ini b/tox.ini
index 8eb3120ed..e17f52069 100644
--- a/tox.ini
+++ b/tox.ini
@@ -20,4 +20,5 @@ commands =
   unit: coverage run -p --source={envsitepackagesdir}/torba -m twisted.trial unit
   integration: orchstr8 download
   integration: coverage run -p --source={envsitepackagesdir}/torba -m twisted.trial --reactor=asyncio integration.test_transactions
-  integration: coverage run -p --source={envsitepackagesdir}/torba -m twisted.trial --reactor=asyncio integration.test_blockchain_reorganization
+  # Too slow on Travis
+  # integration: coverage run -p --source={envsitepackagesdir}/torba -m twisted.trial --reactor=asyncio integration.test_blockchain_reorganization