Make max read size configurable (#273)

* Make max read size configurable

* Add init cmd flag

* migrate cmd

* fix test

Author: agouin

cmd/horcrux/cmd/config.go

@@ -18,7 +18,8 @@ const (
 	flagGRPCTimeout = "grpc-timeout"
 	flagOverwrite   = "overwrite"
 	flagBare        = "bare"
-	flagGRPCAddress = "flagGRPCAddress"
+	flagGRPCAddress = "gprc-address"
+	flagMaxReadSize = "max-read-size"
 )
 
 func configCmd() *cobra.Command {
@@ -70,6 +71,7 @@ for threshold signer mode, --cosigner flags and --threshold flag are required.
 			}
 			debugAddr, _ := cmdFlags.GetString(flagDebugAddr)
 			grpcAddr, _ := cmdFlags.GetString(flagGRPCAddress)
+			maxReadSize, _ := cmdFlags.GetInt(flagMaxReadSize)
 			if signMode == string(signer.SignModeThreshold) {
 				// Threshold Mode Config
 				cosignersFlag, _ := cmdFlags.GetStringSlice(flagCosigner)
@@ -90,9 +92,10 @@ for threshold signer mode, --cosigner flags and --threshold flag are required.
 						GRPCTimeout: grpcTimeout,
 						RaftTimeout: raftTimeout,
 					},
-					ChainNodes: cn,
-					DebugAddr:  debugAddr,
-					GRPCAddr:   grpcAddr,
+					ChainNodes:  cn,
+					DebugAddr:   debugAddr,
+					GRPCAddr:    grpcAddr,
+					MaxReadSize: maxReadSize,
 				}
 
 				if !bare {
@@ -107,6 +110,7 @@ for threshold signer mode, --cosigner flags and --threshold flag are required.
 					PrivValKeyDir: keyDir,
 					ChainNodes:    cn,
 					DebugAddr:     debugAddr,
+					MaxReadSize:   maxReadSize,
 				}
 				if !bare {
 					if err = cfg.ValidateSingleSignerConfig(); err != nil {
@@ -162,5 +166,6 @@ for threshold signer mode, --cosigner flags and --threshold flag are required.
 		"allows initialization without providing any flags. If flags are provided, will not perform final validation",
 	)
 	f.StringP(flagGRPCAddress, "g", "", "GRPC address if listener should be enabled")
+	f.Int(flagMaxReadSize, 1024*1024, "max read size for remote signer connection")
 	return cmd
 }

cmd/horcrux/cmd/config_test.go

@@ -48,6 +48,7 @@ chainNodes:
 - privValAddr: tcp://10.168.0.2:1234
 debugAddr: ""
 grpcAddr: ""
