diff --git a/cmd/cmd.go b/cmd/cmd.go index 25123cce1..82124b5cf 100644 --- a/cmd/cmd.go +++ b/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 } diff --git a/cmd/gtk/main.go b/cmd/gtk/main.go index 216ab874f..0f522c207 100644 --- a/cmd/gtk/main.go +++ b/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) diff --git a/cmd/wallet/main.go b/cmd/wallet/main.go index 1282cc6ef..6da00a9fe 100644 --- a/cmd/wallet/main.go +++ b/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) diff --git a/wallet/client.go b/wallet/client.go index 6864a35a8..9ec400414 100644 --- a/wallet/client.go +++ b/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 } diff --git a/wallet/manager.go b/wallet/manager.go index 5bdcebb39..a96f5f8ab 100644 --- a/wallet/manager.go +++ b/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 diff --git a/wallet/options.go b/wallet/options.go new file mode 100644 index 000000000..87e3614f9 --- /dev/null +++ b/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 + } +} diff --git a/wallet/wallet.go b/wallet/wallet.go index 3f0bf528a..0ae20e2a0 100644 --- a/wallet/wallet.go +++ b/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) } diff --git a/wallet/wallet_test.go b/wallet/wallet_test.go index ce3715fef..7cb629a54 100644 --- a/wallet/wallet_test.go +++ b/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,