From 73e7f4d6e2dd642f85e11fb0ce94955948442224 Mon Sep 17 00:00:00 2001
From: Jeffrey Czyz <jkczyz@gmail.com>
Date: Mon, 26 Jul 2021 12:31:24 -0400
Subject: [PATCH] Support all shutdown scripts defined in BOLT 2

KeysInterface::get_shutdown_pubkey is used to form P2WPKH shutdown
scripts. However, BOLT 2 allows for a wider variety of scripts. Refactor
KeysInterface to allow any supported script while still maintaining
serialization backwards compatibility with P2WPKH script pubkeys stored
simply as the PublicKey.

Add an optional TLV field to Channel and ChannelMonitor to support the
new format, but continue to serialize the legacy PublicKey format.
---
 fuzz/src/chanmon_consistency.rs       |  7 +++--
 fuzz/src/full_stack.rs                |  7 +++--
 lightning/src/chain/channelmonitor.rs | 30 ++++++++++++-------
 lightning/src/chain/keysinterface.rs  | 14 ++++-----
 lightning/src/ln/channel.rs           | 42 ++++++++++++++++++---------
 lightning/src/ln/script.rs            | 36 ++++++++++++++++++++++-
 lightning/src/util/test_utils.rs      |  5 ++--
 7 files changed, 103 insertions(+), 38 deletions(-)

diff --git a/fuzz/src/chanmon_consistency.rs b/fuzz/src/chanmon_consistency.rs
index b879ec9b2..e808dc581 100644
--- a/fuzz/src/chanmon_consistency.rs
+++ b/fuzz/src/chanmon_consistency.rs
@@ -39,6 +39,7 @@ use lightning::ln::{PaymentHash, PaymentPreimage, PaymentSecret};
 use lightning::ln::channelmanager::{ChainParameters, ChannelManager, PaymentSendFailure, ChannelManagerReadArgs};
 use lightning::ln::features::{ChannelFeatures, InitFeatures, NodeFeatures};
 use lightning::ln::msgs::{CommitmentUpdate, ChannelMessageHandler, DecodeError, UpdateAddHTLC, Init};
+use lightning::ln::script::ShutdownScript;
 use lightning::util::enforcing_trait_impls::{EnforcingSigner, INITIAL_REVOKED_COMMITMENT_NUMBER};
 use lightning::util::errors::APIError;
 use lightning::util::events;
@@ -164,9 +165,11 @@ impl KeysInterface for KeyProvider {
 		Builder::new().push_opcode(opcodes::all::OP_PUSHBYTES_0).push_slice(&our_channel_monitor_claim_key_hash[..]).into_script()
 	}
 
