rpcserver: simpler gRPC generators with rapid fork

lightninglabs · Aug 16, 2024 · 8b579da · 8b579da
1 parent 3c85836
commit 8b579da
Show file tree

Hide file tree

Showing 5 changed files with 101 additions and 133 deletions.
diff --git a/go.mod b/go.mod
@@ -208,3 +208,5 @@ require (
 // We want to format raw bytes as hex instead of base64. The forked version
 // allows us to specify that as an option.
 replace google.golang.org/protobuf => github.com/lightninglabs/protobuf-go-hex-display v1.30.0-hex-display
+
+replace pgregory.net/rapid v1.1.0 => github.com/chrisseto/rapid v0.0.0-20240815210052-cdeef406c65c // indirect
diff --git a/go.sum b/go.sum
@@ -124,6 +124,8 @@ github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XL
 github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
 github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/chrisseto/rapid v0.0.0-20240815210052-cdeef406c65c h1:GZtcJAFTBCr16eM7ytFwWMg9oLaMsRfSsVyi3lTo+mw=
+github.com/chrisseto/rapid v0.0.0-20240815210052-cdeef406c65c/go.mod h1:PY5XlDGj0+V1FCq0o192FdRhpKHGTRIWBgqjDBTrq04=
 github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
 github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
 github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=
@@ -1169,8 +1171,6 @@ modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
 modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
 nhooyr.io/websocket v1.8.7 h1:usjR2uOr/zjjkVMy0lW+PPohFok7PCow5sDjLgX4P4g=
 nhooyr.io/websocket v1.8.7/go.mod h1:B70DZP8IakI65RVQ51MsWP/8jndNma26DVA/nFSCgW0=
-pgregory.net/rapid v1.1.0 h1:CMa0sjHSru3puNx+J0MIAuiiEV4N0qj8/cMWGBBCsjw=
-pgregory.net/rapid v1.1.0/go.mod h1:PY5XlDGj0+V1FCq0o192FdRhpKHGTRIWBgqjDBTrq04=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=
 rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=

diff --git a/rpcserver_test.go b/rpcserver_test.go
@@ -5,92 +5,117 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
+	"maps"
+	"reflect"
 	"testing"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
 	"github.com/lightninglabs/taproot-assets/taprpc"
 	"github.com/lightninglabs/taproot-assets/taprpc/universerpc"
 	"github.com/lightninglabs/taproot-assets/universe"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	"pgregory.net/rapid"
 )
 
