From 5c2ea2d5bbccd8d164c99cb57c2c2cdb69c0926c Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Mon, 8 Feb 2021 12:26:02 +0300
Subject: [PATCH] network: refactor P2PNotaryRequest decoding

We need to provide magic for both main and fallback transactions during
decoding, because transactions hashes depend on it.
---
 pkg/network/payload/notary_request.go      | 29 ++++++++++++++++++++--
 pkg/network/payload/notary_request_test.go | 14 ++++++++---
 pkg/network/server_test.go                 |  2 ++
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/pkg/network/payload/notary_request.go b/pkg/network/payload/notary_request.go
index 1a7f5e560..137320fb1 100644
--- a/pkg/network/payload/notary_request.go
+++ b/pkg/network/payload/notary_request.go
@@ -24,6 +24,31 @@ type P2PNotaryRequest struct {
 	signedHash util.Uint256
 }
 
+// NewP2PNotaryRequestFromBytes decodes P2PNotaryRequest from the given bytes.
+func NewP2PNotaryRequestFromBytes(network netmode.Magic, b []byte) (*P2PNotaryRequest, error) {
+	req := &P2PNotaryRequest{Network: network}
+	br := io.NewBinReaderFromBuf(b)
+	req.DecodeBinary(br)
+	if br.Err != nil {
+		return nil, br.Err
+	}
+	_ = br.ReadB()
+	if br.Err == nil {
+		return nil, errors.New("additional data after the payload")
+	}
+	return req, nil
+}
+
+// Bytes returns serialized P2PNotaryRequest payload.
+func (r *P2PNotaryRequest) Bytes() ([]byte, error) {
+	buf := io.NewBufBinWriter()
+	r.EncodeBinary(buf.BinWriter)
+	if buf.Err != nil {
+		return nil, buf.Err
+	}
+	return buf.Bytes(), nil
+}
+
 // Hash returns payload's hash.
 func (r *P2PNotaryRequest) Hash() util.Uint256 {
 	if r.hash.Equals(util.Uint256{}) {
@@ -74,8 +99,8 @@ func (r *P2PNotaryRequest) updateHashes(b []byte) {
 
 // DecodeBinaryUnsigned reads payload from w excluding signature.
 func (r *P2PNotaryRequest) decodeHashableFields(br *io.BinReader) {
-	r.MainTransaction = new(transaction.Transaction)
-	r.FallbackTransaction = new(transaction.Transaction)
+	r.MainTransaction = &transaction.Transaction{Network: r.Network}
+	r.FallbackTransaction = &transaction.Transaction{Network: r.Network}
 	r.MainTransaction.DecodeBinary(br)
 	r.FallbackTransaction.DecodeBinary(br)
 	if br.Err == nil {
diff --git a/pkg/network/payload/notary_request_test.go b/pkg/network/payload/notary_request_test.go
index fd204676b..5c29bb294 100644
--- a/pkg/network/payload/notary_request_test.go
+++ b/pkg/network/payload/notary_request_test.go
@@ -4,7 +4,7 @@ import (
 	"testing"
 
 	"github.com/nspcc-dev/neo-go/internal/random"
-	"github.com/nspcc-dev/neo-go/internal/testserdes"
+	"github.com/nspcc-dev/neo-go/pkg/config/netmode"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/nspcc-dev/neo-go/pkg/util"
@@ -143,8 +143,9 @@ func TestNotaryRequestIsValid(t *testing.T) {
 	})
 }
 
-func TestNotaryRequestEncodeDecodeBinary(t *testing.T) {
+func TestNotaryRequestBytesFromBytes(t *testing.T) {
 	mainTx := &transaction.Transaction{
+		Network:         netmode.UnitTestNet,
 		Attributes:      []transaction.Attribute{{Type: transaction.NotaryAssistedT, Value: &transaction.NotaryAssisted{NKeys: 1}}},
 		Script:          []byte{0, 1, 2},
 		ValidUntilBlock: 123,
@@ -157,6 +158,7 @@ func TestNotaryRequestEncodeDecodeBinary(t *testing.T) {
 	_ = mainTx.Hash()
 	_ = mainTx.Size()
 	fallbackTx := &transaction.Transaction{
+		Network:         netmode.UnitTestNet,
 		Script:          []byte{3, 2, 1},
 		ValidUntilBlock: 123,
 		Attributes: []transaction.Attribute{
@@ -172,6 +174,7 @@ func TestNotaryRequestEncodeDecodeBinary(t *testing.T) {
 	_ = fallbackTx.Hash()
 	_ = fallbackTx.Size()
 	p := &P2PNotaryRequest{
+		Network:             netmode.UnitTestNet,
 		MainTransaction:     mainTx,
 		FallbackTransaction: fallbackTx,
 		Witness: transaction.Witness{
@@ -180,5 +183,10 @@ func TestNotaryRequestEncodeDecodeBinary(t *testing.T) {
 		},
 	}
 	require.Equal(t, hash.Sha256(p.GetSignedHash().BytesBE()), p.Hash())
-	testserdes.EncodeDecodeBinary(t, p, new(P2PNotaryRequest))
+
+	bytes, err := p.Bytes()
+	require.NoError(t, err)
+	actual, err := NewP2PNotaryRequestFromBytes(netmode.UnitTestNet, bytes)
+	require.NoError(t, err)
+	require.Equal(t, p, actual)
 }
diff --git a/pkg/network/server_test.go b/pkg/network/server_test.go
index ef3d6432f..85fd51c36 100644
--- a/pkg/network/server_test.go
+++ b/pkg/network/server_test.go
@@ -524,6 +524,7 @@ func TestGetData(t *testing.T) {
 	})
 	t.Run("p2pNotaryRequest", func(t *testing.T) {
 		mainTx := &transaction.Transaction{
+			Network:         netmode.UnitTestNet,
 			Attributes:      []transaction.Attribute{{Type: transaction.NotaryAssistedT, Value: &transaction.NotaryAssisted{NKeys: 1}}},
 			Script:          []byte{0, 1, 2},
 			ValidUntilBlock: 123,
@@ -533,6 +534,7 @@ func TestGetData(t *testing.T) {
 		mainTx.Size()
 		mainTx.Hash()
 		fallbackTx := &transaction.Transaction{
+			Network:         netmode.UnitTestNet,
 			Script:          []byte{1, 2, 3},
 			ValidUntilBlock: 123,
 			Attributes: []transaction.Attribute{