-	fn get_shutdown_pubkey(&self) -> PublicKey {
+	fn get_shutdown_scriptpubkey(&self) -> ShutdownScript {
 		let secp_ctx = Secp256k1::signing_only();
-		PublicKey::from_secret_key(&secp_ctx, &SecretKey::from_slice(&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, self.node_id]).unwrap())
+		let secret_key = SecretKey::from_slice(&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, self.node_id]).unwrap();
+		let pubkey_hash = WPubkeyHash::hash(&PublicKey::from_secret_key(&secp_ctx, &secret_key).serialize());
+		ShutdownScript::new_p2wpkh(&pubkey_hash)
 	}
 
 	fn get_channel_signer(&self, _inbound: bool, channel_value_satoshis: u64) -> EnforcingSigner {
diff --git a/fuzz/src/full_stack.rs b/fuzz/src/full_stack.rs
index 9398dcb0b..b108d2a06 100644
--- a/fuzz/src/full_stack.rs
+++ b/fuzz/src/full_stack.rs
@@ -36,6 +36,7 @@ use lightning::ln::{PaymentHash, PaymentPreimage, PaymentSecret};
 use lightning::ln::channelmanager::{ChainParameters, ChannelManager};
 use lightning::ln::peer_handler::{MessageHandler,PeerManager,SocketDescriptor};
 use lightning::ln::msgs::DecodeError;
+use lightning::ln::script::ShutdownScript;
 use lightning::routing::router::get_route;
 use lightning::routing::network_graph::NetGraphMsgHandler;
 use lightning::util::config::UserConfig;
@@ -271,9 +272,11 @@ impl KeysInterface for KeyProvider {
 		Builder::new().push_opcode(opcodes::all::OP_PUSHBYTES_0).push_slice(&our_channel_monitor_claim_key_hash[..]).into_script()
 	}
 
-	fn get_shutdown_pubkey(&self) -> PublicKey {
+	fn get_shutdown_scriptpubkey(&self) -> ShutdownScript {
 		let secp_ctx = Secp256k1::signing_only();
-		PublicKey::from_secret_key(&secp_ctx, &SecretKey::from_slice(&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]).unwrap())
+		let secret_key = SecretKey::from_slice(&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]).unwrap();
+		let pubkey_hash = WPubkeyHash::hash(&PublicKey::from_secret_key(&secp_ctx, &secret_key).serialize());
+		ShutdownScript::new_p2wpkh(&pubkey_hash)
 	}
 
 	fn get_channel_signer(&self, inbound: bool, channel_value_satoshis: u64) -> EnforcingSigner {
diff --git a/lightning/src/chain/channelmonitor.rs b/lightning/src/chain/channelmonitor.rs
index a8ec8ee97..cfbd0e64f 100644
--- a/lightning/src/chain/channelmonitor.rs
+++ b/lightning/src/chain/channelmonitor.rs
@@ -489,7 +489,7 @@ pub(crate) struct ChannelMonitorImpl<Signer: Sign> {
 	destination_script: Script,
 	broadcasted_holder_revokable_script: Option<(Script, PublicKey, PublicKey)>,
 	counterparty_payment_script: Script,
-	shutdown_script: Script,
+	shutdown_script: Option<Script>,
 
 	channel_keys_id: [u8; 32],
 	holder_revocation_basepoint: PublicKey,
@@ -665,7 +665,10 @@ impl<Signer: Sign> Writeable for ChannelMonitorImpl<Signer> {
 		}
 
 		self.counterparty_payment_script.write(writer)?;
-		self.shutdown_script.write(writer)?;
+		match &self.shutdown_script {
+			Some(script) => script.write(writer)?,
+			None => Script::new().write(writer)?,
+		}
 
 		self.channel_keys_id.write(writer)?;
 		self.holder_revocation_basepoint.write(writer)?;
@@ -788,14 +791,16 @@ impl<Signer: Sign> Writeable for ChannelMonitorImpl<Signer> {
 		self.lockdown_from_offchain.write(writer)?;
 		self.holder_tx_signed.write(writer)?;
 
-		write_tlv_fields!(writer, {});
+		write_tlv_fields!(writer, {
+			(1, self.shutdown_script, option),
+		});
 
 		Ok(())
 	}
 }
 
 impl<Signer: Sign> ChannelMonitor<Signer> {
-	pub(crate) fn new(secp_ctx: Secp256k1<secp256k1::All>, keys: Signer, shutdown_pubkey: &PublicKey,
+	pub(crate) fn new(secp_ctx: Secp256k1<secp256k1::All>, keys: Signer, shutdown_script: Option<Script>,
 	                  on_counterparty_tx_csv: u16, destination_script: &Script, funding_info: (OutPoint, Script),
 	                  channel_parameters: &ChannelTransactionParameters,
 	                  funding_redeemscript: Script, channel_value_satoshis: u64,
@@ -804,8 +809,6 @@ impl<Signer: Sign> ChannelMonitor<Signer> {
 	                  best_block: BestBlock) -> ChannelMonitor<Signer> {
 
 		assert!(commitment_transaction_number_obscure_factor <= (1 << 48));
-		let our_channel_close_key_hash = WPubkeyHash::hash(&shutdown_pubkey.serialize());
-		let shutdown_script = Builder::new().push_opcode(opcodes::all::OP_PUSHBYTES_0).push_slice(&our_channel_close_key_hash[..]).into_script();
 		let payment_key_hash = WPubkeyHash::hash(&keys.pubkeys().payment_point.serialize());
 		let counterparty_payment_script = Builder::new().push_opcode(opcodes::all::OP_PUSHBYTES_0).push_slice(&payment_key_hash[..]).into_script();
 
@@ -2472,7 +2475,7 @@ impl<Signer: Sign> ChannelMonitorImpl<Signer> {
 					channel_value_satoshis: self.channel_value_satoshis,
 				}));
 				break;
-			} else if outp.script_pubkey == self.shutdown_script {
+			} else if self.shutdown_script.as_ref() == Some(&outp.script_pubkey) {
 				spendable_output = Some(SpendableOutputDescriptor::StaticOutput {
 					outpoint: OutPoint { txid: tx.txid(), index: i as u16 },
 					output: outp.clone(),
@@ -2608,7 +2611,10 @@ impl<'a, Signer: Sign, K: KeysInterface<Signer = Signer>> ReadableArgs<&'a K>
 			_ => return Err(DecodeError::InvalidValue),
 		};
 		let counterparty_payment_script = Readable::read(reader)?;
-		let shutdown_script = Readable::read(reader)?;
+		let mut shutdown_script = {
+			let script = <Script as Readable>::read(reader)?;
+			if script.is_empty() { None } else { Some(script) }
+		};
 
 		let channel_keys_id = Readable::read(reader)?;
 		let holder_revocation_basepoint = Readable::read(reader)?;
@@ -2762,7 +2768,9 @@ impl<'a, Signer: Sign, K: KeysInterface<Signer = Signer>> ReadableArgs<&'a K>
 		let lockdown_from_offchain = Readable::read(reader)?;
 		let holder_tx_signed = Readable::read(reader)?;
 
-		read_tlv_fields!(reader, {});
+		read_tlv_fields!(reader, {
+			(1, shutdown_script, option),
+		});
 
 		let mut secp_ctx = Secp256k1::new();
 		secp_ctx.seeded_randomize(&keys_manager.get_secure_random_bytes());
@@ -2840,6 +2848,7 @@ mod tests {
 	use ln::{PaymentPreimage, PaymentHash};
 	use ln::chan_utils;
 	use ln::chan_utils::{HTLCOutputInCommitment, ChannelPublicKeys, ChannelTransactionParameters, HolderCommitmentTransaction, CounterpartyChannelTransactionParameters};
+	use ln::script::ShutdownScript;
 	use util::test_utils::{TestLogger, TestBroadcaster, TestFeeEstimator};
 	use bitcoin::secp256k1::key::{SecretKey,PublicKey};
 	use bitcoin::secp256k1::Secp256k1;
@@ -2933,9 +2942,10 @@ mod tests {
 		};
 		// Prune with one old state and a holder commitment tx holding a few overlaps with the
 		// old state.
+		let shutdown_pubkey = PublicKey::from_secret_key(&secp_ctx, &SecretKey::from_slice(&[42; 32]).unwrap());
 		let best_block = BestBlock::from_genesis(Network::Testnet);
 		let monitor = ChannelMonitor::new(Secp256k1::new(), keys,
-		                                  &PublicKey::from_secret_key(&secp_ctx, &SecretKey::from_slice(&[42; 32]).unwrap()), 0, &Script::new(),
+		                                  Some(ShutdownScript::new_p2wpkh_from_pubkey(shutdown_pubkey).into_inner()), 0, &Script::new(),
 		                                  (OutPoint { txid: Txid::from_slice(&[43; 32]).unwrap(), index: 0 }, Script::new()),
 		                                  &channel_parameters,
 		                                  Script::new(), 46, 0,
diff --git a/lightning/src/chain/keysinterface.rs b/lightning/src/chain/keysinterface.rs
index d7ff2a633..e88e78ecf 100644
--- a/lightning/src/chain/keysinterface.rs
+++ b/lightning/src/chain/keysinterface.rs
@@ -36,6 +36,7 @@ use chain::transaction::OutPoint;
 use ln::chan_utils;
 use ln::chan_utils::{HTLCOutputInCommitment, make_funding_redeemscript, ChannelPublicKeys, HolderCommitmentTransaction, ChannelTransactionParameters, CommitmentTransaction};
 use ln::msgs::UnsignedChannelAnnouncement;
+use ln::script::ShutdownScript;
 
 use prelude::*;
 use core::sync::atomic::{AtomicUsize, Ordering};
@@ -118,8 +119,8 @@ impl_writeable_tlv_based!(StaticPaymentOutputDescriptor, {
 #[derive(Clone, Debug, PartialEq)]
 pub enum SpendableOutputDescriptor {
 	/// An output to a script which was provided via KeysInterface directly, either from
-	/// `get_destination_script()` or `get_shutdown_pubkey()`, thus you should already know how to
-	/// spend it. No secret keys are provided as rust-lightning was never given any key.
+	/// `get_destination_script()` or `get_shutdown_scriptpubkey()`, thus you should already know
+	/// how to spend it. No secret keys are provided as rust-lightning was never given any key.
 	/// These may include outputs from a transaction punishing our counterparty or claiming an HTLC
 	/// on-chain using the payment preimage or after it has timed out.
 	StaticOutput {
@@ -351,12 +352,11 @@ pub trait KeysInterface {
 	/// This method should return a different value each time it is called, to avoid linking
 	/// on-chain funds across channels as controlled to the same user.
 	fn get_destination_script(&self) -> Script;
-	/// Get a public key which we will send funds to (in the form of a P2WPKH output) when closing
-	/// a channel.
+	/// Get a script pubkey which we will send funds to when closing a channel.
 	///
 	/// This method should return a different value each time it is called, to avoid linking
 	/// on-chain funds across channels as controlled to the same user.
-	fn get_shutdown_pubkey(&self) -> PublicKey;
+	fn get_shutdown_scriptpubkey(&self) -> ShutdownScript;
 	/// Get a new set of Sign for per-channel secrets. These MUST be unique even if you
 	/// restarted with some stale data!
 	///
@@ -1013,8 +1013,8 @@ impl KeysInterface for KeysManager {
 		self.destination_script.clone()
 	}
 
-	fn get_shutdown_pubkey(&self) -> PublicKey {
-		self.shutdown_pubkey.clone()
+	fn get_shutdown_scriptpubkey(&self) -> ShutdownScript {
+		ShutdownScript::new_p2wpkh_from_pubkey(self.shutdown_pubkey.clone())
 	}
 
 	fn get_channel_signer(&self, _inbound: bool, channel_value_satoshis: u64) -> Self::Signer {
diff --git a/lightning/src/ln/channel.rs b/lightning/src/ln/channel.rs
index 013c45982..f6313c075 100644
--- a/lightning/src/ln/channel.rs
+++ b/lightning/src/ln/channel.rs
@@ -9,15 +9,15 @@
 
 use bitcoin::blockdata::script::{Script,Builder};
 use bitcoin::blockdata::transaction::{TxIn, TxOut, Transaction, SigHashType};
-use bitcoin::blockdata::opcodes;
 use bitcoin::util::bip143;
 use bitcoin::consensus::encode;
 
 use bitcoin::hashes::Hash;
 use bitcoin::hashes::sha256::Hash as Sha256;
 use bitcoin::hashes::sha256d::Hash as Sha256d;
-use bitcoin::hash_types::{Txid, BlockHash, WPubkeyHash};
+use bitcoin::hash_types::{Txid, BlockHash};
 
+use bitcoin::secp256k1::constants::PUBLIC_KEY_SIZE;
 use bitcoin::secp256k1::key::{PublicKey,SecretKey};
 use bitcoin::secp256k1::{Secp256k1,Signature};
 use bitcoin::secp256k1;
@@ -349,7 +349,7 @@ pub(super) struct Channel<Signer: Sign> {
 	latest_monitor_update_id: u64,
 
 	holder_signer: Signer,
-	shutdown_pubkey: PublicKey,
+	shutdown_scriptpubkey: Option<ShutdownScript>,
 	destination_script: Script,
 
 	// Our commitment numbers start at 2^48-1 and count down, whereas the ones used in transaction
@@ -607,7 +607,7 @@ impl<Signer: Sign> Channel<Signer> {
 			latest_monitor_update_id: 0,
 
 			holder_signer,
-			shutdown_pubkey: keys_provider.get_shutdown_pubkey(),
+			shutdown_scriptpubkey: Some(keys_provider.get_shutdown_scriptpubkey()),
 			destination_script: keys_provider.get_destination_script(),
 
 			cur_holder_commitment_transaction_number: INITIAL_COMMITMENT_NUMBER,
@@ -851,7 +851,7 @@ impl<Signer: Sign> Channel<Signer> {
 			latest_monitor_update_id: 0,
 
 			holder_signer,
-			shutdown_pubkey: keys_provider.get_shutdown_pubkey(),
+			shutdown_scriptpubkey: Some(keys_provider.get_shutdown_scriptpubkey()),
 			destination_script: keys_provider.get_destination_script(),
 
 			cur_holder_commitment_transaction_number: INITIAL_COMMITMENT_NUMBER,
@@ -1130,8 +1130,7 @@ impl<Signer: Sign> Channel<Signer> {
 
 	#[inline]
 	fn get_closing_scriptpubkey(&self) -> Script {
-		let channel_close_key_hash = WPubkeyHash::hash(&self.shutdown_pubkey.serialize());
-		Builder::new().push_opcode(opcodes::all::OP_PUSHBYTES_0).push_slice(&channel_close_key_hash[..]).into_script()
+		self.shutdown_scriptpubkey.clone().unwrap().into_inner()
 	}
 
 	#[inline]
@@ -1676,8 +1675,9 @@ impl<Signer: Sign> Channel<Signer> {
 		let funding_redeemscript = self.get_funding_redeemscript();
 		let funding_txo_script = funding_redeemscript.to_v0_p2wsh();
 		let obscure_factor = get_commitment_transaction_number_obscure_factor(&self.get_holder_pubkeys().payment_point, &self.get_counterparty_pubkeys().payment_point, self.is_outbound());
+		let shutdown_script = self.shutdown_scriptpubkey.clone().map(|script| script.into_inner());
 		let channel_monitor = ChannelMonitor::new(self.secp_ctx.clone(), self.holder_signer.clone(),
-		                                          &self.shutdown_pubkey, self.get_holder_selected_contest_delay(),
+		                                          shutdown_script, self.get_holder_selected_contest_delay(),
 		                                          &self.destination_script, (funding_txo, funding_txo_script.clone()),
 		                                          &self.channel_transaction_parameters,
 		                                          funding_redeemscript.clone(), self.channel_value_satoshis,
@@ -1749,8 +1749,9 @@ impl<Signer: Sign> Channel<Signer> {
 		let funding_txo = self.get_funding_txo().unwrap();
 		let funding_txo_script = funding_redeemscript.to_v0_p2wsh();
 		let obscure_factor = get_commitment_transaction_number_obscure_factor(&self.get_holder_pubkeys().payment_point, &self.get_counterparty_pubkeys().payment_point, self.is_outbound());
+		let shutdown_script = self.shutdown_scriptpubkey.clone().map(|script| script.into_inner());
 		let channel_monitor = ChannelMonitor::new(self.secp_ctx.clone(), self.holder_signer.clone(),
-		                                          &self.shutdown_pubkey, self.get_holder_selected_contest_delay(),
+		                                          shutdown_script, self.get_holder_selected_contest_delay(),
 		                                          &self.destination_script, (funding_txo, funding_txo_script),
 		                                          &self.channel_transaction_parameters,
 		                                          funding_redeemscript.clone(), self.channel_value_satoshis,
@@ -4549,7 +4550,12 @@ impl<Signer: Sign> Writeable for Channel<Signer> {
 		(key_data.0.len() as u32).write(writer)?;
 		writer.write_all(&key_data.0[..])?;
 
-		self.shutdown_pubkey.write(writer)?;
+		// Write out the old serialization for shutdown_pubkey for backwards compatibility, if
+		// deserialized from that format.
+		match self.shutdown_scriptpubkey.as_ref().and_then(|script| script.as_legacy_pubkey()) {
+			Some(shutdown_pubkey) => shutdown_pubkey.write(writer)?,
+			None => [0u8; PUBLIC_KEY_SIZE].write(writer)?,
+		}
 		self.destination_script.write(writer)?;
 
 		self.cur_holder_commitment_transaction_number.write(writer)?;
@@ -4745,6 +4751,7 @@ impl<Signer: Sign> Writeable for Channel<Signer> {
 			(1, self.minimum_depth, option),
 			(3, self.counterparty_selected_channel_reserve_satoshis, option),
 			(5, self.config, required),
+			(7, self.shutdown_scriptpubkey, option),
 		});
 
 		Ok(())
@@ -4788,7 +4795,12 @@ impl<'a, Signer: Sign, K: Deref> ReadableArgs<&'a K> for Channel<Signer>
 		}
 		let holder_signer = keys_source.read_chan_signer(&keys_data)?;
 
-		let shutdown_pubkey = Readable::read(reader)?;
+		// Read the old serialization for shutdown_pubkey, preferring it for shutdown_scriptpubkey
+		// over the TLV if valid.
+		let mut shutdown_scriptpubkey = match <PublicKey as Readable>::read(reader) {
+			Ok(pubkey) => Some(ShutdownScript::new_p2wpkh_from_pubkey(pubkey)),
+			Err(_) => None,
+		};
 		let destination_script = Readable::read(reader)?;
 
 		let cur_holder_commitment_transaction_number = Readable::read(reader)?;
@@ -4959,6 +4971,7 @@ impl<'a, Signer: Sign, K: Deref> ReadableArgs<&'a K> for Channel<Signer>
 			(1, minimum_depth, option),
 			(3, counterparty_selected_channel_reserve_satoshis, option),
 			(5, config, option), // Note that if none is provided we will *not* overwrite the existing one.
+			(7, shutdown_scriptpubkey, option),
 		});
 
 		let mut secp_ctx = Secp256k1::new();
@@ -4976,7 +4989,7 @@ impl<'a, Signer: Sign, K: Deref> ReadableArgs<&'a K> for Channel<Signer>
 			latest_monitor_update_id,
 
 			holder_signer,
-			shutdown_pubkey,
+			shutdown_scriptpubkey,
 			destination_script,
 
 			cur_holder_commitment_transaction_number,
@@ -5069,6 +5082,7 @@ mod tests {
 	use ln::channel::MAX_FUNDING_SATOSHIS;
 	use ln::features::InitFeatures;
 	use ln::msgs::{ChannelUpdate, DataLossProtect, DecodeError, OptionalField, UnsignedChannelUpdate};
+	use ln::script::ShutdownScript;
 	use ln::chan_utils;
 	use ln::chan_utils::{ChannelPublicKeys, HolderCommitmentTransaction, CounterpartyChannelTransactionParameters, HTLC_SUCCESS_TX_WEIGHT, HTLC_TIMEOUT_TX_WEIGHT};
 	use chain::BestBlock;
@@ -5118,10 +5132,10 @@ mod tests {
 			Builder::new().push_opcode(opcodes::all::OP_PUSHBYTES_0).push_slice(&channel_monitor_claim_key_hash[..]).into_script()
 		}
 
-		fn get_shutdown_pubkey(&self) -> PublicKey {
+		fn get_shutdown_scriptpubkey(&self) -> ShutdownScript {
 			let secp_ctx = Secp256k1::signing_only();
 			let channel_close_key = SecretKey::from_slice(&hex::decode("0fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff").unwrap()[..]).unwrap();
-			PublicKey::from_secret_key(&secp_ctx, &channel_close_key)
+			ShutdownScript::new_p2wpkh_from_pubkey(PublicKey::from_secret_key(&secp_ctx, &channel_close_key))
 		}
 
 		fn get_channel_signer(&self, _inbound: bool, _channel_value_satoshis: u64) -> InMemorySigner {
diff --git a/lightning/src/ln/script.rs b/lightning/src/ln/script.rs
index 589c3f311..a01866de6 100644
--- a/lightning/src/ln/script.rs
+++ b/lightning/src/ln/script.rs
@@ -8,19 +8,24 @@ use bitcoin::hash_types::{PubkeyHash, ScriptHash, WPubkeyHash, WScriptHash};
 use bitcoin::secp256k1::key::PublicKey;
 
 use ln::features::InitFeatures;
+use ln::msgs::DecodeError;
+use util::ser::{Readable, Writeable, Writer};
 
 use std::convert::TryFrom;
+use std::io::Read;
 use core::num::NonZeroU8;
 
 /// A script pubkey for shutting down a channel as defined by [BOLT #2].
 ///
 /// [BOLT #2]: https://github.com/lightningnetwork/lightning-rfc/blob/master/02-peer-protocol.md
+#[derive(Clone)]
 pub struct ShutdownScript(ShutdownScriptImpl);
 
 /// An error occurring when converting from [`Script`] to [`ShutdownScript`].
 #[derive(Debug)]
 pub struct InvalidShutdownScript(Script);
 
+#[derive(Clone)]
 enum ShutdownScriptImpl {
 	/// [`PublicKey`] used to form a P2WPKH script pubkey. Used to support backward-compatible
 	/// serialization.
@@ -30,9 +35,30 @@ enum ShutdownScriptImpl {
 	Bolt2(Script),
 }
 
+impl Writeable for ShutdownScript {
+	fn write<W: Writer>(&self, w: &mut W) -> Result<(), ::std::io::Error> {
+		self.0.write(w)
+	}
+
+	fn serialized_length(&self) -> usize {
+		self.0.serialized_length()
+	}
+}
+
+impl Readable for ShutdownScript {
+	fn read<R: Read>(r: &mut R) -> Result<Self, DecodeError> {
+		Ok(ShutdownScript(ShutdownScriptImpl::read(r)?))
+	}
+}
+
+impl_writeable_tlv_based_enum!(ShutdownScriptImpl, ;
+	(0, Legacy),
+	(1, Bolt2),
+);
+
 impl ShutdownScript {
 	/// Generates a P2WPKH script pubkey from the given [`PublicKey`].
-	pub fn new_p2wpkh_from_pubkey(pubkey: PublicKey) -> Self {
+	pub(crate) fn new_p2wpkh_from_pubkey(pubkey: PublicKey) -> Self {
 		Self(ShutdownScriptImpl::Legacy(pubkey))
 	}
 
@@ -71,6 +97,14 @@ impl ShutdownScript {
 	pub fn into_inner(self) -> Script {
 		self.into()
 	}
+
+	/// Returns the [`PublicKey`] used for a P2WPKH shutdown script if constructed directly from it.
+	pub fn as_legacy_pubkey(&self) -> Option<&PublicKey> {
+		match &self.0 {
+			ShutdownScriptImpl::Legacy(pubkey) => Some(pubkey),
+			ShutdownScriptImpl::Bolt2(_) => None,
+		}
+	}
 }
 
 impl TryFrom<Script> for ShutdownScript {
diff --git a/lightning/src/util/test_utils.rs b/lightning/src/util/test_utils.rs
index fbd5a26d9..743bf6b24 100644
--- a/lightning/src/util/test_utils.rs
+++ b/lightning/src/util/test_utils.rs
@@ -19,6 +19,7 @@ use chain::keysinterface;
 use ln::features::{ChannelFeatures, InitFeatures};
 use ln::msgs;
 use ln::msgs::OptionalField;
+use ln::script::ShutdownScript;
 use util::enforcing_trait_impls::{EnforcingSigner, INITIAL_REVOKED_COMMITMENT_NUMBER};
 use util::events;
 use util::logger::{Logger, Level, Record};
@@ -70,7 +71,7 @@ impl keysinterface::KeysInterface for OnlyReadsKeysInterface {
 
 	fn get_node_secret(&self) -> SecretKey { unreachable!(); }
 	fn get_destination_script(&self) -> Script { unreachable!(); }
-	fn get_shutdown_pubkey(&self) -> PublicKey { unreachable!(); }
+	fn get_shutdown_scriptpubkey(&self) -> ShutdownScript { unreachable!(); }
 	fn get_channel_signer(&self, _inbound: bool, _channel_value_satoshis: u64) -> EnforcingSigner { unreachable!(); }
 	fn get_secure_random_bytes(&self) -> [u8; 32] { [0; 32] }
 
@@ -458,7 +459,7 @@ impl keysinterface::KeysInterface for TestKeysInterface {
 
 	fn get_node_secret(&self) -> SecretKey { self.backing.get_node_secret() }
 	fn get_destination_script(&self) -> Script { self.backing.get_destination_script() }
-	fn get_shutdown_pubkey(&self) -> PublicKey { self.backing.get_shutdown_pubkey() }
+	fn get_shutdown_scriptpubkey(&self) -> ShutdownScript { self.backing.get_shutdown_scriptpubkey() }
 	fn get_channel_signer(&self, inbound: bool, channel_value_satoshis: u64) -> EnforcingSigner {
 		let keys = self.backing.get_channel_signer(inbound, channel_value_satoshis);
 		let revoked_commitment = self.make_revoked_commitment_cell(keys.commitment_seed);
-- 
2.39.5