From 5d413c8e793a439540d064d24fddfc868e1817d0 Mon Sep 17 00:00:00 2001
From: Pieter Wuille <pieter@wuille.net>
Date: Mon, 24 Oct 2022 14:26:29 -0400
Subject: [PATCH] Add feature_taproot case involved invalid internal pubkey

---
 test/functional/feature_taproot.py       | 47 +++++++++++++++++++++++-
 test/functional/test_framework/script.py |  9 +++--
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/test/functional/feature_taproot.py b/test/functional/feature_taproot.py
index cbb2e0338b8..31a6b31225a 100755
--- a/test/functional/feature_taproot.py
+++ b/test/functional/feature_taproot.py
@@ -96,7 +96,14 @@ from test_framework.util import (
     assert_equal,
     random_bytes,
 )
-from test_framework.key import generate_privkey, compute_xonly_pubkey, sign_schnorr, tweak_add_privkey, ECKey
+from test_framework.key import (
+    generate_privkey,
+    compute_xonly_pubkey,
+    sign_schnorr,
+    tweak_add_privkey,
+    ECKey,
+    SECP256K1
+)
 from test_framework.address import (
     hash160,
     program_to_witness,
@@ -661,6 +668,44 @@ def spenders_taproot_active():
     # Test with signature with bit flipped.
     add_spender(spenders, "sig/bitflip", tap=tap, key=secs[0], failure={"signature": bitflipper(default_signature)}, **ERR_SIG_SCHNORR)
 
+    # == Test involving an internal public key not on the curve ==
+
+    # X-only public keys are 32 bytes, but not every 32-byte array is a valid public key; only
+    # around 50% of them are. This does not affect users using correct software; these "keys" have
+    # no corresponding private key, and thus will never appear as output of key
+    # generation/derivation/tweaking.
+    #
+    # Using an invalid public key as P2TR output key makes the UTXO unspendable. Revealing an
+    # invalid public key as internal key in a P2TR script path spend also makes the spend invalid.
+    # These conditions are explicitly spelled out in BIP341.
+    #
+    # It is however hard to create test vectors for this, because it involves "guessing" how a
+    # hypothetical incorrect implementation deals with an obviously-invalid condition, and making
+    # sure that guessed behavior (accepting it in certain condition) doesn't occur.
+    #
+    # The test case added here tries to detect a very specific bug a verifier could have: if they
+    # don't verify whether or not a revealed internal public key in a script path spend is valid,
+    # and (correctly) implement output_key == tweak(internal_key, tweakval) but (incorrectly) treat
+    # tweak(invalid_key, tweakval) as equal the public key corresponding to private key tweakval.
+    # This may seem like a far-fetched edge condition to test for, but in fact, the BIP341 wallet
+    # pseudocode did exactly that (but obviously only triggerable by someone invoking the tweaking
+    # function with an invalid public key, which shouldn't happen).
+
+    # Generate an invalid public key
+    while True:
+        invalid_pub = random_bytes(32)
+        if not SECP256K1.is_x_coord(int.from_bytes(invalid_pub, 'big')):
+            break
+
+    # Implement a test case that detects validation logic which maps invalid public keys to the
+    # point at infinity in the tweaking logic.
+    tap = taproot_construct(invalid_pub, [("true", CScript([OP_1]))], treat_internal_as_infinity=True)
+    add_spender(spenders, "output/invalid_x", tap=tap, key_tweaked=tap.tweak, failure={"leaf": "true", "inputs": []}, **ERR_WITNESS_PROGRAM_MISMATCH)
+
+    # Do the same thing without invalid point, to make sure there is no mistake in the test logic.
+    tap = taproot_construct(pubs[0], [("true", CScript([OP_1]))])
+    add_spender(spenders, "output/invalid_x_mock", tap=tap, key=secs[0], leaf="true", inputs=[])
+
     # == Tests for signature hashing ==
 
     # Run all tests once with no annex, and once with a valid random annex.
diff --git a/test/functional/test_framework/script.py b/test/functional/test_framework/script.py
index 2b70eab4e4c..f531ccc0301 100644
--- a/test/functional/test_framework/script.py
+++ b/test/functional/test_framework/script.py
@@ -12,7 +12,7 @@ import struct
 import unittest
 from typing import List, Dict
 
-from .key import TaggedHash, tweak_add_pubkey
+from .key import TaggedHash, tweak_add_pubkey, compute_xonly_pubkey
 
 from .messages import (
     CTransaction,
@@ -872,7 +872,7 @@ TaprootInfo = namedtuple("TaprootInfo", "scriptPubKey,internal_pubkey,negflag,tw
 # - merklebranch: the merkle branch to use for this leaf (32*N bytes)
 TaprootLeafInfo = namedtuple("TaprootLeafInfo", "script,version,merklebranch,leaf_hash")
 
-def taproot_construct(pubkey, scripts=None):
+def taproot_construct(pubkey, scripts=None, treat_internal_as_infinity=False):
     """Construct a tree of Taproot spending conditions
 
     pubkey: a 32-byte xonly pubkey for the internal pubkey (bytes)
@@ -891,7 +891,10 @@ def taproot_construct(pubkey, scripts=None):
 
     ret, h = taproot_tree_helper(scripts)
     tweak = TaggedHash("TapTweak", pubkey + h)
-    tweaked, negated = tweak_add_pubkey(pubkey, tweak)
+    if treat_internal_as_infinity:
+        tweaked, negated = compute_xonly_pubkey(tweak)
+    else:
+        tweaked, negated = tweak_add_pubkey(pubkey, tweak)
     leaves = dict((name, TaprootLeafInfo(script, version, merklebranch, leaf)) for name, version, script, merklebranch, leaf in ret)
     return TaprootInfo(CScript([OP_1, tweaked]), pubkey, negated + 0, tweak, leaves, h, tweaked)