From ca01b83adf75b00395fb0bc846837e5675fa8384 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 15 Nov 2021 12:29:47 +0300
Subject: [PATCH] [#78] owner: check id for validity

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 owner/id.go      | 26 +++++++++++++++++++++++++-
 owner/id_test.go | 29 +++++++++++++++++++++++++++++
 owner/test/id.go | 13 ++++++++++---
 3 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/owner/id.go b/owner/id.go
index b454286..6aed52e 100644
--- a/owner/id.go
+++ b/owner/id.go
@@ -6,6 +6,8 @@ import (
 	"fmt"
 
 	"github.com/mr-tron/base58"
+	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
+	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
 	"github.com/nspcc-dev/neofs-api-go/v2/refs"
 )
 
@@ -71,7 +73,7 @@ func (id *ID) Parse(s string) error {
 	data, err := base58.Decode(s)
 	if err != nil {
 		return fmt.Errorf("could not parse owner.ID from string: %w", err)
-	} else if len(data) != NEO3WalletSize {
+	} else if !valid(data) {
 		return errInvalidIDString
 	}
 
@@ -80,6 +82,28 @@ func (id *ID) Parse(s string) error {
 	return nil
 }
 
+// Valid returns true if id is a valid owner id.
+// The rules for v2 are the following:
+// 1. Must be 25 bytes in length.
+// 2. Must have N3 address prefix-byte.
+// 3. Last 4 bytes must contain the address checksum.
+func (id *ID) Valid() bool {
+	rawID := id.ToV2().GetValue()
+	return valid(rawID)
+}
+
+func valid(rawID []byte) bool {
+	if len(rawID) != NEO3WalletSize {
+		return false
+	}
+	if rawID[0] != address.NEO3Prefix {
+		return false
+	}
+
+	const boundIndex = NEO3WalletSize - 4
+	return bytes.Equal(rawID[boundIndex:], hash.Checksum(rawID[:boundIndex]))
+}
+
 // Marshal marshals ID into a protobuf binary form.
 func (id *ID) Marshal() ([]byte, error) {
 	return (*refs.OwnerID)(id).StableMarshal(nil)
diff --git a/owner/id_test.go b/owner/id_test.go
index 877a9bb..1737551 100644
--- a/owner/id_test.go
+++ b/owner/id_test.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/mr-tron/base58"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
+	"github.com/nspcc-dev/neo-go/pkg/util/slice"
 	"github.com/nspcc-dev/neofs-api-go/v2/refs"
 	. "github.com/nspcc-dev/neofs-sdk-go/owner"
 	ownertest "github.com/nspcc-dev/neofs-sdk-go/owner/test"
@@ -20,6 +21,34 @@ func TestIDV2(t *testing.T) {
 	require.Equal(t, id, NewIDFromV2(idV2))
 }
 
+func TestID_Valid(t *testing.T) {
+	id := ownertest.GenerateID()
+	require.True(t, id.Valid())
+
+	val := id.ToV2().GetValue()
+
+	t.Run("invalid prefix", func(t *testing.T) {
+		val := slice.Copy(val)
+		val[0] ^= 0xFF
+
+		id := ownertest.GenerateIDFromBytes(val)
+		require.False(t, id.Valid())
+	})
+	t.Run("invalid size", func(t *testing.T) {
+		val := val[:NEO3WalletSize-1]
+
+		id := ownertest.GenerateIDFromBytes(val)
+		require.False(t, id.Valid())
+	})
+	t.Run("invalid checksum", func(t *testing.T) {
+		val := slice.Copy(val)
+		val[NEO3WalletSize-1] ^= 0xFF
+
+		id := ownertest.GenerateIDFromBytes(val)
+		require.False(t, id.Valid())
+	})
+}
+
 func TestNewIDFromNeo3Wallet(t *testing.T) {
 	p, err := keys.NewPrivateKey()
 	require.NoError(t, err)
diff --git a/owner/test/id.go b/owner/test/id.go
index fe360c9..0ffcd80 100644
--- a/owner/test/id.go
+++ b/owner/test/id.go
@@ -3,6 +3,9 @@ package ownertest
 import (
 	"math/rand"
 
+	"github.com/mr-tron/base58"
+	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
+	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/neofs-api-go/v2/refs"
 	"github.com/nspcc-dev/neofs-sdk-go/owner"
 )
@@ -10,10 +13,14 @@ import (
 // GenerateID returns owner.ID calculated
 // from a random owner.NEO3Wallet.
 func GenerateID() *owner.ID {
-	data := make([]byte, owner.NEO3WalletSize)
-
-	rand.Read(data)
+	u := util.Uint160{}
+	rand.Read(u[:])
 
+	addr := address.Uint160ToString(u)
+	data, err := base58.Decode(addr)
+	if err != nil {
+		panic(err)
+	}
 	return GenerateIDFromBytes(data)
 }