From 8f2ce1e638f596cccd89543f2e719a4bf0f3c598 Mon Sep 17 00:00:00 2001
From: Christian Decker <decker.christian@gmail.com>
Date: Thu, 7 May 2020 10:24:27 +0930
Subject: [PATCH] pytest: Add a test for the commitment_revocation hook

---
 tests/plugins/watchtower.py | 14 ++++++++
 tests/test_plugin.py        | 72 +++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)
 create mode 100755 tests/plugins/watchtower.py

diff --git a/tests/plugins/watchtower.py b/tests/plugins/watchtower.py
new file mode 100755
index 000000000..e36ce3577
--- /dev/null
+++ b/tests/plugins/watchtower.py
@@ -0,0 +1,14 @@
+#!/usr/bin/env python3
+
+from pyln.client import Plugin
+
+plugin = Plugin()
+
+
+@plugin.hook('commitment_revocation')
+def on_commitment_revocation(commitment_txid, penalty_tx, plugin, **kwargs):
+    with open('watchtower.csv', 'a') as f:
+        f.write("{}, {}\n".format(commitment_txid, penalty_tx))
+
+
+plugin.run()
diff --git a/tests/test_plugin.py b/tests/test_plugin.py
index 65dc9702a..f8a6e5d70 100644
--- a/tests/test_plugin.py
+++ b/tests/test_plugin.py
@@ -1270,6 +1270,78 @@ def test_replacement_payload(node_factory):
     assert l2.daemon.wait_for_log("Attept to pay.*with wrong secret")
 
 
+@unittest.skipIf(not DEVELOPER, "Requires dev_sign_last_tx")
+def test_watchtower(node_factory, bitcoind, directory, chainparams):
+    """Test watchtower hook.
+
+    l1 and l2 open a channel, make a couple of updates and then l1 cheats on
+    l2 while that one is offline. The watchtower plugin meanwhile stashes all
+    the penalty transactions and we release the one matching the offending
+    commitment transaction.
+
+    """
+    p = os.path.join(os.path.dirname(__file__), "plugins/watchtower.py")
+    l1, l2 = node_factory.line_graph(
+        2,
+        opts=[{'may_fail': True, 'allow_broken_log': True}, {'plugin': p}]
+    )
+
+    # Force a new commitment
+    l1.rpc.pay(l2.rpc.invoice(25000000, 'lbl1', 'desc1')['bolt11'])
+
+    tx = l1.rpc.dev_sign_last_tx(l2.info['id'])['tx']
+
+    # Now make sure it is out of date
+    l1.rpc.pay(l2.rpc.invoice(25000000, 'lbl2', 'desc2')['bolt11'])
+
+    # l2 stops watching the chain, allowing the watchtower to react
+    l2.stop()
+
+    # Now l1 cheats
+    bitcoind.rpc.sendrawtransaction(tx)
+    time.sleep(1)
+    bitcoind.generate_block(1)
+
+    wt_file = os.path.join(
+        l2.daemon.lightning_dir,
+        chainparams['name'],
+        'watchtower.csv'
+    )
+
+    cheat_tx = bitcoind.rpc.decoderawtransaction(tx)
+    for l in open(wt_file, 'r'):
+        txid, penalty = l.strip().split(', ')
+        if txid == cheat_tx['txid']:
+            # This one should succeed, since it is a response to the cheat_tx
+            bitcoind.rpc.sendrawtransaction(penalty)
+            break
+
+    # Need this to check that l2 gets the funds
+    penalty_meta = bitcoind.rpc.decoderawtransaction(penalty)
+
+    time.sleep(1)
+    bitcoind.generate_block(1)
+
+    # Make sure l2's normal penalty_tx doesn't reach the network
+    def mock_sendrawtransaction(tx):
+        print("NOT broadcasting", tx)
+
+    l2.daemon.rpcproxy.mock_rpc('sendrawtransaction', mock_sendrawtransaction)
+
+    # Restart l2, and it should continue where the watchtower left off:
+    l2.start()
+
+    # l2 will still try to broadcast its latest commitment tx, but it'll fail
+    # since l1 has cheated. All commitments share the same prefix, so look for
+    # that.
+    penalty_prefix = tx[:(4 + 1 + 36) * 2]  # version, txin_count, first txin in hex
+    l2.daemon.wait_for_log(r'Expected error broadcasting tx {}'.format(penalty_prefix))
+
+    # Now make sure the penalty output ends up in our wallet
+    fund_txids = [o['txid'] for o in l2.rpc.listfunds()['outputs']]
+    assert(penalty_meta['txid'] in fund_txids)
+
+
 def test_plugin_fail(node_factory):
     """Test that a plugin which fails (not during a command)"""
     plugin = os.path.join(os.path.dirname(__file__), 'plugins/fail_by_itself.py')