Add custom msg type registry map to client config, pull into codec du…

…ring client creation
bryanlabs · Sep 14, 2024 · f28ea36 · f28ea36
1 parent ffbbb18
commit f28ea36
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 13 deletions.
diff --git a/client/client.go b/client/client.go
@@ -27,17 +27,26 @@ type ChainClient struct {
 
 func NewChainClient(ccc *ChainClientConfig, homepath string, input io.Reader, output io.Writer, kro ...keyring.Option) (*ChainClient, error) {
 	ccc.KeyDirectory = keysDir(homepath, ccc.ChainID)
+
+	codec, err := MakeCodec(ccc.Modules, ccc.CustomMsgTypeRegistry)
+
+	if err != nil {
+		return nil, err
+	}
+
 	cc := &ChainClient{
 		KeyringOptions: kro,
 		Config:         ccc,
 		Input:          input,
 		Output:         output,
-		Codec:          MakeCodec(ccc.Modules),
+		Codec:          codec,
 		Logger:         log.NewTMLogger(log.NewSyncWriter(output)),
 	}
+
 	if err := cc.Init(); err != nil {
 		return nil, err
 	}
+
 	return cc, nil
 }
 

diff --git a/client/codec/types/interface_registry.go b/client/codec/types/interface_registry.go
@@ -77,6 +77,11 @@ func (registry *ProbeInterfaceRegistry) RegisterCustomTypeURL(iface interface{},
 	registry.registerImpl(iface, typeURL, impl)
 }
 
+func (registry *ProbeInterfaceRegistry) TypeURLIsRegistered(typeURL string) bool {
+	_, found := registry.typeURLMap[typeURL]
+	return found
+}
+
 // registerImpl registers a concrete type which implements the given
 // interface under `typeURL`.
 //

diff --git a/client/config.go b/client/config.go
@@ -4,6 +4,7 @@ package client
 
 import (
 	"github.com/CosmWasm/wasmd/x/wasm"
+	sdkTypes "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/x/auth"
 	"github.com/cosmos/cosmos-sdk/x/auth/vesting"
@@ -48,14 +49,15 @@ var (
 )
 
 type ChainClientConfig struct {
-	Key            string                  `json:"key" yaml:"key"`
-	ChainID        string                  `json:"chain-id" yaml:"chain-id"`
-	RPCAddr        string                  `json:"rpc-addr" yaml:"rpc-addr"`
-	AccountPrefix  string                  `json:"account-prefix" yaml:"account-prefix"`
-	KeyringBackend string                  `json:"keyring-backend" yaml:"keyring-backend"`
-	KeyDirectory   string                  `json:"key-directory" yaml:"key-directory"`
-	Debug          bool                    `json:"debug" yaml:"debug"`
-	Timeout        string                  `json:"timeout" yaml:"timeout"`
-	OutputFormat   string                  `json:"output-format" yaml:"output-format"`
-	Modules        []module.AppModuleBasic `json:"-" yaml:"-"`
+	Key                   string                  `json:"key" yaml:"key"`
+	ChainID               string                  `json:"chain-id" yaml:"chain-id"`
+	RPCAddr               string                  `json:"rpc-addr" yaml:"rpc-addr"`
+	AccountPrefix         string                  `json:"account-prefix" yaml:"account-prefix"`
+	KeyringBackend        string                  `json:"keyring-backend" yaml:"keyring-backend"`
+	KeyDirectory          string                  `json:"key-directory" yaml:"key-directory"`
+	Debug                 bool                    `json:"debug" yaml:"debug"`
+	Timeout               string                  `json:"timeout" yaml:"timeout"`
+	OutputFormat          string                  `json:"output-format" yaml:"output-format"`
+	Modules               []module.AppModuleBasic `json:"-" yaml:"-"`
+	CustomMsgTypeRegistry map[string]sdkTypes.Msg `json:"-" yaml:"-"`
 }
diff --git a/client/encoding.go b/client/encoding.go
@@ -3,6 +3,8 @@
 package client
 
 import (
+	"fmt"
+
 	osmosisGammTypes "github.com/DefiantLabs/probe/client/codec/osmosis/v15/x/gamm/types"
 	osmosisLockupTypes "github.com/DefiantLabs/probe/client/codec/osmosis/v15/x/lockup/types"
 	osmosisPoolManagerTypes "github.com/DefiantLabs/probe/client/codec/osmosis/v15/x/poolmanager/types"
@@ -12,6 +14,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/codec/types"
 	"github.com/cosmos/cosmos-sdk/std"
+	sdkTypes "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/x/auth/tx"
 )
@@ -24,15 +27,22 @@ type Codec struct {
 	Amino                  *codec.LegacyAmino
 }
 
-func MakeCodec(moduleBasics []module.AppModuleBasic) Codec {
+func MakeCodec(moduleBasics []module.AppModuleBasic, customMsgTypeRegistry map[string]sdkTypes.Msg) (Codec, error) {
 	modBasic := module.NewBasicManager(moduleBasics...)
 	encodingConfig := MakeCodecConfig()
 	std.RegisterLegacyAminoCodec(encodingConfig.Amino)
 	std.RegisterInterfaces(encodingConfig.InterfaceRegistry)
 	modBasic.RegisterLegacyAminoCodec(encodingConfig.Amino)
 	modBasic.RegisterInterfaces(encodingConfig.InterfaceRegistry)
 
-	return encodingConfig
+	for typeURL, msg := range customMsgTypeRegistry {
+		if encodingConfig.ProbeInterfaceRegistry.TypeURLIsRegistered(typeURL) {
+			return Codec{}, fmt.Errorf("error registering custom message type in codec, typeURL %s is already registered", typeURL)
+		}
+		encodingConfig.ProbeInterfaceRegistry.RegisterCustomTypeURL((*sdkTypes.Msg)(nil), typeURL, msg)
+	}
+
+	return encodingConfig, nil
 }
 
 // Split out from base codec to not include explicitly.