feat(wallet): add timeout client connection (#1396)

cmd/cmd.go

@@ -388,7 +388,8 @@ func StartNode(workingDir string, passwordFetcher func(*wallet.Wallet) (string,
 	}
 
 	defaultWalletPath := PactusDefaultWalletPath(workingDir)
-	walletInstance, err := wallet.Open(defaultWalletPath, true)
+	walletInstance, err := wallet.Open(defaultWalletPath, true,
+		wallet.WithCustomServers([]string{conf.GRPC.Listen}))
 	if err != nil {
 		return nil, nil, err
 	}

cmd/gtk/main.go

@@ -170,8 +170,6 @@ func run(n *node.Node, wlt *wallet.Wallet, app *gtk.Application) {
 	grpcAddr := n.GRPC().Address()
 	cmd.PrintInfoMsgf("connect wallet to grpc server: %s\n", grpcAddr)
 
-	wlt.SetServerAddr(grpcAddr)
-
 	nodeModel := newNodeModel(n)
 	walletModel := newWalletModel(wlt, n)
 

cmd/wallet/main.go

@@ -1,15 +1,18 @@
 package main
 
 import (
+	"time"
+
 	"github.com/pactus-project/pactus/cmd"
 	"github.com/pactus-project/pactus/wallet"
 	"github.com/spf13/cobra"
 )
 
 var (
-	pathOpt       *string
-	offlineOpt    *bool
-	serverAddrOpt *string
+	pathOpt        *string
+	offlineOpt     *bool
+	serverAddrsOpt *[]string
+	timeoutOpt     *int
 )
 
 func addPasswordOption(c *cobra.Command) *string {
@@ -18,13 +21,19 @@ func addPasswordOption(c *cobra.Command) *string {
 }
 
 func openWallet() (*wallet.Wallet, error) {
-	wlt, err := wallet.Open(*pathOpt, *offlineOpt)
-	if err != nil {
-		return nil, err
+	opts := make([]wallet.Option, 0)
+
+	if *serverAddrsOpt != nil {
+		opts = append(opts, wallet.WithCustomServers(*serverAddrsOpt))
 	}
 
-	if *serverAddrOpt != "" {
-		wlt.SetServerAddr(*serverAddrOpt)
+	if *timeoutOpt > 0 {
+		opts = append(opts, wallet.WithTimeout(time.Duration(*timeoutOpt)*time.Second))
+	}
+
+	wlt, err := wallet.Open(*pathOpt, *offlineOpt, opts...)
+	if err != nil {
+		return nil, err
 	}
 
 	return wlt, err
@@ -43,7 +52,9 @@ func main() {
 	pathOpt = rootCmd.PersistentFlags().String("path",
 		cmd.PactusDefaultWalletPath(cmd.PactusDefaultHomeDir()), "the path to the wallet file")
 	offlineOpt = rootCmd.PersistentFlags().Bool("offline", false, "offline mode")
-	serverAddrOpt = rootCmd.PersistentFlags().String("server", "", "server gRPC address")
+	serverAddrsOpt = rootCmd.PersistentFlags().StringSlice("servers", []string{}, "servers gRPC address")
+	timeoutOpt = rootCmd.PersistentFlags().Int("timeout", 1,
+		"specifies the timeout duration for the client connection in seconds")
 
 	buildCreateCmd(rootCmd)
 	buildRecoverCmd(rootCmd)

wallet/client.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"encoding/hex"
 	"errors"
+	"net"
+	"time"
 
 	"github.com/pactus-project/pactus/crypto/hash"
 	"github.com/pactus-project/pactus/types/amount"
@@ -20,23 +22,27 @@ type grpcClient struct {
 	ctx               context.Context
 	servers           []string
 	conn              *grpc.ClientConn
+	timeout           time.Duration
 	blockchainClient  pactus.BlockchainClient
 	transactionClient pactus.TransactionClient
 }
 
-func newGrpcClient() *grpcClient {
-	ctx := context.WithoutCancel(context.Background())
+func newGrpcClient(timeout time.Duration, servers []string) *grpcClient {
+	ctx := context.Background()
 
-	return &grpcClient{
+	cli := &grpcClient{
 		ctx:               ctx,
+		timeout:           timeout,
 		conn:              nil,
 		blockchainClient:  nil,
 		transactionClient: nil,
 	}
-}
 
-func (c *grpcClient) SetServerAddrs(servers []string) {
-	c.servers = servers
+	if len(servers) > 0 {
+		cli.servers = servers
+	}
+
+	return cli
 }
 
 func (c *grpcClient) connect() error {
@@ -46,7 +52,10 @@ func (c *grpcClient) connect() error {
 
 	for _, server := range c.servers {
 		conn, err := grpc.NewClient(server,
-			grpc.WithTransportCredentials(insecure.NewCredentials()))
+			grpc.WithTransportCredentials(insecure.NewCredentials()),
+			grpc.WithContextDialer(func(_ context.Context, s string) (net.Conn, error) {
+				return net.DialTimeout("tcp", s, c.timeout)
+			}))
 		if err != nil {
 			continue
 		}
@@ -58,6 +67,8 @@ func (c *grpcClient) connect() error {
 		_, err = blockchainClient.GetBlockchainInfo(c.ctx,
 			&pactus.GetBlockchainInfoRequest{})
 		if err != nil {
+			_ = conn.Close()
+
 			continue
 		}
 

wallet/manager.go

@@ -84,11 +84,10 @@ func (wm *Manager) LoadWallet(walletName, serverAddr string) error {
 	}
 
 	walletPath := util.MakeAbs(filepath.Join(wm.walletDirectory, walletName))
-	wlt, err := Open(walletPath, true)
+	wlt, err := Open(walletPath, true, WithCustomServers([]string{serverAddr}))
 	if err != nil {
 		return err
 	}
-	wlt.SetServerAddr(serverAddr)
 
 	wm.wallets[walletName] = wlt
 

wallet/options.go

@@ -0,0 +1,27 @@
+package wallet
+
+import "time"
+
+type walletOpt struct {
+	timeout time.Duration
+	servers []string
+}
+
+type Option func(*walletOpt)
+
+var defaultWalletOpt = &walletOpt{
+	timeout: 1 * time.Second,
+	servers: make([]string, 0),
+}
+
+func WithTimeout(timeout time.Duration) Option {
+	return func(opt *walletOpt) {
+		opt.timeout = timeout
+	}
+}
+
+func WithCustomServers(servers []string) Option {
+	return func(opt *walletOpt) {
+		opt.servers = servers
+	}
+}

wallet/wallet.go

@@ -47,7 +47,7 @@ func CheckMnemonic(mnemonic string) error {
 // A wallet can be opened in offline or online modes.
 // Offline wallet doesn’t have any connection to any node.
 // Online wallet has a connection to one of the pre-defined servers.
-func Open(walletPath string, offline bool) (*Wallet, error) {
+func Open(walletPath string, offline bool, options ...Option) (*Wallet, error) {
 	data, err := util.ReadFile(walletPath)
 	if err != nil {
 		return nil, err
@@ -59,12 +59,24 @@ func Open(walletPath string, offline bool) (*Wallet, error) {
 		return nil, err
 	}
 
-	return newWallet(walletPath, store, offline)
+	opts := defaultWalletOpt
+
+	for _, opt := range options {
+		opt(opts)
+	}
+
+	return newWallet(walletPath, store, offline, opts)
 }
 
 // Create creates a wallet from mnemonic (seed phrase) and save it at the
 // given path.
-func Create(walletPath, mnemonic, password string, chain genesis.ChainType) (*Wallet, error) {
+func Create(walletPath, mnemonic, password string, chain genesis.ChainType, options ...Option) (*Wallet, error) {
+	opts := defaultWalletOpt
+
+	for _, opt := range options {
+		opt(opts)
+	}
+
 	walletPath = util.MakeAbs(walletPath)
 	if util.PathExists(walletPath) {
 		return nil, ExitsError{
@@ -89,7 +101,7 @@ func Create(walletPath, mnemonic, password string, chain genesis.ChainType) (*Wa
 		Network:   chain,
 		Vault:     nil,
 	}
-	wallet, err := newWallet(walletPath, store, true)
+	wallet, err := newWallet(walletPath, store, true, opts)
 	if err != nil {
 		return nil, err
 	}
@@ -106,7 +118,7 @@ func Create(walletPath, mnemonic, password string, chain genesis.ChainType) (*Wa
 	return wallet, nil
 }
 
-func newWallet(walletPath string, store *store, offline bool) (*Wallet, error) {
+func newWallet(walletPath string, store *store, offline bool, option *walletOpt) (*Wallet, error) {
 	if !store.Network.IsMainnet() {
 		crypto.AddressHRP = "tpc"
 		crypto.PublicKeyHRP = "tpublic"
@@ -115,7 +127,7 @@ func newWallet(walletPath string, store *store, offline bool) (*Wallet, error) {
 		crypto.XPrivateKeyHRP = "txsecret"
 	}
 
-	client := newGrpcClient()
+	client := newGrpcClient(option.timeout, option.servers)
 
 	w := &Wallet{
 		store:      store,
@@ -146,16 +158,15 @@ func newWallet(walletPath string, store *store, offline bool) (*Wallet, error) {
 		}
 
 		util.Shuffle(netServers)
-		w.grpcClient.SetServerAddrs(netServers)
+
+		if client.servers == nil {
+			client.servers = netServers
+		}
 	}
 
 	return w, nil
 }
 
-func (w *Wallet) SetServerAddr(addr string) {
-	w.grpcClient.SetServerAddrs([]string{addr})
-}
-
 func (w *Wallet) Name() string {
 	return path.Base(w.path)
 }

wallet/wallet_test.go

@@ -34,11 +34,6 @@ func setup(t *testing.T) *testData {
 	password := ""
 	walletPath := util.TempFilePath()
 	mnemonic, _ := wallet.GenerateMnemonic(128)
-	wlt, err := wallet.Create(walletPath, mnemonic, password, genesis.Mainnet)
-	assert.NoError(t, err)
-	assert.False(t, wlt.IsEncrypted())
-	assert.Equal(t, wlt.Path(), walletPath)
-	assert.Equal(t, wlt.Name(), path.Base(walletPath))
 
 	grpcConf := &grpc.Config{
 		Enable: true,
@@ -57,7 +52,12 @@ func setup(t *testing.T) *testData {
 
 	assert.NoError(t, gRPCServer.StartServer())
 
-	wlt.SetServerAddr(gRPCServer.Address())
+	wlt, err := wallet.Create(walletPath, mnemonic, password, genesis.Mainnet,
+		wallet.WithCustomServers([]string{gRPCServer.Address()}))
+	assert.NoError(t, err)
+	assert.False(t, wlt.IsEncrypted())
+	assert.Equal(t, wlt.Path(), walletPath)
+	assert.Equal(t, wlt.Name(), path.Base(walletPath))
 
 	return &testData{
 		TestSuite: ts,