Don't allow constructing unusable serai_client::bitcoin::Address es

This commit is contained in:
Luke Parker 2024-01-31 17:54:43 -05:00
parent 4913873b10
commit cc75b52a43
No known key found for this signature in database
4 changed files with 73 additions and 56 deletions

View file

@ -134,8 +134,7 @@ impl OutputTrait<Bitcoin> for Output {
fn write<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
self.kind.write(writer)?;
let presumed_origin: Option<Vec<u8>> =
self.presumed_origin.clone().map(|presumed_origin| presumed_origin.try_into().unwrap());
let presumed_origin: Option<Vec<u8>> = 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::<NetworkChecked>::new(BNetwork::Bitcoin, address_payload(key).unwrap()))
Address::new(BAddress::<NetworkChecked>::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<Vec<u8>, _> = 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(),
}],
};

View file

@ -12,9 +12,8 @@ use bitcoin::{
type BAddress = BAddressGeneric<NetworkChecked>;
// 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<Address, Error> {
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<Vec<u8>> for Address {
}
}
#[allow(clippy::from_over_into)]
impl TryInto<Vec<u8>> for Address {
type Error = ();
fn try_into(self) -> Result<Vec<u8>, ()> {
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<Vec<u8>, ()> {
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<Address> for Vec<u8> {
fn from(addr: Address) -> Vec<u8> {
// Safe since only encodable addresses can be created
try_to_vec(&addr).unwrap()
}
}
impl From<Address> for BAddress {
fn from(addr: Address) -> BAddress {
addr.0
}
}
impl AsRef<BAddress> for Address {
fn as_ref(&self) -> &BAddress {
&self.0
}
}
impl Address {
pub fn new(address: BAddress) -> Option<Self> {
let res = Self(address);
if try_to_vec(&res).is_ok() {
return Some(res);
}
None
}
}

View file

@ -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(),
)

View file

@ -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()
}