From 676e15910c1a681e5ff2019bc7e53540119ceee5 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Fri, 14 Feb 2020 20:06:31 -0500
Subject: [PATCH] Add an fn to get a ChannelMonitor reference from a
 ManyChannelMonitor

This is really awkwrd, and relies on some (at least rather simple)
unsafe code (essentially avoiding the need for Pin), but it works
pretty well.
---
 lightning/src/lib.rs               |  1 -
 lightning/src/ln/channelmonitor.rs | 54 +++++++++++++++++++++++++++++-
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/lightning/src/lib.rs b/lightning/src/lib.rs
index 0cea7ad5d..7161d8881 100644
--- a/lightning/src/lib.rs
+++ b/lightning/src/lib.rs
@@ -10,7 +10,6 @@
 //! instead of having a rather-separate lightning appendage to a wallet.
 
 #![cfg_attr(not(feature = "fuzztarget"), deny(missing_docs))]
-#![forbid(unsafe_code)]
 
 // In general, rust is absolutely horrid at supporting users doing things like,
 // for example, compiling Rust code for real environments. Disable useless lints
diff --git a/lightning/src/ln/channelmonitor.rs b/lightning/src/ln/channelmonitor.rs
index 9e42749ce..b18a49e01 100644
--- a/lightning/src/ln/channelmonitor.rs
+++ b/lightning/src/ln/channelmonitor.rs
@@ -40,7 +40,7 @@ use util::ser::{Readable, MaybeReadable, Writer, Writeable, U48};
 use util::{byte_utils, events};
 
 use std::collections::{HashMap, hash_map};
-use std::sync::Mutex;
+use std::sync::{Mutex, MutexGuard};
 use std::{hash,cmp, mem};
 use std::ops::Deref;
 
@@ -209,6 +209,27 @@ pub trait ManyChannelMonitor<ChanSigner: ChannelKeys>: Send + Sync {
 	fn get_and_clear_pending_htlcs_updated(&self) -> Vec<HTLCUpdate>;
 }
 
+/// A lock held on a specific ManyChannelMonitor that includes a reference to the current version
+/// of a ChannelMonitor contained within.
+pub struct ManyChannelMonitorLock<'a, Key, ChanSigner: ChannelKeys> {
+	lock_ptr: *mut MutexGuard<'a, HashMap<Key, ChannelMonitor<ChanSigner>>>,
+	monitor: &'a ChannelMonitor<ChanSigner>,
+}
+impl<'a, Key, ChanSigner: ChannelKeys> ::std::ops::Deref for ManyChannelMonitorLock<'a, Key, ChanSigner> {
+	type Target = ChannelMonitor<ChanSigner>;
+	fn deref(&self) -> &ChannelMonitor<ChanSigner> {
+		self.monitor
+	}
+}
+impl<'a, Key, ChanSigner: ChannelKeys> Drop for ManyChannelMonitorLock<'a, Key, ChanSigner> {
+	fn drop(&mut self) {
+		// Dereferencing the lock_ptr is trivially safe here - it is created when this object is
+		// created, is never null, and is not modified at any point other than creation.
+		let _ = unsafe { Box::from_raw(self.lock_ptr) };
+		// Drop the box, freeing the lock
+	}
+}
+
 /// A simple implementation of a ManyChannelMonitor and ChainListener. Can be used to create a
 /// watchtower or watch our own channels.
 ///
@@ -318,6 +339,37 @@ impl<Key : Send + cmp::Eq + hash::Hash + 'static, ChanSigner: ChannelKeys, T: De
 			None => Err(MonitorUpdateError("No such monitor registered"))
 		}
 	}
+
+	/// Gets a reference to the latest copy of a given ChannelMonitor given a &Key, if any has been
+	/// registered.
+	///
+	/// The returned value contains a lock on this object, and other calls into this object will
+	/// almost certainly block until the returned value is dropped!
+	pub fn get_monitor_ref_by_key<'a>(&'a self, key: &Key) -> Option<ManyChannelMonitorLock<'a, Key, ChanSigner>> {
+		// Rust doesn't natively allow self-referential structs, and the only way to return a
+		// reference to something inside our Mutex is to return a struct that contains the lock and
+		// a reference to something pulled out of said lock.
+		// To avoid this, we have to fall back to some use of unsafe, but luckily its incredibly
+		// trivial - we simply Box up the MutexGuard and Box::leak() it, ensuring that its sitting
+		// in our heap without Rust having any reference to drop it.
+		// Then, we do a map lookup against the raw pointer, either returning a
+		// ManyChannelMonitorMonRef (which will drop the lock by recreating the Box when it gets
+		// dropped), or we will recreate the Box immediately and drop the lock before returning
+		// None.
+		//
+		// The returned ManyChannelMonitorMonRef is templated by a lifetime for which &self is
+		// valid, ensuring this object cannot be dropped until after the returned value is.
+		let lock = Box::new(self.monitors.lock().unwrap());
+		let lock_ptr: *mut MutexGuard<HashMap<Key, ChannelMonitor<ChanSigner>>> = Box::leak(lock);
+		let mon = unsafe { (*lock_ptr).get(key) };
+		if let Some(monitor) = mon {
+			Some(ManyChannelMonitorLock { lock_ptr, monitor })
+		} else {
+			let _ = unsafe { Box::from_raw(lock_ptr) };
+			// Drop the lock again
+			None
+		}
+	}
 }
 
 impl<ChanSigner: ChannelKeys, T: Deref + Sync + Send, F: Deref + Sync + Send, L: Deref + Sync + Send, C: Deref + Sync + Send> ManyChannelMonitor<ChanSigner> for SimpleManyChannelMonitor<OutPoint, ChanSigner, T, F, L, C>
-- 
2.39.5