From b8024dbfa8fa03b40be12562d8d05bde87405f3f Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 17 Feb 2021 17:48:41 +0300
Subject: [PATCH] cli: support Uint160 hashes with 0x prefix

---
 cli/flags/address_test.go           | 31 +++++++++++++++++++++++++++++
 cli/smartcontract/smart_contract.go | 21 ++++---------------
 cli/wallet/nep17.go                 |  2 +-
 cli/wallet/wallet.go                |  5 ++---
 4 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/cli/flags/address_test.go b/cli/flags/address_test.go
index a5019216d..c990f7729 100644
--- 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
index f7ea2fed2..a03f95b15 100644
--- a/cli/smartcontract/smart_contract.go
+++ b/cli/smartcontract/smart_contract.go
@@ -459,15 +459,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")
@@ -524,7 +518,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)
 	}
@@ -885,14 +879,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
index 007845368..86440c3e9 100644
--- 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
index c384178e7..f9334a84b 100644
--- 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)
 	}