+maxReadSize: 1048576
 `,
 		},
 		{
@@ -64,6 +65,7 @@ chainNodes:
 - privValAddr: tcp://10.168.0.2:1234
 debugAddr: ""
 grpcAddr: ""
+maxReadSize: 1048576
 `,
 		},
 		{

cmd/horcrux/cmd/migrate.go

@@ -296,6 +296,7 @@ func migrateCmd() *cobra.Command {
 				}
 
 				config.Config.SignMode = signMode
+				config.Config.MaxReadSize = 1024 * 1024
 
 				if err := config.WriteConfigFile(); err != nil {
 					return err

cmd/horcrux/cmd/start.go

@@ -71,7 +71,7 @@ func startCmd() *cobra.Command {
 
 			go EnableDebugAndMetrics(cmd.Context(), out)
 
-			services, err = signer.StartRemoteSigners(services, logger, val, config.Config.Nodes())
+			services, err = signer.StartRemoteSigners(services, logger, val, config.Config.Nodes(), config.Config.MaxReadSize)
 			if err != nil {
 				return fmt.Errorf("failed to start remote signer(s): %w", err)
 			}

cmd/horcrux/cmd/testdata/config-migrated.yaml

@@ -16,3 +16,4 @@ chainNodes:
 - privValAddr: tcp://127.0.0.1:3456
 debugAddr: ""
 grpcAddr: ""
+maxReadSize: 1048576

signer/config.go

@@ -35,6 +35,7 @@ type Config struct {
 	ChainNodes          ChainNodes           `yaml:"chainNodes"`
 	DebugAddr           string               `yaml:"debugAddr"`
 	GRPCAddr            string               `yaml:"grpcAddr"`
+	MaxReadSize         int                  `yaml:"maxReadSize"`
 }
 
 func (c *Config) Nodes() (out []string) {

signer/config_test.go

@@ -380,6 +380,7 @@ func TestRuntimeConfigWriteConfigFile(t *testing.T) {
 					PrivValAddr: "tcp://127.0.0.1:3456",
 				},
 			},
+			MaxReadSize: 1024 * 1024,
 		},
 	}
 
@@ -404,6 +405,7 @@ chainNodes:
 - privValAddr: tcp://127.0.0.1:3456
 debugAddr: ""
 grpcAddr: ""
+maxReadSize: 1048576
 `, string(configYamlBz))
 }
 

signer/io.go

@@ -8,9 +8,11 @@ import (
 )
 
 // ReadMsg reads a message from an io.Reader
-func ReadMsg(reader io.Reader) (msg cometprotoprivval.Message, err error) {
-	const maxRemoteSignerMsgSize = 1024 * 10
-	protoReader := protoio.NewDelimitedReader(reader, maxRemoteSignerMsgSize)
+func ReadMsg(reader io.Reader, maxReadSize int) (msg cometprotoprivval.Message, err error) {
+	if maxReadSize <= 0 {
+		maxReadSize = 1024 * 1024 // 1MB
+	}
+	protoReader := protoio.NewDelimitedReader(reader, maxReadSize)
 	_, err = protoReader.ReadMsg(&msg)
 	return msg, err
 }

signer/remote_signer.go

@@ -37,6 +37,8 @@ type ReconnRemoteSigner struct {
 	privVal PrivValidator
 
 	dialer net.Dialer
+
+	maxReadSize int
 }
 
 // NewReconnRemoteSigner return a ReconnRemoteSigner that will dial using the given
@@ -49,12 +51,14 @@ func NewReconnRemoteSigner(
 	logger cometlog.Logger,
 	privVal PrivValidator,
 	dialer net.Dialer,
+	maxReadSize int,
 ) *ReconnRemoteSigner {
 	rs := &ReconnRemoteSigner{
-		address: address,
-		privVal: privVal,
-		dialer:  dialer,
-		privKey: cometcryptoed25519.GenPrivKey(),
+		address:     address,
+		privVal:     privVal,
+		dialer:      dialer,
+		privKey:     cometcryptoed25519.GenPrivKey(),
+		maxReadSize: maxReadSize,
 	}
 
 	rs.BaseService = *cometservice.NewBaseService(logger, "RemoteSigner", rs)
@@ -136,7 +140,7 @@ func (rs *ReconnRemoteSigner) loop(ctx context.Context) {
 			return
 		}
 
-		req, err := ReadMsg(conn)
+		req, err := ReadMsg(conn, rs.maxReadSize)
 		if err != nil {
 			rs.Logger.Error(
 				"Failed to read message from connection",
@@ -288,6 +292,7 @@ func StartRemoteSigners(
 	logger cometlog.Logger,
 	privVal PrivValidator,
 	nodes []string,
+	maxReadSize int,
 ) ([]cometservice.Service, error) {
 	var err error
 	go StartMetrics()
@@ -296,7 +301,7 @@ func StartRemoteSigners(
 		// A long timeout such as 30 seconds would cause the sentry to fail in loops
 		// Use a short timeout and dial often to connect within 3 second window
 		dialer := net.Dialer{Timeout: 2 * time.Second}
-		s := NewReconnRemoteSigner(node, logger, privVal, dialer)
+		s := NewReconnRemoteSigner(node, logger, privVal, dialer, maxReadSize)
 
 		err = s.Start()
 		if err != nil {