WIP: add group key tapscript design exploration unit test

Author: ffranr

taprpc/group_key_tapscript_test.go

@@ -0,0 +1,237 @@
+package taprpc
+
+import (
+	"bytes"
+	"crypto/rand"
+	"fmt"
+	"testing"
+
+	"github.com/btcsuite/btcd/btcec/v2"
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcd/txscript"
+	"github.com/lightninglabs/taproot-assets/commitment"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGroupKeyTapscript(t *testing.T) {
+	t.Parallel()
+
+	// Define parameters which are available to all examples.
+	//
+	// Generate a random internal key.
+	internalKey := RandInternalPubKey()
+
+	// Construct an asset ID leaf.
+	assetIDLeaf := txscript.NewBaseTapLeaf(
+		[]byte("something something OP_RETURN <asset ID>"),
+	)
+	assetIDLeafHash := assetIDLeaf.TapHash()
+
+	// Construct a custom user script leaf. This is used to validate the
+	// control block.
+	customScriptLeaf := txscript.NewBaseTapLeaf(
+		[]byte("I'm a custom user script"),
+	)
+
+	// ---------------------------------------------------------------------
+	//
+	// Example with user-provided script tree.
+
+	fullUserTree := &TapscriptFullTree{
+		AllLeaves: []*TapLeaf{
+			{
+				Script: customScriptLeaf.Script,
+			},
+		},
+	}
+	treeNodes, err := UnmarshalTapscriptFullTree(fullUserTree)
+	require.NoError(t, err)
+
+	// Validate leaves similar to IsTaprootAssetCommitmentScript().
+	treePreimage, err := commitment.NewPreimageFromTapscriptTreeNodes(
+		*treeNodes,
+	)
+	require.NoError(t, err)
+
+	userTreeRootHash, err := treePreimage.TapHash()
+	require.NoError(t, err)
+
+	actualTaprootHash1 := tapBranchHash(assetIDLeafHash, *userTreeRootHash)
+	fmt.Printf("Taproot hash with user tree: %v\n", actualTaprootHash1)
+
+	// Construct the user subtree control block.
+	outputKey := txscript.ComputeTaprootOutputKey(
+		&internalKey, actualTaprootHash1[:],
+	)
+	outputKeyIsOdd := outputKey.SerializeCompressed()[0] == 0x03
+
+	inclusionProof := bytes.Join(
+		[][]byte{
+			assetIDLeafHash[:],
+		}, nil,
+	)
+
+	userSubtreeControlBlock := txscript.ControlBlock{
+		InternalKey:     &internalKey,
+		OutputKeyYIsOdd: outputKeyIsOdd,
+		LeafVersion:     txscript.BaseLeafVersion,
+		InclusionProof:  inclusionProof,
+	}
+
+	// Ensure the custom script control block is correct by computing the
+	// root hash given the control block and the custom script leaf.
+	rootCheckBytes := userSubtreeControlBlock.RootHash(
+		customScriptLeaf.Script,
+	)
+
+	var rootCheck chainhash.Hash
+	copy(rootCheck[:], rootCheckBytes)
+
+	require.Equal(t, actualTaprootHash1, rootCheck)
+
+	// ---------------------------------------------------------------------
+	//
+	// Example with user-provided branch only.
+
+	// Construct a branch from two node hashes. (A node hash can be a leaf
+	// hash or a branch hash.) We must use two nodes so that we can be sure
+	// that there is only one asset ID leaf in layer 1.
+	//
+	// In this case, we use the customScriptLeaf as one of the nodes so that
+	// we can test the control block.
+	userLeafHash := customScriptLeaf.TapHash()
+	userSiblingNodeHash := chainhash.Hash([32]byte{3, 4, 5})
+
+	userBranch := &TapBranch{
+		LeftTaphash:  userLeafHash[:],
+		RightTaphash: userSiblingNodeHash[:],
+	}
+	branchNodes, err := UnmarshalTapscriptBranch(userBranch)
+	require.NoError(t, err)
+
+	branchPreimage, err := commitment.NewPreimageFromTapscriptTreeNodes(
+		*branchNodes,
+	)
+	require.NoError(t, err)
+
+	userTreeRootHash, err = branchPreimage.TapHash()
+	require.NoError(t, err)
+
+	actualTaprootHash2 := tapBranchHash(assetIDLeafHash, *userTreeRootHash)
+	fmt.Printf("Taproot hash with user branch: %v\n", actualTaprootHash2)
+
+	// Construct the user subtree control block. This block targets the
+	// custom script leaf.
+	outputKey = txscript.ComputeTaprootOutputKey(
+		&internalKey, actualTaprootHash2[:],
+	)
+	outputKeyIsOdd = outputKey.SerializeCompressed()[0] == 0x03
+
+	inclusionProof = bytes.Join(
+		[][]byte{
+			userSiblingNodeHash[:],
+			assetIDLeafHash[:],
+		}, nil,
+	)
+
+	userSubtreeControlBlock = txscript.ControlBlock{
+		InternalKey:     &internalKey,
+		OutputKeyYIsOdd: outputKeyIsOdd,
+		LeafVersion:     txscript.BaseLeafVersion,
+		InclusionProof:  inclusionProof,
+	}
+
+	// Ensure the custom script control block is correct by computing the
+	// root hash given the control block and the custom script leaf.
+	rootCheckBytes = userSubtreeControlBlock.RootHash(
+		customScriptLeaf.Script,
+	)
+	copy(rootCheck[:], rootCheckBytes)
+
+	require.Equal(t, actualTaprootHash2, rootCheck)
+
+	// ---------------------------------------------------------------------
+	//
+	// Example with an optional user-provided root hash only.
+
+	// Formulate an internal key leaf.
+	internalKeyLeafHash, err := InternalKeyLeafHash(internalKey)
+	require.NoError(t, err)
+
+	// This is the user's custom tapscript tree root hash. It can be set or
+	// unset, up to the user. If unset, using [32]byte{} as the hash will be
+	// fine.
+	//
+	// It can also be a leaf hash or a branch hash.
+	userNodeHash := customScriptLeaf.TapHash()
+
+	branchHash := tapBranchHash(*internalKeyLeafHash, userNodeHash)
+	actualTapscriptRootHash3 := tapBranchHash(assetIDLeafHash, branchHash)
+	fmt.Printf("Tapscript root hash with optional user tree root: %v\n",
+		actualTapscriptRootHash3)
+
+	// Construct the user subtree control block.
+	outputKey = txscript.ComputeTaprootOutputKey(
+		&internalKey, actualTapscriptRootHash3[:],
+	)
+	outputKeyIsOdd = outputKey.SerializeCompressed()[0] == 0x03
+
+	inclusionProof = bytes.Join(
+		[][]byte{
+			internalKeyLeafHash[:],
+			assetIDLeafHash[:],
+		}, nil,
+	)
+
+	userSubtreeControlBlock = txscript.ControlBlock{
+		InternalKey:     &internalKey,
+		OutputKeyYIsOdd: outputKeyIsOdd,
+		LeafVersion:     txscript.BaseLeafVersion,
+		InclusionProof:  inclusionProof,
+	}
+
+	// Ensure the custom script control block is correct by computing the
+	// root hash given the control block and the custom script leaf.
+	rootCheckBytes = userSubtreeControlBlock.RootHash(
+		customScriptLeaf.Script,
+	)
+	copy(rootCheck[:], rootCheckBytes)
+
+	require.Equal(t, actualTapscriptRootHash3, rootCheck)
+}
+
+// Copy of commitment.tapBranchHash, should probably export.
+func tapBranchHash(l, r chainhash.Hash) chainhash.Hash {
+	if bytes.Compare(l[:], r[:]) > 0 {
+		l, r = r, l
+	}
+	return *chainhash.TaggedHash(chainhash.TagTapBranch, l[:], r[:])
+}
+
+func RandInternalPubKey() btcec.PublicKey {
+	randBytes := make([]byte, 32)
+	_, _ = rand.Read(randBytes)
+
+	privateKey, _ := btcec.PrivKeyFromBytes(randBytes)
+	pubKey := privateKey.PubKey()
+	return *pubKey
+}
+
+func InternalKeyLeafHash(internalKey btcec.PublicKey) (*chainhash.Hash, error) {
+	// Construct a tapscript leaf for the internal key. Use OP_RETURN to
+	// ensure that the script can not be executed.
+	leafScript, err := txscript.NewScriptBuilder().
+		AddOp(txscript.OP_RETURN).
+		AddData(internalKey.SerializeCompressed()).
+		Script()
+	if err != nil {
+		return nil, err
+	}
+
+	leaf := txscript.TapLeaf{
+		Script:      leafScript,
+		LeafVersion: txscript.BaseLeafVersion,
+	}
+	leafHash := leaf.TapHash()
+	return &leafHash, nil
+}