From 625f4b2e3dd546188ae49ae454f1930de28ba09d Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Fri, 4 Sep 2020 22:16:07 +0200
Subject: [PATCH] go/runtime/transaction: Ensure consistent batch order indices

---
 go/runtime/transaction/transaction.go      |  9 ++++++++-
 go/runtime/transaction/transaction_test.go | 21 +++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/go/runtime/transaction/transaction.go b/go/runtime/transaction/transaction.go
index 9b1369e3b08..65e15fbe9d4 100644
--- a/go/runtime/transaction/transaction.go
+++ b/go/runtime/transaction/transaction.go
@@ -45,7 +45,7 @@ const (
 
 // MarshalBinary encodes an artifact kind into binary form.
 func (ak artifactKind) MarshalBinary() (data []byte, err error) {
-	// kindInvalid should not be marshaleld.
+	// kindInvalid should not be marshaled.
 	if ak == kindInvalid {
 		return nil, errMalformedArtifactKind
 	}
@@ -258,6 +258,13 @@ func (t *Tree) GetInputBatch(ctx context.Context, maxBatchSize, maxBatchSizeByte
 
 	// Sort transactions to be in batch order.
 	sort.Stable(bo)
+
+	// Make sure that item orders are consistent.
+	for i, v := range bo.order {
+		if uint32(i) != v {
+			return nil, fmt.Errorf("transaction: inconsistent order: item %d has batch order %d", i, v)
+		}
+	}
 	bo.order = nil
 
 	return bo.batch, nil
diff --git a/go/runtime/transaction/transaction_test.go b/go/runtime/transaction/transaction_test.go
index d8392f73af4..229b8035770 100644
--- a/go/runtime/transaction/transaction_test.go
+++ b/go/runtime/transaction/transaction_test.go
@@ -159,3 +159,24 @@ func TestTransaction(t *testing.T) {
 		require.True(t, txnsByHash[checkTx.Hash()].Equal(&checkTx), "transaction should have the correct artifacts") // nolint: gosec
 	}
 }
+
+func TestTransactionInvalidBatchOrder(t *testing.T) {
+	ctx := context.Background()
+	store := mkvs.New(nil, nil)
+
+	var emptyRoot node.Root
+	emptyRoot.Empty()
+
+	tree := NewTree(store, emptyRoot)
+
+	tx := Transaction{
+		Input:      []byte("this goes in"),
+		Output:     []byte("and this comes out"),
+		BatchOrder: 1, // Invalid batch order as the first transaction should be at index zero.
+	}
+	err := tree.AddTransaction(ctx, tx, nil)
+	require.NoError(t, err, "AddTransaction")
+
+	_, err = tree.GetInputBatch(ctx, 0, 0)
+	require.Error(t, err, "GetInputBatch should fail with inconsistent order")
+}