From 6d27d9244786eb6d95319725565503c35069886b Mon Sep 17 00:00:00 2001
From: ffranr <ffranr@proton.me>
Date: Tue, 17 Dec 2024 13:51:34 +0000
Subject: [PATCH] WIP: GroupKeyRevealV1 unit tests

---
 asset/group_key_reveal_test.go | 180 +++++++++++++++++++++++++++++++++
 1 file changed, 180 insertions(+)
 create mode 100644 asset/group_key_reveal_test.go

diff --git a/asset/group_key_reveal_test.go b/asset/group_key_reveal_test.go
new file mode 100644
index 000000000..20f97f0f4
--- /dev/null
+++ b/asset/group_key_reveal_test.go
@@ -0,0 +1,180 @@
+package asset
+
+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/fn"
+	"github.com/lightninglabs/taproot-assets/internal/test"
+	"github.com/stretchr/testify/require"
+)
+
+//// MockGKRRecord creates a TLV record for a GroupKeyReveal.
+//func MockGKRRecord(reveal *GroupKeyReveal) tlv.Record {
+//	recordSize := func() uint64 {
+//		if reveal == nil || *reveal == nil {
+//			return 0
+//		}
+//		r := *reveal
+//		return uint64(
+//			btcec.PubKeyBytesLenCompressed + len(r.TapscriptRoot()),
+//		)
+//	}
+//	return tlv.MakeDynamicRecord(
+//		0, reveal, recordSize, GroupKeyRevealEncoder, GroupKeyRevealDecoder,
+//	)
+//}
+
+type testCaseGkrEncodeDecode struct {
+	testName string
+
+	internalKey       btcec.PublicKey
+	genesisAssetID    ID
+	customSubtreeRoot fn.Option[chainhash.Hash]
+}
+
+// GroupKeyReveal generates a GroupKeyReveal instance from the test case.
+func (tc testCaseGkrEncodeDecode) GroupKeyReveal() (GroupKeyReveal, error) {
+	gkr, err := NewGroupKeyRevealV1(
+		tc.internalKey, tc.genesisAssetID, tc.customSubtreeRoot,
+	).Unpack()
+
+	return &gkr, err
+}
+
+func TestGroupKeyTapscriptEncodeDecode(t *testing.T) {
+	t.Parallel()
+
+	// Create a random internal public key.
+	internalKey := *(test.RandPubKey(t))
+
+	// Create a random genesis asset ID.
+	randomAssetIdBytes := test.RandBytes(32)
+	genesisAssetID := ID(randomAssetIdBytes)
+
+	testCases := []testCaseGkrEncodeDecode{
+		{
+			testName: "no custom root",
+
+			internalKey:       internalKey,
+			genesisAssetID:    genesisAssetID,
+			customSubtreeRoot: fn.None[chainhash.Hash](),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.testName, func(tt *testing.T) {
+			gkr, err := tc.GroupKeyReveal()
+			require.NoError(tt, err)
+
+			// Encode the GroupKeyReveal into buffer.
+			var buffer bytes.Buffer
+			var scratchBuffEncode [8]byte
+			err = GroupKeyRevealEncoder(&buffer, &gkr, &scratchBuffEncode)
+			require.NoError(tt, err)
+
+			// Decode the GroupKeyReveal from buffer.
+			var gkrDecoded GroupKeyReveal
+			var scratchBuffDecode [8]byte
+			err = GroupKeyRevealDecoder(
+				&buffer, &gkrDecoded, &scratchBuffDecode, uint64(buffer.Len()),
+			)
+			require.NoError(tt, err)
+
+			// Ensure the decoded GroupKeyReveal matches the original.
+			require.Equal(
+				tt, gkr, gkrDecoded,
+				"decoded GroupKeyReveal does not match original",
+			)
+		})
+	}
+}
+
+func TestGroupKeyTapscript(t *testing.T) {
+	t.Parallel()
+
+	// 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"),
+	)
+
+	// Construct a non-executable script leaf.
+	noExecScript, err := txscript.NewScriptBuilder().
+		AddOp(txscript.OP_RETURN).
+		Script()
+	require.NoError(t, err)
+
+	noExecScriptLeaf := txscript.TapLeaf{
+		Script:      noExecScript,
+		LeafVersion: txscript.BaseLeafVersion,
+	}
+
+	noExecScriptLeafHash := noExecScriptLeaf.TapHash()
+
+	// 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(noExecScriptLeafHash, userNodeHash)
+	actualTapscriptRootHash4 := TapBranchHash(assetIDLeafHash, branchHash)
+	fmt.Printf("Tapscript root hash with optional user tree root: %v\n",
+		actualTapscriptRootHash4)
+
+	// Construct the user subtree control block.
+	internalKey := RandInternalPubKey()
+	outputKey := txscript.ComputeTaprootOutputKey(
+		&internalKey, actualTapscriptRootHash4[:],
+	)
+	outputKeyIsOdd := outputKey.SerializeCompressed()[0] == 0x03
+
+	inclusionProof := bytes.Join(
+		[][]byte{
+			noExecScriptLeafHash[:],
+			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.
+	// 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, actualTapscriptRootHash4, rootCheck)
+}
+
+// RandInternalPubKey generates a random internal public key.
+func RandInternalPubKey() btcec.PublicKey {
+	randBytes := make([]byte, 32)
+	_, _ = rand.Read(randBytes)
+
+	privateKey, _ := btcec.PrivKeyFromBytes(randBytes)
+	pubKey := privateKey.PubKey()
+	return *pubKey
+}