From 4e7f9e1f306eae1d67f871bb76e07bb9cfa8e08e Mon Sep 17 00:00:00 2001
From: alrevuelta <alvrevuelta@gmail.com>
Date: Tue, 16 Jan 2024 16:01:01 +0100
Subject: [PATCH] Ensure no out of bound

---
 rln/rln.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/rln/rln.go b/rln/rln.go
index 78b3860..282227a 100644
--- a/rln/rln.go
+++ b/rln/rln.go
@@ -406,6 +406,11 @@ func (r *RLN) GetMerkleProof(index MembershipIndex) (MerkleProof, error) {
 		return MerkleProof{}, err
 	}
 
+	// Check if we can read the first byte
+	if len(proofBytes) < 8 {
+		return MerkleProof{}, errors.New(fmt.Sprintf("wrong output size: %d", len(proofBytes)))
+	}
+
 	var result MerkleProof
 	var numElements big.Int
 	var numIndexes big.Int
@@ -416,6 +421,14 @@ func (r *RLN) GetMerkleProof(index MembershipIndex) (MerkleProof, error) {
 	numElements.SetBytes(revert(proofBytes[offset : offset+8]))
 	offset += 8
 
+	// With numElements we can determine the expected length of the proof.
+	expectedLen := 8 + int(32*numElements.Uint64()) + 8 + int(numElements.Uint64())
+	if len(proofBytes) != expectedLen {
+		return MerkleProof{}, errors.New(fmt.Sprintf("wrong output size expected: %d, current: %d",
+			expectedLen,
+			len(proofBytes)))
+	}
+
 	result.PathElements = make([]MerkleNode, numElements.Uint64())
 
 	for i := uint64(0); i < numElements.Uint64(); i++ {