From cc75b52a4353eb0d52468da98a8778777e178564 Mon Sep 17 00:00:00 2001 From: Luke Parker Date: Wed, 31 Jan 2024 17:54:43 -0500 Subject: [PATCH] Don't allow constructing unusable serai_client::bitcoin::Address es --- processor/src/networks/bitcoin.rs | 26 ++---- substrate/client/src/networks/bitcoin.rs | 95 +++++++++++++-------- tests/full-stack/src/tests/mint_and_burn.rs | 2 +- tests/processor/src/networks.rs | 6 +- 4 files changed, 73 insertions(+), 56 deletions(-) diff --git a/processor/src/networks/bitcoin.rs b/processor/src/networks/bitcoin.rs index b47f0f03..65908aa3 100644 --- a/processor/src/networks/bitcoin.rs +++ b/processor/src/networks/bitcoin.rs @@ -134,8 +134,7 @@ impl OutputTrait for Output { fn write(&self, writer: &mut W) -> io::Result<()> { self.kind.write(writer)?; - let presumed_origin: Option> = - self.presumed_origin.clone().map(|presumed_origin| presumed_origin.try_into().unwrap()); + let presumed_origin: Option> = self.presumed_origin.clone().map(Into::into); writer.write_all(&presumed_origin.encode())?; self.output.write(writer)?; writer.write_all(&u16::try_from(self.data.len()).unwrap().to_le_bytes())?; @@ -415,7 +414,7 @@ impl Bitcoin { .iter() .map(|payment| { ( - payment.address.0.clone(), + payment.address.clone().into(), // If we're solely estimating the fee, don't specify the actual amount // This won't affect the fee calculation yet will ensure we don't hit a not enough funds // error @@ -427,7 +426,7 @@ impl Bitcoin { match BSignableTransaction::new( inputs.iter().map(|input| input.output.clone()).collect(), &payments, - change.as_ref().map(|change| &change.0), + change.as_ref().map(AsRef::as_ref), None, fee.0, ) { @@ -532,7 +531,8 @@ impl Network for Bitcoin { } fn external_address(key: ProjectivePoint) -> Address { - Address(BAddress::::new(BNetwork::Bitcoin, address_payload(key).unwrap())) + Address::new(BAddress::::new(BNetwork::Bitcoin, address_payload(key).unwrap())) + .unwrap() } fn branch_address(key: ProjectivePoint) -> Address { @@ -603,17 +603,9 @@ impl Network for Bitcoin { } tx.unwrap().output.swap_remove(usize::try_from(spent_output.vout).unwrap()) }; - BAddress::from_script(&spent_output.script_pubkey, BNetwork::Bitcoin).ok().and_then( - |address| { - let address = Address(address); - let encoded: Result, _> = address.clone().try_into(); - if encoded.is_ok() { - Some(address) - } else { - None - } - }, - ) + BAddress::from_script(&spent_output.script_pubkey, BNetwork::Bitcoin) + .ok() + .and_then(Address::new) }; let output = Output { kind, presumed_origin, output, data }; @@ -802,7 +794,7 @@ impl Network for Bitcoin { }], output: vec![TxOut { value: tx.output[0].value - BAmount::from_sat(10000), - script_pubkey: address.0.script_pubkey(), + script_pubkey: address.as_ref().script_pubkey(), }], }; diff --git a/substrate/client/src/networks/bitcoin.rs b/substrate/client/src/networks/bitcoin.rs index 336b1731..42cf41bf 100644 --- a/substrate/client/src/networks/bitcoin.rs +++ b/substrate/client/src/networks/bitcoin.rs @@ -12,9 +12,8 @@ use bitcoin::{ type BAddress = BAddressGeneric; -// TODO: Add a new so you can't create an address which can't be encoded #[derive(Clone, Eq, Debug)] -pub struct Address(pub BAddress); +pub struct Address(BAddress); impl PartialEq for Address { fn eq(&self, other: &Self) -> bool { @@ -27,11 +26,12 @@ impl PartialEq for Address { impl FromStr for Address { type Err = Error; fn from_str(str: &str) -> Result { - Ok(Address( + Address::new( BAddressGeneric::from_str(str) .map_err(|_| Error::UnrecognizedScript)? .require_network(Network::Bitcoin)?, - )) + ) + .ok_or(Error::UnrecognizedScript) } } @@ -77,38 +77,63 @@ impl TryFrom> for Address { } } -#[allow(clippy::from_over_into)] -impl TryInto> for Address { - type Error = (); - fn try_into(self) -> Result, ()> { - Ok( - (match self.0.payload() { - Payload::PubkeyHash(hash) => EncodedAddress::P2PKH(*hash.as_raw_hash().as_byte_array()), - Payload::ScriptHash(hash) => EncodedAddress::P2SH(*hash.as_raw_hash().as_byte_array()), - Payload::WitnessProgram(program) => match program.version() { - WitnessVersion::V0 => { - let program = program.program(); - if program.len() == 20 { - let mut buf = [0; 20]; - buf.copy_from_slice(program.as_ref()); - EncodedAddress::P2WPKH(buf) - } else if program.len() == 32 { - let mut buf = [0; 32]; - buf.copy_from_slice(program.as_ref()); - EncodedAddress::P2WSH(buf) - } else { - Err(())? - } +fn try_to_vec(addr: &Address) -> Result, ()> { + Ok( + (match addr.0.payload() { + Payload::PubkeyHash(hash) => EncodedAddress::P2PKH(*hash.as_raw_hash().as_byte_array()), + Payload::ScriptHash(hash) => EncodedAddress::P2SH(*hash.as_raw_hash().as_byte_array()), + Payload::WitnessProgram(program) => match program.version() { + WitnessVersion::V0 => { + let program = program.program(); + if program.len() == 20 { + let mut buf = [0; 20]; + buf.copy_from_slice(program.as_ref()); + EncodedAddress::P2WPKH(buf) + } else if program.len() == 32 { + let mut buf = [0; 32]; + buf.copy_from_slice(program.as_ref()); + EncodedAddress::P2WSH(buf) + } else { + Err(())? } - WitnessVersion::V1 => { - let program_ref: &[u8] = program.program().as_ref(); - EncodedAddress::P2TR(program_ref.try_into().map_err(|_| ())?) - } - _ => Err(())?, - }, + } + WitnessVersion::V1 => { + let program_ref: &[u8] = program.program().as_ref(); + EncodedAddress::P2TR(program_ref.try_into().map_err(|_| ())?) + } _ => Err(())?, - }) - .encode(), - ) + }, + _ => Err(())?, + }) + .encode(), + ) +} + +impl From
for Vec { + fn from(addr: Address) -> Vec { + // Safe since only encodable addresses can be created + try_to_vec(&addr).unwrap() + } +} + +impl From
for BAddress { + fn from(addr: Address) -> BAddress { + addr.0 + } +} + +impl AsRef for Address { + fn as_ref(&self) -> &BAddress { + &self.0 + } +} + +impl Address { + pub fn new(address: BAddress) -> Option { + let res = Self(address); + if try_to_vec(&res).is_ok() { + return Some(res); + } + None } } diff --git a/tests/full-stack/src/tests/mint_and_burn.rs b/tests/full-stack/src/tests/mint_and_burn.rs index 4fe1378e..9a18ac54 100644 --- a/tests/full-stack/src/tests/mint_and_burn.rs +++ b/tests/full-stack/src/tests/mint_and_burn.rs @@ -514,7 +514,7 @@ async fn mint_and_burn_test() { Coin::Bitcoin, 1_000_000_00, ExternalAddress::new( - serai_client::networks::bitcoin::Address(bitcoin_addr.clone()).try_into().unwrap(), + serai_client::networks::bitcoin::Address::new(bitcoin_addr.clone()).unwrap().into(), ) .unwrap(), ) diff --git a/tests/processor/src/networks.rs b/tests/processor/src/networks.rs index a54703ce..dc234476 100644 --- a/tests/processor/src/networks.rs +++ b/tests/processor/src/networks.rs @@ -389,9 +389,9 @@ impl Wallet { Wallet::Bitcoin { public_key, .. } => { use bitcoin_serai::bitcoin::{Network, Address}; ExternalAddress::new( - networks::bitcoin::Address(Address::p2pkh(public_key, Network::Regtest)) - .try_into() - .unwrap(), + networks::bitcoin::Address::new(Address::p2pkh(public_key, Network::Regtest)) + .unwrap() + .into(), ) .unwrap() }