From a47345bd6f0485fb705b2eb95dc592a46803bbbe Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Tue, 28 Feb 2023 19:56:44 +0100
Subject: [PATCH] input: add script size unit and benchmark tests

---
 input/script_utils_test.go | 212 +++++++++++++++++++++++++++++++++++++
 1 file changed, 212 insertions(+)

diff --git a/input/script_utils_test.go b/input/script_utils_test.go
index 41f004a4b..05bc6b5d1 100644
--- a/input/script_utils_test.go
+++ b/input/script_utils_test.go
@@ -2034,3 +2034,215 @@ func TestSpecificationKeyDerivation(t *testing.T) {
 			actualRevocationPrivKeyHex)
 	}
 }
+
+// BenchmarkScriptBuilderAlloc benchmarks the script builder's default
+// allocation.
+func BenchmarkScriptBuilderAlloc(t *testing.B) {
+	dummyData := []byte("dummy data")
+	randomPriv, err := btcec.NewPrivateKey()
+	require.NoError(t, err)
+	randomPub := randomPriv.PubKey()
+	randomPubBytes := randomPub.SerializeCompressed()
+
+	t.ReportAllocs()
+	t.ResetTimer()
+
+	// We run each iteration 1000 times to make the script allocation stand
+	// out a bit more vs. the overhead of the rest of the function calls.
+	for i := 0; i < t.N*1000; i++ {
+		err := runScriptAllocTest(dummyData, randomPubBytes, randomPub)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// TestScriptBuilderAlloc makes sure the scripts generated by the utils have the
+// correct length.
+func TestScriptBuilderAlloc(t *testing.T) {
+	t.Parallel()
+
+	dummyData := []byte("dummy data")
+	randomPriv, err := btcec.NewPrivateKey()
+	require.NoError(t, err)
+	randomPub := randomPriv.PubKey()
+	randomPubBytes := randomPub.SerializeCompressed()
+
+	err = runScriptAllocTest(dummyData, randomPubBytes, randomPub)
+	require.NoError(t, err)
+}
+
+// runScriptAllocTest runs the script size test. We don't use the "require"
+// library to assert the length to not influence the benchmark too much.
+func runScriptAllocTest(dummyData, randomPubBytes []byte,
+	randomPub *btcec.PublicKey) error {
+
+	const (
+		maxCLTV = 500000000
+		maxCSV  = (1 << 31) - 1
+	)
+	script, err := WitnessScriptHash(dummyData)
+	if err != nil {
+		return err
+	}
+	if len(script) != P2WSHSize {
+		return fmt.Errorf("expected script size of %d", P2WSHSize)
+	}
+
+	script, err = WitnessPubKeyHash(randomPubBytes)
+	if err != nil {
+		return err
+	}
+	if len(script) != P2WPKHSize {
+		return fmt.Errorf("expected script size of %d", P2WPKHSize)
+	}
+
+	script, err = GenerateP2SH(dummyData)
+	if err != nil {
+		return err
+	}
+	if len(script) != NestedP2WPKHSize {
+		return fmt.Errorf("expected script size of %d",
+			NestedP2WPKHSize)
+	}
+
+	script, err = GenerateP2PKH(dummyData)
+	if err != nil {
+		return err
+	}
+	if len(script) != P2PKHSize {
+		return fmt.Errorf("expected script size of %d", P2PKHSize)
+	}
+
+	script, err = GenMultiSigScript(randomPubBytes, randomPubBytes)
+	if err != nil {
+		return err
+	}
+	if len(script) != MultiSigSize {
+		return fmt.Errorf("expected script size of %d", MultiSigSize)
+	}
+
+	script, err = SenderHTLCScript(
+		randomPub, randomPub, randomPub, dummyData, false,
+	)
+	if err != nil {
+		return err
+	}
+	if len(script) != OfferedHtlcScriptSize {
+		return fmt.Errorf("expected script size of %d",
+			OfferedHtlcScriptSize)
+	}
+
+	script, err = SenderHTLCScript(
+		randomPub, randomPub, randomPub, dummyData, true,
+	)
+	if err != nil {
+		return err
+	}
+	if len(script) != OfferedHtlcScriptSizeConfirmed {
+		return fmt.Errorf("expected script size of %d",
+			OfferedHtlcScriptSizeConfirmed)
+	}
+
+	script, err = ReceiverHTLCScript(
+		maxCLTV, randomPub, randomPub, randomPub, dummyData, false,
+	)
+	if err != nil {
+		return err
+	}
+	if len(script) != AcceptedHtlcScriptSize {
+		return fmt.Errorf("expected script size of %d",
+			AcceptedHtlcScriptSize)
+	}
+
+	script, err = ReceiverHTLCScript(
+		maxCLTV, randomPub, randomPub, randomPub, dummyData, true,
+	)
+	if err != nil {
+		return err
+	}
+	if len(script) != AcceptedHtlcScriptSizeConfirmed {
+		return fmt.Errorf("expected script size of %d",
+			AcceptedHtlcScriptSizeConfirmed)
+	}
+
+	script, err = SecondLevelHtlcScript(randomPub, randomPub, maxCSV)
+	if err != nil {
+		return err
+	}
+	if len(script) != ToLocalScriptSize {
+		return fmt.Errorf("expected script size of %d",
+			ToLocalScriptSize)
+	}
+
+	script, err = LeaseSecondLevelHtlcScript(
+		randomPub, randomPub, maxCSV, maxCLTV,
+	)
+	if err != nil {
+		return err
+	}
+	if len(script) != ToLocalScriptSize+LeaseWitnessScriptSizeOverhead {
+		return fmt.Errorf("expected script size of %d",
+			ToLocalScriptSize+LeaseWitnessScriptSizeOverhead)
+	}
+
+	script, err = CommitScriptToSelf(maxCSV, randomPub, randomPub)
+	if err != nil {
+		return err
+	}
+	if len(script) != ToLocalScriptSize {
+		return fmt.Errorf("expected script size of %d",
+			ToLocalScriptSize)
+	}
+
+	script, err = LeaseCommitScriptToSelf(
+		randomPub, randomPub, maxCSV, maxCLTV,
+	)
+	if err != nil {
+		return err
+	}
+	if len(script) != ToLocalScriptSize+LeaseWitnessScriptSizeOverhead {
+		return fmt.Errorf("expected script size of %d",
+			ToLocalScriptSize+LeaseWitnessScriptSizeOverhead)
+	}
+
+	script, err = CommitScriptUnencumbered(randomPub)
+	if err != nil {
+		return err
+	}
+	if len(script) != P2WPKHSize {
+		return fmt.Errorf("expected script size of %d", P2WPKHSize)
+	}
+
+	script, err = CommitScriptToRemoteConfirmed(randomPub)
+	if err != nil {
+		return err
+	}
+	if len(script) != ToRemoteConfirmedScriptSize {
+		return fmt.Errorf("expected script size of %d",
+			ToRemoteConfirmedScriptSize)
+	}
+
+	script, err = LeaseCommitScriptToRemoteConfirmed(randomPub, maxCSV)
+	if err != nil {
+		return err
+	}
+	if len(script) !=
+		ToRemoteConfirmedScriptSize+LeaseWitnessScriptSizeOverhead {
+
+		return fmt.Errorf("expected script size of %d",
+			ToRemoteConfirmedScriptSize+
+				LeaseWitnessScriptSizeOverhead)
+	}
+
+	script, err = CommitScriptAnchor(randomPub)
+	if err != nil {
+		return err
+	}
+	if len(script) != AnchorScriptSize {
+		return fmt.Errorf("expected script size of %d",
+			AnchorScriptSize)
+	}
+
+	return nil
+}