Merge pull request #33 from tnull/2022-11-initial-wallet

Author: tnull

src/wallet.rs

@@ -0,0 +1,339 @@
+use crate::logger::{
+	log_error, log_given_level, log_internal, log_trace, FilesystemLogger, Logger,
+};
+
+use crate::Error;
+
+use lightning::chain::chaininterface::{
+	BroadcasterInterface, ConfirmationTarget, FeeEstimator, FEERATE_FLOOR_SATS_PER_KW,
+};
+
+use lightning::chain::keysinterface::{
+	EntropySource, InMemorySigner, KeyMaterial, KeysManager, NodeSigner, Recipient, SignerProvider,
+	SpendableOutputDescriptor,
+};
+use lightning::ln::msgs::DecodeError;
+use lightning::ln::script::ShutdownScript;
+
+use bdk::blockchain::{Blockchain, EsploraBlockchain};
+use bdk::database::BatchDatabase;
+use bdk::wallet::AddressIndex;
+use bdk::{FeeRate, SignOptions, SyncOptions};
+
+use bitcoin::bech32::u5;
+use bitcoin::secp256k1::ecdh::SharedSecret;
+use bitcoin::secp256k1::ecdsa::RecoverableSignature;
+use bitcoin::secp256k1::{PublicKey, Scalar, Secp256k1, SecretKey, Signing};
+use bitcoin::{Script, Transaction, TxOut};
+
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex, RwLock};
+
+pub struct Wallet<D>
+where
+	D: BatchDatabase,
+{
+	// A BDK blockchain used for wallet sync.
+	blockchain: EsploraBlockchain,
+	// A BDK on-chain wallet.
+	inner: Mutex<bdk::Wallet<D>>,
+	// A cache storing the most recently retrieved fee rate estimations.
+	fee_rate_cache: RwLock<HashMap<ConfirmationTarget, FeeRate>>,
+	tokio_runtime: RwLock<Option<Arc<tokio::runtime::Runtime>>>,
+	logger: Arc<FilesystemLogger>,
+}
+
+impl<D> Wallet<D>
+where
+	D: BatchDatabase,
+{
+	pub(crate) fn new(
+		blockchain: EsploraBlockchain, wallet: bdk::Wallet<D>, logger: Arc<FilesystemLogger>,
+	) -> Self {
+		let inner = Mutex::new(wallet);
+		let fee_rate_cache = RwLock::new(HashMap::new());
+		let tokio_runtime = RwLock::new(None);
+		Self { blockchain, inner, fee_rate_cache, tokio_runtime, logger }
+	}
+
+	pub(crate) async fn sync(&self) -> Result<(), Error> {
+		match self.update_fee_estimates().await {
+			Ok(()) => (),
+			Err(e) => {
+				log_error!(self.logger, "Fee estimation error: {}", e);
+				return Err(e);
+			}
+		}
+
+		let sync_options = SyncOptions { progress: None };
+		match self.inner.lock().unwrap().sync(&self.blockchain, sync_options).await {
+			Ok(()) => Ok(()),
+			Err(e) => {
+				log_error!(self.logger, "Wallet sync error: {}", e);
+				Err(From::from(e))
+			}
+		}
+	}
+
+	pub(crate) fn set_runtime(&self, tokio_runtime: Arc<tokio::runtime::Runtime>) {
+		*self.tokio_runtime.write().unwrap() = Some(tokio_runtime);
+	}
+
+	pub(crate) fn drop_runtime(&self) {
+		*self.tokio_runtime.write().unwrap() = None;
+	}
+
+	pub(crate) async fn update_fee_estimates(&self) -> Result<(), Error> {
+		let mut locked_fee_rate_cache = self.fee_rate_cache.write().unwrap();
+
+		let confirmation_targets = vec![
+			ConfirmationTarget::Background,
+			ConfirmationTarget::Normal,
+			ConfirmationTarget::HighPriority,
+		];
+		for target in confirmation_targets {
+			let num_blocks = match target {
+				ConfirmationTarget::Background => 12,
+				ConfirmationTarget::Normal => 6,
+				ConfirmationTarget::HighPriority => 3,
+			};
+
+			let est_fee_rate = self.blockchain.estimate_fee(num_blocks).await;
+
+			match est_fee_rate {
+				Ok(rate) => {
+					locked_fee_rate_cache.insert(target, rate);
+					log_trace!(
+						self.logger,
+						"Fee rate estimation updated: {} sats/kwu",
+						rate.fee_wu(1000)
+					);
+				}
+				Err(e) => {
+					log_error!(
+						self.logger,
+						"Failed to update fee rate estimation for {:?}: {}",
+						target,
+						e
+					);
+				}
+			}
+		}
+		Ok(())
+	}
+
+	pub(crate) fn create_funding_transaction(
+		&self, output_script: Script, value_sats: u64, confirmation_target: ConfirmationTarget,
+	) -> Result<Transaction, Error> {
+		let fee_rate = self.estimate_fee_rate(confirmation_target);
+
+		let locked_wallet = self.inner.lock().unwrap();
+		let mut tx_builder = locked_wallet.build_tx();
+
+		tx_builder.add_recipient(output_script, value_sats).fee_rate(fee_rate).enable_rbf();
+
+		let mut psbt = match tx_builder.finish() {
+			Ok((psbt, _)) => {
+				log_trace!(self.logger, "Created funding PSBT: {:?}", psbt);
+				psbt
+			}
+			Err(err) => {
+				log_error!(self.logger, "Failed to create funding transaction: {}", err);
+				Err(err)?
+			}
+		};
+
+		if !locked_wallet.sign(&mut psbt, SignOptions::default())? {
+			return Err(Error::FundingTxCreationFailed);
+		}
+
+		Ok(psbt.extract_tx())
+	}
+
+	pub(crate) fn get_new_address(&self) -> Result<bitcoin::Address, Error> {
+		let address_info = self.inner.lock().unwrap().get_address(AddressIndex::New)?;
+		Ok(address_info.address)
+	}
+
+	#[cfg(any(test))]
+	pub(crate) fn get_balance(&self) -> Result<bdk::Balance, Error> {
+		Ok(self.inner.lock().unwrap().get_balance()?)
+	}
+
+	fn estimate_fee_rate(&self, confirmation_target: ConfirmationTarget) -> FeeRate {
+		let locked_fee_rate_cache = self.fee_rate_cache.read().unwrap();
+
+		let fallback_sats_kwu = match confirmation_target {
+			ConfirmationTarget::Background => FEERATE_FLOOR_SATS_PER_KW,
+			ConfirmationTarget::Normal => 2000,
+			ConfirmationTarget::HighPriority => 5000,
+		};
+
+		// We'll fall back on this, if we really don't have any other information.
+		let fallback_rate = FeeRate::from_sat_per_kwu(fallback_sats_kwu as f32);
+
+		*locked_fee_rate_cache.get(&confirmation_target).unwrap_or(&fallback_rate)
+	}
+}
+
+impl<D> FeeEstimator for Wallet<D>
+where
+	D: BatchDatabase,
+{
+	fn get_est_sat_per_1000_weight(&self, confirmation_target: ConfirmationTarget) -> u32 {
+		(self.estimate_fee_rate(confirmation_target).fee_wu(1000) as u32)
+			.max(FEERATE_FLOOR_SATS_PER_KW)
+	}
+}
+
+impl<D> BroadcasterInterface for Wallet<D>
+where
+	D: BatchDatabase,
+{
+	fn broadcast_transaction(&self, tx: &Transaction) {
+		let locked_runtime = self.tokio_runtime.read().unwrap();
+		if locked_runtime.as_ref().is_none() {
+			log_error!(self.logger, "Failed to broadcast transaction: No runtime.");
+			unreachable!("Failed to broadcast transaction: No runtime.");
+		}
+
+		let res = tokio::task::block_in_place(move || {
+			locked_runtime
+				.as_ref()
+				.unwrap()
+				.block_on(async move { self.blockchain.broadcast(tx).await })
+		});
+
+		match res {
+			Ok(_) => {}
+			Err(err) => {
+				log_error!(self.logger, "Failed to broadcast transaction: {}", err);
+			}
+		}
+	}
+}
+
+/// Similar to [`KeysManager`], but overrides the destination and shutdown scripts so they are
+/// directly spendable by the BDK wallet.
+pub struct WalletKeysManager<D>
+where
+	D: BatchDatabase,
+{
+	inner: KeysManager,
+	wallet: Arc<Wallet<D>>,
+}
+
+impl<D> WalletKeysManager<D>
+where
+	D: BatchDatabase,
+{
+	/// Constructs a `WalletKeysManager` that overrides the destination and shutdown scripts.
+	///
+	/// See [`KeysManager::new`] for more information on `seed`, `starting_time_secs`, and
+	/// `starting_time_nanos`.
+	pub fn new(
+		seed: &[u8; 32], starting_time_secs: u64, starting_time_nanos: u32, wallet: Arc<Wallet<D>>,
+	) -> Self {
+		let inner = KeysManager::new(seed, starting_time_secs, starting_time_nanos);
+		Self { inner, wallet }
+	}
+
+	/// See [`KeysManager::spend_spendable_outputs`] for documentation on this method.
+	pub fn spend_spendable_outputs<C: Signing>(
+		&self, descriptors: &[&SpendableOutputDescriptor], outputs: Vec<TxOut>,
+		change_destination_script: Script, feerate_sat_per_1000_weight: u32,
+		secp_ctx: &Secp256k1<C>,
+	) -> Result<Transaction, ()> {
+		let only_non_static = &descriptors
+			.iter()
+			.filter(|desc| !matches!(desc, SpendableOutputDescriptor::StaticOutput { .. }))
+			.copied()
+			.collect::<Vec<_>>();
+		self.inner.spend_spendable_outputs(
+			only_non_static,
+			outputs,
+			change_destination_script,
+			feerate_sat_per_1000_weight,
+			secp_ctx,
+		)
+	}
+}
+
+impl<D> NodeSigner for WalletKeysManager<D>
+where
+	D: BatchDatabase,
+{
+	fn get_node_secret(&self, recipient: Recipient) -> Result<SecretKey, ()> {
+		self.inner.get_node_secret(recipient)
+	}
+
+	fn get_node_id(&self, recipient: Recipient) -> Result<PublicKey, ()> {
+		self.inner.get_node_id(recipient)
+	}
+
+	fn ecdh(
+		&self, recipient: Recipient, other_key: &PublicKey, tweak: Option<&Scalar>,
+	) -> Result<SharedSecret, ()> {
+		self.inner.ecdh(recipient, other_key, tweak)
+	}
+
+	fn get_inbound_payment_key_material(&self) -> KeyMaterial {
+		self.inner.get_inbound_payment_key_material()
+	}
+
+	fn sign_invoice(
+		&self, hrp_bytes: &[u8], invoice_data: &[u5], recipient: Recipient,
+	) -> Result<RecoverableSignature, ()> {
+		self.inner.sign_invoice(hrp_bytes, invoice_data, recipient)
+	}
+}
+
+impl<D> EntropySource for WalletKeysManager<D>
+where
+	D: BatchDatabase,
+{
+	fn get_secure_random_bytes(&self) -> [u8; 32] {
+		self.inner.get_secure_random_bytes()
+	}
+}
+
+impl<D> SignerProvider for WalletKeysManager<D>
+where
+	D: BatchDatabase,
+{
+	type Signer = InMemorySigner;
+
+	fn generate_channel_keys_id(
+		&self, inbound: bool, channel_value_satoshis: u64, user_channel_id: u128,
+	) -> [u8; 32] {
+		self.inner.generate_channel_keys_id(inbound, channel_value_satoshis, user_channel_id)
+	}
+
+	fn derive_channel_signer(
+		&self, channel_value_satoshis: u64, channel_keys_id: [u8; 32],
+	) -> Self::Signer {
+		self.inner.derive_channel_signer(channel_value_satoshis, channel_keys_id)
+	}
+
+	fn read_chan_signer(&self, reader: &[u8]) -> Result<Self::Signer, DecodeError> {
+		self.inner.read_chan_signer(reader)
+	}
+
+	fn get_destination_script(&self) -> Script {
+		let address =
+			self.wallet.get_new_address().expect("Failed to retrieve new address from wallet.");
+		address.script_pubkey()
+	}
+
+	fn get_shutdown_scriptpubkey(&self) -> ShutdownScript {
+		let address =
+			self.wallet.get_new_address().expect("Failed to retrieve new address from wallet.");
+		match address.payload {
+			bitcoin::util::address::Payload::WitnessProgram { version, program } => {
+				return ShutdownScript::new_witness_program(version, &program)
+					.expect("Invalid shutdown script.");
+			}
+			_ => panic!("Tried to use a non-witness address. This must not ever happen."),
+		}
+	}
+}