From 0d7728239849e2cd0cb2d77f70faf26e47b1098f Mon Sep 17 00:00:00 2001
From: Yeastplume <yeastplume@protonmail.com>
Date: Tue, 3 Dec 2019 08:45:49 +0000
Subject: [PATCH] Check stored transaction for payment proof when checking
 proof (#266)

* check original transaction when checking payment proof

* rustfmt

* tweak to trigger build
---
 README.md                    |  1 -
 libwallet/src/internal/tx.rs | 39 ++++++++++++++++++++++--------------
 tests/cmd_line_basic.rs      |  4 ++++
 3 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index f375062f..de10ffeb 100644
--- a/README.md
+++ b/README.md
@@ -23,4 +23,3 @@ Full documentation outlining how to use the command line wallet can be found on
 # License
 
 Apache License v2.0
-
diff --git a/libwallet/src/internal/tx.rs b/libwallet/src/internal/tx.rs
index bc43d0e8..8e7350c3 100644
--- a/libwallet/src/internal/tx.rs
+++ b/libwallet/src/internal/tx.rs
@@ -474,7 +474,30 @@ where
 	C: NodeClient + 'a,
 	K: Keychain + 'a,
 {
+	let tx_vec = updater::retrieve_txs(wallet, None, Some(slate.id), Some(parent_key_id), false)?;
+	if tx_vec.len() == 0 {
+		return Err(ErrorKind::PaymentProof(
+			"TxLogEntry with original proof info not found (is account correct?)".to_owned(),
+		))?;
+	}
+
+	let orig_proof_info = tx_vec[0].clone().payment_proof;
+
+	if orig_proof_info.is_some() && slate.payment_proof.is_none() {
+		return Err(ErrorKind::PaymentProof(
+			"Expected Payment Proof for this Transaction is not present".to_owned(),
+		))?;
+	}
+
 	if let Some(ref p) = slate.payment_proof {
+		let orig_proof_info = match orig_proof_info {
+			Some(p) => p,
+			None => {
+				return Err(ErrorKind::PaymentProof(
+					"Original proof info not stored in tx".to_owned(),
+				))?;
+			}
+		};
 		let keychain = wallet.keychain(keychain_mask)?;
 		let index = match context.payment_proof_derivation_index {
 			Some(i) => i,
@@ -492,21 +515,7 @@ where
 				"Sender address on slate does not match original sender address".to_owned(),
 			))?;
 		}
-		let tx_vec =
-			updater::retrieve_txs(wallet, None, Some(slate.id), Some(&parent_key_id), false)?;
-		if tx_vec.len() != 1 {
-			return Err(ErrorKind::PaymentProof(
-				"TxLogEntry with original proof info not found".to_owned(),
-			))?;
-		}
-		let orig_proof_info = match tx_vec[0].clone().payment_proof {
-			Some(o) => o,
-			None => {
-				return Err(ErrorKind::PaymentProof(
-					"Original proof info not stored in tx".to_owned(),
-				))?;
-			}
-		};
+
 		if orig_proof_info.receiver_address != p.receiver_address {
 			return Err(ErrorKind::PaymentProof(
 				"Recipient address on slate does not match original recipient address".to_owned(),
diff --git a/tests/cmd_line_basic.rs b/tests/cmd_line_basic.rs
index 2a23526f..0eb5eb98 100644
--- a/tests/cmd_line_basic.rs
+++ b/tests/cmd_line_basic.rs
@@ -213,6 +213,8 @@ fn command_line_test_impl(test_dir: &str) -> Result<(), grin_wallet_controller::
 
 	let arg_vec = vec![
 		"grin-wallet",
+		"-a",
+		"mining",
 		"-p",
 		"password",
 		"finalize",
@@ -307,6 +309,8 @@ fn command_line_test_impl(test_dir: &str) -> Result<(), grin_wallet_controller::
 
 	let arg_vec = vec![
 		"grin-wallet",
+		"-a",
+		"mining",
 		"-p",
 		"password",
 		"finalize",