From 988dfcc7fe1d71c82129879e6e33290ea76c66c3 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Wed, 24 Feb 2021 11:35:09 +0300
Subject: [PATCH] services: update Notary transaction size after completion

Close #1766
---
 pkg/core/notary_test.go       | 15 +++++++++++++++
 pkg/services/notary/notary.go | 18 +++++++++++++++++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/pkg/core/notary_test.go b/pkg/core/notary_test.go
index f1f438395..d3ba93ec8 100644
--- a/pkg/core/notary_test.go
+++ b/pkg/core/notary_test.go
@@ -20,6 +20,7 @@ import (
 	"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/crypto/keys"
+	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/network/payload"
 	"github.com/nspcc-dev/neo-go/pkg/services/notary"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
@@ -184,12 +185,18 @@ func TestNotary(t *testing.T) {
 			}
 			scripts[i].InvocationScript = append([]byte{byte(opcode.PUSHDATA1), 64}, requesters[i].PrivateKey().Sign(main.GetSignedPart())...)
 			main.Scripts = scripts
+
+			_ = main.Size() // for size update test
+
 			var fallback *transaction.Transaction
 			if len(NVBincrements) == len(requesters) {
 				fallback = createFallbackTx(requesters[i], main, NVBincrements[i])
 			} else {
 				fallback = createFallbackTx(requesters[i], main)
 			}
+
+			_ = fallback.Size() // for size update test
+
 			payloads[i] = &payload.P2PNotaryRequest{
 				MainTransaction:     main,
 				FallbackTransaction: fallback,
@@ -255,6 +262,10 @@ func TestNotary(t *testing.T) {
 			require.NotNil(t, completedTx, errors.New("main transaction expected to be completed"))
 			require.Equal(t, nKeys+1, len(completedTx.Signers))
 			require.Equal(t, nKeys+1, len(completedTx.Scripts))
+
+			// check that tx size was updated
+			require.Equal(t, io.GetVarSize(completedTx), completedTx.Size())
+
 			interopCtx := bc.newInteropContext(trigger.Verification, bc.dao, nil, completedTx)
 			for i, req := range requests {
 				require.Equal(t, req.MainTransaction.Scripts[i], completedTx.Scripts[i])
@@ -305,6 +316,10 @@ func TestNotary(t *testing.T) {
 					InvocationScript:   append([]byte{byte(opcode.PUSHDATA1), 64}, acc1.PrivateKey().Sign(req.FallbackTransaction.GetSignedPart())...),
 					VerificationScript: []byte{},
 				}, completedTx.Scripts[0])
+
+				// check that tx size was updated
+				require.Equal(t, io.GetVarSize(completedTx), completedTx.Size())
+
 				interopCtx := bc.newInteropContext(trigger.Verification, bc.dao, nil, completedTx)
 				_, err := bc.verifyHashAgainstScript(completedTx.Signers[1].Account, &completedTx.Scripts[1], interopCtx, -1)
 				require.NoError(t, err)
diff --git a/pkg/services/notary/notary.go b/pkg/services/notary/notary.go
index b8da872be..eb5b83cf3 100644
--- a/pkg/services/notary/notary.go
+++ b/pkg/services/notary/notary.go
@@ -15,6 +15,7 @@ import (
 	"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/crypto/keys"
+	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/network/payload"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/neo-go/pkg/vm"
@@ -316,7 +317,22 @@ func (n *Notary) finalize(tx *transaction.Transaction) error {
 			break
 		}
 	}
-	return n.onTransaction(tx)
+	newTx, err := updateTxSize(tx)
+	if err != nil {
+		return fmt.Errorf("failed to update completed transaction's size: %w", err)
+	}
+
+	return n.onTransaction(newTx)
+}
+
+// updateTxSize returns transaction with re-calculated size and an error.
+func updateTxSize(tx *transaction.Transaction) (*transaction.Transaction, error) {
+	bw := io.NewBufBinWriter()
+	tx.EncodeBinary(bw.BinWriter)
+	if bw.Err != nil {
+		return nil, fmt.Errorf("encode binary: %w", bw.Err)
+	}
+	return transaction.NewTransactionFromBytes(tx.Network, tx.Bytes())
 }
 
 // verifyIncompleteWitnesses checks that tx either doesn't have all witnesses attached (in this case none of them