From b8c68dc2e4c6ce633e5702d0b909d903fedebd0b Mon Sep 17 00:00:00 2001
From: SChernykh <sergey.v.chernykh@gmail.com>
Date: Mon, 23 Oct 2023 19:21:45 +0200
Subject: [PATCH] Added `verify_merkle_proof` by index

---
 src/merkle.cpp             | 78 ++++++++++++++++++++++++++++++++++++++
 src/merkle.h               |  1 +
 tests/src/merkle_tests.cpp | 13 ++++++-
 3 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/src/merkle.cpp b/src/merkle.cpp
index 75959c4..8fbced0 100644
--- a/src/merkle.cpp
+++ b/src/merkle.cpp
@@ -193,4 +193,82 @@ bool verify_merkle_proof(hash h, const std::vector<std::pair<bool, hash>>& proof
 	return (h == root);
 }
 
+bool verify_merkle_proof(hash h, const std::vector<hash>& proof, size_t index, size_t count, const hash& root)
+{
+	if (index >= count) {
+		return false;
+	}
+
+	hash tmp[2];
+
+	if (count == 1) {
+	}
+	else if (count == 2) {
+		if (proof.empty()) {
+			return false;
+		}
+
+		if (index & 1) {
+			tmp[0] = proof[0];
+			tmp[1] = h;
+		}
+		else {
+			tmp[0] = h;
+			tmp[1] = proof[0];
+		}
+
+		keccak(tmp[0].h, HASH_SIZE * 2, h.h);
+	}
+	else {
+		size_t cnt = 1;
+		do { cnt <<= 1; } while (cnt <= count);
+		cnt >>= 1;
+
+		size_t proof_index = 0;
+
+		const size_t k = cnt * 2 - count;
+
+		if (index >= k) {
+			index -= k;
+
+			if (proof_index >= proof.size()) {
+				return false;
+			}
+
+			if (index & 1) {
+				tmp[0] = proof[proof_index];
+				tmp[1] = h;
+			}
+			else {
+				tmp[0] = h;
+				tmp[1] = proof[proof_index];
+			}
+
+			keccak(tmp[0].h, HASH_SIZE * 2, h.h);
+
+			index = (index >> 1) + k;
+			++proof_index;
+		}
+
+		for (; cnt >= 2; ++proof_index, index >>= 1, cnt >>= 1) {
+			if (proof_index >= proof.size()) {
+				return false;
+			}
+
+			if (index & 1) {
+				tmp[0] = proof[proof_index];
+				tmp[1] = h;
+			}
+			else {
+				tmp[0] = h;
+				tmp[1] = proof[proof_index];
+			}
+
+			keccak(tmp[0].h, HASH_SIZE * 2, h.h);
+		}
+	}
+
+	return (h == root);
+}
+
 } // namespace p2pool
diff --git a/src/merkle.h b/src/merkle.h
index 4c7da19..a9c6cbb 100644
--- a/src/merkle.h
+++ b/src/merkle.h
@@ -23,5 +23,6 @@ void merkle_hash(const std::vector<hash>& hashes, hash& root);
 void merkle_hash_full_tree(const std::vector<hash>& hashes, std::vector<std::vector<hash>>& tree);
 bool get_merkle_proof(const std::vector<std::vector<hash>>& tree, const hash& h, std::vector<std::pair<bool, hash>>& proof);
 bool verify_merkle_proof(hash h, const std::vector<std::pair<bool, hash>>& proof, const hash& root);
+bool verify_merkle_proof(hash h, const std::vector<hash>& proof, size_t index, size_t count, const hash& root);
 
 } // namespace p2pool
diff --git a/tests/src/merkle_tests.cpp b/tests/src/merkle_tests.cpp
index 9446176..9a0514f 100644
--- a/tests/src/merkle_tests.cpp
+++ b/tests/src/merkle_tests.cpp
@@ -70,10 +70,21 @@ TEST(merkle, root_hash)
 			}
 		}
 
-		for (const hash& h : hashes) {
+		for (size_t i = 0, n = hashes.size(); i < n; ++i) {
+			const hash& h = hashes[i];
 			std::vector<std::pair<bool, hash>> proof;
+
 			ASSERT_TRUE(get_merkle_proof(tree, h, proof));
 			ASSERT_TRUE(verify_merkle_proof(h, proof, root));
+
+			std::vector<hash> proof2;
+			proof2.reserve(proof.size());
+
+			for (const auto& i : proof) {
+				proof2.emplace_back(i.second);
+			}
+
+			ASSERT_TRUE(verify_merkle_proof(h, proof2, i, n, root));
 		}
 	};