+type rapidFieldGen = map[string]*rapid.Generator[any]
+type rapidFieldMap = map[reflect.Type]rapidFieldGen
+type rapidTypeMap = map[reflect.Type]*rapid.Generator[any]
+
 // Custom generators.
 var (
-	ByteSliceGen   = rapid.SliceOf(rapid.Byte())
-	GenesisInfoGen = rapid.Custom(func(t *rapid.T) taprpc.GenesisInfo {
-		return taprpc.GenesisInfo{
-			GenesisPoint: rapid.String().Draw(t, "genesis_point"),
-			Name:         rapid.String().Draw(t, "name"),
-			MetaHash:     ByteSliceGen.Draw(t, "meta_hash"),
-			AssetId:      ByteSliceGen.Draw(t, "id"),
-			AssetType: taprpc.AssetType(
-				rapid.Int32().Draw(t, "asset_type"),
-			),
-			OutputIndex: rapid.Uint32().Draw(t, "output_index"),
-		}
-	})
-	AssetGroupGen = rapid.Custom(func(t *rapid.T) taprpc.AssetGroup {
-		return taprpc.AssetGroup{
-			RawGroupKey: ByteSliceGen.Draw(t, "raw_group_key"),
-			TweakedGroupKey: ByteSliceGen.Draw(
-				t, "tweaked_group_key",
-			),
-			AssetWitness:  ByteSliceGen.Draw(t, "asset_witness"),
-			TapscriptRoot: ByteSliceGen.Draw(t, "tapscript_root"),
+	ByteSliceGen = rapid.SliceOf(rapid.Byte())
+
+	// Ignore private gRPC fields of messages, which we don't read when
+	// unmarshalling and cause issues with rapid.Make().
+	ignorePrivateRPCFields = rapidFieldGen{
+		"state":         rapid.Just(protoimpl.MessageState{}).AsAny(),
+		"sizeCache":     rapid.Just(protoimpl.SizeCache(0)).AsAny(),
+		"unknownFields": rapid.Just(protoimpl.UnknownFields{}).AsAny(),
+	}
+
+	// Create a generator config for a gRPC message, which may include
+	// custom generators or type generator overrides.
+	genMakeConfig = func(rpcType any, customGens rapidFieldGen,
+		genOverrides rapidTypeMap) rapid.MakeConfig {
+
+		cfg := rapid.MakeConfig{
+			Types:  make(rapidTypeMap),
+			Fields: make(rapidFieldMap),
 		}
-	})
-	AnchorInfoGen = rapid.Custom(func(t *rapid.T) taprpc.AnchorInfo {
-		return taprpc.AnchorInfo{
-			AnchorTx: ByteSliceGen.Draw(t, "anchor_tx"),
-			AnchorBlockHash: rapid.String().Draw(
-				t, "anchor_block_hash",
-			),
-			AnchorOutpoint: rapid.String().Draw(
-				t, "anchor_outpoint",
-			),
-			InternalKey: ByteSliceGen.Draw(t, "internal_key"),
-			MerkleRoot:  ByteSliceGen.Draw(t, "merkle_root"),
-			TapscriptSibling: ByteSliceGen.Draw(
-				t, "tapscript_sibling",
-			),
-			BlockHeight: rapid.Uint32().Draw(t, "block_height"),
+
+		// Add custom generators for fields, by field name, to override
+		// default rapid.Make() behavior.
+		ignoredFields := maps.Clone(ignorePrivateRPCFields)
+		for k, v := range customGens {
+			ignoredFields[k] = v
 		}
-	})
-	PrevInputAssetGen = rapid.Custom(
-		func(t *rapid.T) taprpc.PrevInputAsset {
-
-			return taprpc.PrevInputAsset{
-				AnchorPoint: rapid.String().Draw(
-					t, "anchor_point",
-				),
-				AssetId:   ByteSliceGen.Draw(t, "asset_id"),
-				ScriptKey: ByteSliceGen.Draw(t, "script_key"),
-				Amount:    rapid.Uint64().Draw(t, "amount"),
-			}
-		})
-	PrevWitnessGen = rapid.Custom(func(t *rapid.T) taprpc.PrevWitness {
-		// Leave the split commitment as nil.
-		return taprpc.PrevWitness{
-			PrevId: rapid.Ptr(PrevInputAssetGen, true).Draw(
-				t, "prev_id",
-			),
-			TxWitness: rapid.SliceOf(ByteSliceGen).Draw(
-				t, "tx_witnesses",
-			),
+
+		cfg.Fields[reflect.TypeOf(rpcType)] = ignoredFields
+
+		// Add custom generators that will override the generators
+		// rapid.Make() would create for struct member types.
+		for k, v := range genOverrides {
+			cfg.Types[k] = v
 		}
-	})
+
+		return cfg
+	}
+
+	GenesisInfoGen = rapid.Ptr(rapid.MakeCustom[taprpc.GenesisInfo](
+		genMakeConfig(taprpc.GenesisInfo{}, nil, nil),
+	), true)
+	AssetGroupGen = rapid.Ptr(rapid.MakeCustom[taprpc.AssetGroup](
+		genMakeConfig(taprpc.AssetGroup{}, nil, nil),
+	), true)
+	AnchorInfoGen = rapid.Ptr(rapid.MakeCustom[taprpc.AnchorInfo](
+		genMakeConfig(taprpc.AnchorInfo{}, nil, nil),
+	), true)
+	PrevInputAssetGen = rapid.MakeCustom[taprpc.PrevInputAsset](
+		genMakeConfig(taprpc.PrevInputAsset{}, nil, nil),
+	)
+
+	// Leave the split commitment for prev witnesses as nil.
+	emptySplitCommitmentGen = rapid.Just(taprpc.SplitCommitment{})
+	splitCommitPtrGen       = rapid.Ptr(emptySplitCommitmentGen, true)
+	nilSplitCommitment      = rapidTypeMap{
+		//nolint:lll
+		reflect.TypeOf(&taprpc.SplitCommitment{}): splitCommitPtrGen.AsAny(),
+	}
+	PrevWitnessGen = rapid.MakeCustom[taprpc.PrevWitness](
+		genMakeConfig(taprpc.PrevWitness{}, nil, nilSplitCommitment),
+	)
 	PrevWitnessesGen = rapid.Custom(func(t *rapid.T) []*taprpc.PrevWitness {
 		witnessGen := rapid.Ptr(PrevWitnessGen, false)
 		return rapid.SliceOf(witnessGen).Draw(t, "prev_witnesses")
 	})
-	DecDisplayGen = rapid.Custom(func(t *rapid.T) taprpc.DecimalDisplay {
-		return taprpc.DecimalDisplay{
-			DecimalDisplay: rapid.Uint32().Draw(
-				t, "decimal_display",
-			),
-		}
-	})
+	DecDisplayGen = rapid.Ptr(rapid.MakeCustom[taprpc.DecimalDisplay](
+		genMakeConfig(taprpc.DecimalDisplay{}, nil, nil),
+	), true)
+
+	// Set generator overrides for members of taprpc.Asset that are gRPC
+	// messages.
+	assetMemberGens = rapidTypeMap{
+		reflect.TypeOf(&taprpc.GenesisInfo{}): GenesisInfoGen.AsAny(),
+		reflect.TypeOf(&taprpc.AssetGroup{}):  AssetGroupGen.AsAny(),
+		reflect.TypeOf(&taprpc.AnchorInfo{}):  AnchorInfoGen.AsAny(),
+		//nolint:lll
+		reflect.TypeOf([]*taprpc.PrevWitness{}):  PrevWitnessesGen.AsAny(),
+		reflect.TypeOf(&taprpc.DecimalDisplay{}): DecDisplayGen.AsAny(),
+	}
+	AssetGen = rapid.MakeCustom[taprpc.Asset](
+		genMakeConfig(taprpc.Asset{}, nil, assetMemberGens),
+	)
+	AssetPtrGen = rapid.Ptr(AssetGen, true)
+
+	// Use the custom taprpc.Asset generator for *universerpc.AssetLeaf.
+	leafMemberGens = rapidTypeMap{
+		reflect.TypeOf(&taprpc.Asset{}): AssetPtrGen.AsAny(),
+	}
+	AssetLeafGen = rapid.MakeCustom[universerpc.AssetLeaf](
+		genMakeConfig(universerpc.AssetLeaf{}, nil, leafMemberGens),
+	)
+	AssetLeafPtrGen = rapid.Ptr(AssetLeafGen, true)
 )
 
 // Result is used to store the output of a fallible function call.
@@ -530,68 +555,9 @@ func TestUnmarshalUniId(t *testing.T) {
 }
 
 func testUnmarshalAssetLeaf(t *rapid.T) {
-	// rapid.Make failed on the private gRPC-specific fields of
-	// taprpc.Asset, so we'll populate only the public fields.
-	LeafAssetGen := rapid.Custom(func(t *rapid.T) taprpc.Asset {
-		vers := taprpc.AssetVersion(rapid.Int32().Draw(t, "version"))
-		genesis := rapid.Ptr(GenesisInfoGen, true).Draw(t, "genesis")
-		amount := rapid.Uint64().Draw(t, "amount")
-		lockTime := rapid.Int32().Draw(t, "lock_time")
-		relativeLockTime := rapid.Int32().Draw(t, "relative_lock_time")
-		scriptVersion := rapid.Int32().Draw(t, "script_version")
-		scriptKey := ByteSliceGen.Draw(t, "script_key")
-		scriptKeyIsLocal := rapid.Bool().Draw(t, "script_key_is_local")
-		group := rapid.Ptr(AssetGroupGen, true).Draw(t, "asset_group")
-		chainAnchor := rapid.Ptr(AnchorInfoGen, true).Draw(
-			t, "chain_anchor",
-		)
-		prevWitnesses := PrevWitnessesGen.Draw(t, "prev_witnesses")
-		isSpent := rapid.Bool().Draw(t, "is_spent")
-		leaseOwner := ByteSliceGen.Draw(t, "lease_owner")
-		leaseExpiry := rapid.Int64().Draw(t, "lease_expiry")
-		isBurn := rapid.Bool().Draw(t, "is_burn")
-		scriptKeyDeclaredKnown := rapid.Bool().Draw(
-			t, "script_key_declared_known",
-		)
-		scriptKeyHasScriptPath := rapid.Bool().Draw(
-			t, "script_key_has_script_path",
-		)
-		decimalDisplay := rapid.Ptr(DecDisplayGen, true).Draw(
-			t, "decimal_display",
-		)
-
-		return taprpc.Asset{
-			Version:                vers,
-			AssetGenesis:           genesis,
-			Amount:                 amount,
-			LockTime:               lockTime,
-			RelativeLockTime:       relativeLockTime,
-			ScriptVersion:          scriptVersion,
-			ScriptKey:              scriptKey,
-			ScriptKeyIsLocal:       scriptKeyIsLocal,
-			AssetGroup:             group,
-			ChainAnchor:            chainAnchor,
-			PrevWitnesses:          prevWitnesses,
-			IsSpent:                isSpent,
-			LeaseOwner:             leaseOwner,
-			LeaseExpiry:            leaseExpiry,
-			IsBurn:                 isBurn,
-			ScriptKeyDeclaredKnown: scriptKeyDeclaredKnown,
-			ScriptKeyHasScriptPath: scriptKeyHasScriptPath,
-			DecimalDisplay:         decimalDisplay,
-		}
-	})
-
-	leafGen := rapid.Custom(func(t *rapid.T) universerpc.AssetLeaf {
-		return universerpc.AssetLeaf{
-			Asset: rapid.Ptr(LeafAssetGen, true).Draw(t, "Asset"),
-			Proof: ByteSliceGen.Draw(t, "Proof"),
-		}
-	})
-	leaf := rapid.Ptr(leafGen, true).Draw(t, "Leaf")
-
 	// Don't check the unmarshal output, we are only testing if we can
 	// cause unmarshal to panic.
+	leaf := AssetLeafPtrGen.Draw(t, "Leaf")
 	_, _ = unmarshalAssetLeaf(leaf)
 }
 

diff --git a/testdata/rapid/TestUnmarshalAssetLeaf/TestUnmarshalAssetLeaf-20240815223752-3701568.fail b/testdata/rapid/TestUnmarshalAssetLeaf/TestUnmarshalAssetLeaf-20240815223752-3701568.fail
diff --git a/testdata/rapid/TestUnmarshalAssetLeaf/TestUnmarshalAssetLeaf-20240816174811-4026466.fail b/testdata/rapid/TestUnmarshalAssetLeaf/TestUnmarshalAssetLeaf-20240816174811-4026466.fail
@@ -0,0 +1,4 @@
+# 2024/08/16 17:48:11.689274 [TestUnmarshalAssetLeaf] [rapid] draw Leaf: (*universerpc.AssetLeaf)(nil)
+# 
+v0.4.8#7630617197267023936
+0x0