From 0899529c26902ef4af3bd3e9af7a382f4f07662e Mon Sep 17 00:00:00 2001
From: Evan Feenstra <evanfeenstra@gmail.com>
Date: Sat, 21 Oct 2023 09:40:42 -0700
Subject: [PATCH] add a test

---
 lightning/src/ln/onion_utils.rs | 113 +++++++++++++++++++++++++++++++-
 1 file changed, 110 insertions(+), 3 deletions(-)

diff --git a/lightning/src/ln/onion_utils.rs b/lightning/src/ln/onion_utils.rs
index 7bc001bf3ea..0bad1290099 100644
--- a/lightning/src/ln/onion_utils.rs
+++ b/lightning/src/ln/onion_utils.rs
@@ -13,7 +13,7 @@ use crate::ln::msgs;
 use crate::ln::wire::Encode;
 use crate::routing::gossip::NetworkUpdate;
 use crate::routing::router::{BlindedTail, Path, RouteHop};
-use crate::sign::NodeSigner;
+use crate::sign::{NodeSigner, Recipient};
 use crate::util::chacha20::{ChaCha20, ChaChaReader};
 use crate::util::errors::{self, APIError};
 use crate::util::ser::{Readable, ReadableArgs, Writeable, Writer, LengthCalculatingWriter};
@@ -1044,12 +1044,20 @@ pub enum PeeledPayment {
 /// Unwrap one layer of an incoming HTLC, returning either or a received payment, or a another
 /// onion to forward.
 pub fn peel_payment_onion<NS: Deref>(
-	shared_secret: [u8; 32], onion: &msgs::OnionPacket, payment_hash: PaymentHash,
-	node_signer: &NS, secp_ctx: &Secp256k1<secp256k1::All>
+	onion: &msgs::OnionPacket, payment_hash: PaymentHash, node_signer: &NS,
+	secp_ctx: &Secp256k1<secp256k1::All>
 ) -> Result<PeeledPayment, ()> 
 where 
 	NS::Target: NodeSigner,
 {
+	if onion.public_key.is_err() {
+		return Err(());
+	}
+	let shared_secret = node_signer
+		.ecdh(Recipient::Node, &onion.public_key.unwrap(), None)
+		.unwrap()
+		.secret_bytes();
+
 	let hop = decode_next_payment_hop(shared_secret, &onion.hop_data[..], onion.hmac, payment_hash, node_signer).map_err(|_| ())?;
 	let peeled = match hop {
 		Hop::Forward { next_hop_data, next_hop_hmac, new_packet_bytes } => {
@@ -1381,4 +1389,103 @@ mod tests {
 			writer.write_all(&self.data[..])
 		}
 	}
+
+	#[test]
+	fn create_and_peel_payment_onion() {
+		use crate::ln::channelmanager::RecipientOnionFields;
+		use crate::ln::PaymentPreimage;
+		use std::convert::TryInto;
+		use super::{create_payment_onion, peel_payment_onion};
+		use super::{Sha256, Hash, PeeledPayment, ReceivedPayment};
+
+		let secp_ctx = Secp256k1::new();
+		let recipient_onion = RecipientOnionFields::spontaneous_empty();
+		let session_priv_bytes = [42; 32];
+		let session_priv = SecretKey::from_slice(&session_priv_bytes).unwrap();
+		let amt_msat = 1000;
+		let cur_height = 1000;
+		let preimage_bytes = [43; 32];
+		let preimage = PaymentPreimage(preimage_bytes);
+		let rhash = Sha256::hash(&preimage_bytes).to_vec();
+		let rhash_bytes: [u8; 32] = rhash.try_into().unwrap();
+		let payment_hash = PaymentHash(rhash_bytes);
+		let prng_seed = [44; 32];
+
+		// let alice = make_keys_manager(&[1; 32]);
+		let bob = make_keys_manager(&[2; 32]);
+		let bob_pk = PublicKey::from_secret_key(&secp_ctx, &bob.get_node_secret_key());
+		let charlie = make_keys_manager(&[3; 32]);
+		let charlie_pk = PublicKey::from_secret_key(&secp_ctx, &charlie.get_node_secret_key());
+
+		// make a route alice -> bob -> charlie
+		let bob_fee = 1;
+		let recipient_amount = 999;
+		let hops = vec![
+			RouteHop {
+				pubkey: bob_pk,
+				fee_msat: bob_fee,
+				cltv_expiry_delta: 0,
+				short_channel_id: 1,
+				node_features: NodeFeatures::empty(),
+				channel_features: ChannelFeatures::empty(),
+				maybe_announced_channel: false,
+			},
+			RouteHop {
+				pubkey: charlie_pk,
+				fee_msat: recipient_amount,
+				cltv_expiry_delta: 0,
+				short_channel_id: 2,
+				node_features: NodeFeatures::empty(),
+				channel_features: ChannelFeatures::empty(),
+				maybe_announced_channel: false,
+			}
+		];
+		let path = Path {
+			hops: hops,
+			blinded_tail: None,
+		};
+
+		let (_htlc_msat, _htlc_cltv, onion) = create_payment_onion(
+			&secp_ctx, &path, &session_priv, amt_msat, recipient_onion, cur_height, payment_hash, 
+			Some(preimage), prng_seed
+		).unwrap();
+
+		// bob peels to find another onion
+		let next_onion = match peel_payment_onion(&onion, payment_hash, &&bob, &secp_ctx).unwrap() {
+			PeeledPayment::Receive(_) => {
+				panic!("should not be a receive");
+			},
+			PeeledPayment::Forward(forwarded) => {
+				forwarded.onion_packet
+			},
+		};
+
+		// charlie peels to find a received payment
+		match peel_payment_onion(&next_onion, payment_hash, &&charlie, &secp_ctx).unwrap() {
+			PeeledPayment::Receive(received) => match received {
+				ReceivedPayment::Regular{amt_msat, keysend_preimage, ..} => {
+					assert_eq!(amt_msat, recipient_amount, "wrong received amount");
+					assert_eq!(Some(preimage_bytes), keysend_preimage, "wrong received preimage");
+				},
+				ReceivedPayment::Blinded{..} => {
+					panic!("should not be blinded");
+				}
+			},
+			PeeledPayment::Forward(_) => {
+				panic!("should not be a receive");
+			},
+		};
+
+	}
+
+	fn make_keys_manager(seed: &[u8; 32]) -> crate::sign::KeysManager {
+		use crate::sign::KeysManager;
+		use std::time::{SystemTime, UNIX_EPOCH};
+		let start = SystemTime::now();
+		let now = start
+			.duration_since(UNIX_EPOCH)
+			.expect("Time went backwards");
+		KeysManager::new(seed, now.as_secs(), now.subsec_nanos())
+	}
+
 }