Merge pull request #1758 from nspcc-dev/cli/0x

cli: support Uint160 hashes with 0x prefix
nspcc-dev · Feb 19, 2021 · c58dc7e · c58dc7e
2 parents 7f6a0fd + b8024db
commit c58dc7e
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 22 deletions.
diff --git a/cli/flags/address.go b/cli/flags/address.go
@@ -35,7 +35,7 @@ func (a Address) String() string {
 
 // Set implements flag.Value interface.
 func (a *Address) Set(s string) error {
-	addr, err := address.StringToUint160(s)
+	addr, err := ParseAddress(s)
 	if err != nil {
 		return cli.NewExitError(err, 1)
 	}
@@ -89,3 +89,14 @@ func (f AddressFlag) Apply(set *flag.FlagSet) {
 		set.Var(&f.Value, name, f.Usage)
 	})
 }
+
+// ParseAddress parses Uint160 form either LE string or address.
+func ParseAddress(s string) (util.Uint160, error) {
+	const uint160size = 2 * util.Uint160Size
+	switch len(s) {
+	case uint160size, uint160size + 2:
+		return util.Uint160DecodeStringLE(strings.TrimPrefix(s, "0x"))
+	default:
+		return address.StringToUint160(s)
+	}
+}
diff --git a/cli/flags/address_test.go b/cli/flags/address_test.go
@@ -5,11 +5,42 @@ import (
 	"io/ioutil"
 	"testing"
 
+	"github.com/nspcc-dev/neo-go/internal/random"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/stretchr/testify/require"
 )
 
+func TestParseAddress(t *testing.T) {
+	expected := random.Uint160()
+	t.Run("simple LE", func(t *testing.T) {
+		u, err := ParseAddress(expected.StringLE())
+		require.NoError(t, err)
+		require.Equal(t, expected, u)
+	})
+	t.Run("with prefix", func(t *testing.T) {
+		u, err := ParseAddress("0x" + expected.StringLE())
+		require.NoError(t, err)
+		require.Equal(t, expected, u)
+
+		t.Run("bad", func(t *testing.T) {
+			_, err := ParseAddress("0s" + expected.StringLE())
+			require.Error(t, err)
+		})
+	})
+	t.Run("address", func(t *testing.T) {
+		addr := address.Uint160ToString(expected)
+		u, err := ParseAddress(addr)
+		require.NoError(t, err)
+		require.Equal(t, expected, u)
+
+		t.Run("bad", func(t *testing.T) {
+			_, err := ParseAddress(addr[1:])
+			require.Error(t, err)
+		})
+	})
+}
+
 func TestAddress_String(t *testing.T) {
 	value := util.Uint160{1, 2, 3}
 	addr := Address{

diff --git a/cli/smartcontract/smart_contract.go b/cli/smartcontract/smart_contract.go
@@ -465,15 +465,9 @@ func contractCompile(ctx *cli.Context) error {
 
 func calcHash(ctx *cli.Context) error {
 	s := ctx.String("sender")
-	u, err := address.StringToUint160(s)
+	u, err := flags.ParseAddress(s)
 	if err != nil {
-		if strings.HasPrefix(s, "0x") {
-			s = s[2:]
-		}
-		u, err = util.Uint160DecodeStringLE(s)
-		if err != nil {
-			return cli.NewExitError(errors.New("invalid sender: must be either address or Uint160 in LE form"), 1)
-		}
+		return cli.NewExitError(errors.New("invalid sender: must be either address or Uint160 in LE form"), 1)
 	}
 
 	p := ctx.String("in")
@@ -530,7 +524,7 @@ func invokeInternal(ctx *cli.Context, signAndPush bool) error {
 	if !args.Present() {
 		return cli.NewExitError(errNoScriptHash, 1)
 	}
-	script, err := util.Uint160DecodeStringLE(args[0])
+	script, err := flags.ParseAddress(args[0])
 	if err != nil {
 		return cli.NewExitError(fmt.Errorf("incorrect script hash: %w", err), 1)
 	}
@@ -891,14 +885,7 @@ func parseCosigner(c string) (transaction.Signer, error) {
 	)
 	data := strings.SplitN(c, ":", 2)
 	s := data[0]
-	if len(s) == 2*util.Uint160Size+2 && s[0:2] == "0x" {
-		s = s[2:]
-	}
-	if len(s) == util.Uint160Size*2 {
-		res.Account, err = util.Uint160DecodeStringLE(s)
-	} else {
-		res.Account, err = address.StringToUint160(s)
-	}
+	res.Account, err = flags.ParseAddress(s)
 	if err != nil {
 		return res, err
 	}

diff --git a/cli/wallet/nep17.go b/cli/wallet/nep17.go
@@ -260,7 +260,7 @@ func importNEP17Token(ctx *cli.Context) error {
 	}
 	defer wall.Close()
 
-	tokenHash, err := util.Uint160DecodeStringLE(ctx.String("token"))
+	tokenHash, err := flags.ParseAddress(ctx.String("token"))
 	if err != nil {
 		return cli.NewExitError(fmt.Errorf("invalid token contract hash: %w", err), 1)
 	}

diff --git a/cli/wallet/wallet.go b/cli/wallet/wallet.go
@@ -15,7 +15,6 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/manifest"
-	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/neo-go/pkg/wallet"
 	"github.com/urfave/cli"
 )
@@ -397,8 +396,8 @@ func importDeployed(ctx *cli.Context) error {
 
 	defer wall.Close()
 
-	rawHash := strings.TrimPrefix(ctx.String("contract"), "0x")
-	h, err := util.Uint160DecodeStringLE(rawHash)
+	rawHash := ctx.String("contract")
+	h, err := flags.ParseAddress(rawHash)
 	if err != nil {
 		return cli.NewExitError(fmt.Errorf("invalid contract hash: %w", err), 1)
 	}