Skip to main content

brk_types/
addr_bytes.rs

1use std::str::FromStr;
2
3use bitcoin::{Network, PublicKey, ScriptBuf, opcodes, script::Builder};
4use brk_error::Error;
5
6use super::{
7    OutputType, P2ABytes, P2PK33Bytes, P2PK65Bytes, P2PKHBytes, P2SHBytes, P2TRBytes, P2WPKHBytes,
8    P2WSHBytes,
9};
10
11#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
12pub enum AddrBytes {
13    P2PK65(P2PK65Bytes), // 65
14    P2PK33(P2PK33Bytes), // 33
15    P2PKH(P2PKHBytes),   // 20
16    P2SH(P2SHBytes),     // 20
17    P2WPKH(P2WPKHBytes), // 20
18    P2WSH(P2WSHBytes),   // 32
19    P2TR(P2TRBytes),     // 32
20    P2A(P2ABytes),       // 2
21}
22
23impl AddrBytes {
24    pub fn as_slice(&self) -> &[u8] {
25        match self {
26            AddrBytes::P2PK65(bytes) => &bytes[..],
27            AddrBytes::P2PK33(bytes) => &bytes[..],
28            AddrBytes::P2PKH(bytes) => &bytes[..],
29            AddrBytes::P2SH(bytes) => &bytes[..],
30            AddrBytes::P2WPKH(bytes) => &bytes[..],
31            AddrBytes::P2WSH(bytes) => &bytes[..],
32            AddrBytes::P2TR(bytes) => &bytes[..],
33            AddrBytes::P2A(bytes) => &bytes[..],
34        }
35    }
36
37    pub fn hash(&self) -> u64 {
38        rapidhash::v3::rapidhash_v3(self.as_slice())
39    }
40
41    pub(crate) fn script_payload(
42        script: &ScriptBuf,
43        output_type: OutputType,
44    ) -> Result<&[u8], Error> {
45        let bytes = script.as_bytes();
46        match output_type {
47            OutputType::P2PK65 => match bytes.len() {
48                67 => Ok(&bytes[1..66]),
49                received => Err(Error::WrongLength {
50                    expected: 67,
51                    received,
52                }),
53            },
54            OutputType::P2PK33 => match bytes.len() {
55                35 => Ok(&bytes[1..34]),
56                received => Err(Error::WrongLength {
57                    expected: 35,
58                    received,
59                }),
60            },
61            OutputType::P2PKH => Ok(&bytes[3..23]),
62            OutputType::P2SH => Ok(&bytes[2..22]),
63            OutputType::P2WPKH | OutputType::P2WSH | OutputType::P2TR | OutputType::P2A => {
64                Ok(&bytes[2..])
65            }
66            OutputType::P2MS | OutputType::Unknown | OutputType::Empty | OutputType::OpReturn => {
67                Err(Error::WrongAddrType)
68            }
69        }
70    }
71
72    /// Reconstruct the script_pubkey from the address bytes
73    pub fn to_script_pubkey(&self) -> ScriptBuf {
74        match self {
75            AddrBytes::P2PK65(b) => Builder::new()
76                .push_slice(***b)
77                .push_opcode(opcodes::all::OP_CHECKSIG)
78                .into_script(),
79            AddrBytes::P2PK33(b) => Builder::new()
80                .push_slice(***b)
81                .push_opcode(opcodes::all::OP_CHECKSIG)
82                .into_script(),
83            AddrBytes::P2PKH(b) => Builder::new()
84                .push_opcode(opcodes::all::OP_DUP)
85                .push_opcode(opcodes::all::OP_HASH160)
86                .push_slice(***b)
87                .push_opcode(opcodes::all::OP_EQUALVERIFY)
88                .push_opcode(opcodes::all::OP_CHECKSIG)
89                .into_script(),
90            AddrBytes::P2SH(b) => Builder::new()
91                .push_opcode(opcodes::all::OP_HASH160)
92                .push_slice(***b)
93                .push_opcode(opcodes::all::OP_EQUAL)
94                .into_script(),
95            AddrBytes::P2WPKH(b) => Builder::new().push_int(0).push_slice(***b).into_script(),
96            AddrBytes::P2WSH(b) => Builder::new().push_int(0).push_slice(***b).into_script(),
97            AddrBytes::P2TR(b) => Builder::new().push_int(1).push_slice(***b).into_script(),
98            AddrBytes::P2A(b) => Builder::new().push_int(1).push_slice(***b).into_script(),
99        }
100    }
101}
102
103impl TryFrom<&ScriptBuf> for AddrBytes {
104    type Error = Error;
105    fn try_from(script: &ScriptBuf) -> Result<Self, Self::Error> {
106        Self::try_from((script, OutputType::from(script)))
107    }
108}
109
110impl TryFrom<(&ScriptBuf, OutputType)> for AddrBytes {
111    type Error = Error;
112    fn try_from(tuple: (&ScriptBuf, OutputType)) -> Result<Self, Self::Error> {
113        let (script, output_type) = tuple;
114        let bytes = Self::script_payload(script, output_type)?;
115
116        Ok(match output_type {
117            OutputType::P2PK65 => Self::P2PK65(P2PK65Bytes::from(bytes)),
118            OutputType::P2PK33 => Self::P2PK33(P2PK33Bytes::from(bytes)),
119            OutputType::P2PKH => Self::P2PKH(P2PKHBytes::from(bytes)),
120            OutputType::P2SH => Self::P2SH(P2SHBytes::from(bytes)),
121            OutputType::P2WPKH => Self::P2WPKH(P2WPKHBytes::from(bytes)),
122            OutputType::P2WSH => Self::P2WSH(P2WSHBytes::from(bytes)),
123            OutputType::P2TR => Self::P2TR(P2TRBytes::from(bytes)),
124            OutputType::P2A => Self::P2A(P2ABytes::from(bytes)),
125            OutputType::P2MS | OutputType::Unknown | OutputType::Empty | OutputType::OpReturn => {
126                unreachable!()
127            }
128        })
129    }
130}
131
132impl From<P2PK65Bytes> for AddrBytes {
133    #[inline]
134    fn from(value: P2PK65Bytes) -> Self {
135        Self::P2PK65(value)
136    }
137}
138
139impl From<P2PK33Bytes> for AddrBytes {
140    #[inline]
141    fn from(value: P2PK33Bytes) -> Self {
142        Self::P2PK33(value)
143    }
144}
145
146impl From<P2PKHBytes> for AddrBytes {
147    #[inline]
148    fn from(value: P2PKHBytes) -> Self {
149        Self::P2PKH(value)
150    }
151}
152
153impl From<P2SHBytes> for AddrBytes {
154    #[inline]
155    fn from(value: P2SHBytes) -> Self {
156        Self::P2SH(value)
157    }
158}
159
160impl From<P2WPKHBytes> for AddrBytes {
161    #[inline]
162    fn from(value: P2WPKHBytes) -> Self {
163        Self::P2WPKH(value)
164    }
165}
166
167impl From<P2WSHBytes> for AddrBytes {
168    #[inline]
169    fn from(value: P2WSHBytes) -> Self {
170        Self::P2WSH(value)
171    }
172}
173
174impl From<P2TRBytes> for AddrBytes {
175    #[inline]
176    fn from(value: P2TRBytes) -> Self {
177        Self::P2TR(value)
178    }
179}
180
181impl From<P2ABytes> for AddrBytes {
182    #[inline]
183    fn from(value: P2ABytes) -> Self {
184        Self::P2A(value)
185    }
186}
187
188impl AddrBytes {
189    /// Parse an address string to a ScriptBuf
190    pub fn addr_to_script(addr: &str) -> Result<ScriptBuf, Error> {
191        if let Ok(addr) = bitcoin::Address::from_str(addr) {
192            if !addr.is_valid_for_network(Network::Bitcoin) {
193                return Err(Error::InvalidNetwork);
194            }
195            let addr = addr.assume_checked();
196            Ok(addr.script_pubkey())
197        } else if let Ok(pubkey) = PublicKey::from_str(addr) {
198            Ok(ScriptBuf::new_p2pk(&pubkey))
199        } else {
200            Err(Error::InvalidAddr)
201        }
202    }
203}
204
205impl FromStr for AddrBytes {
206    type Err = Error;
207
208    fn from_str(s: &str) -> Result<Self, Self::Err> {
209        let script = Self::addr_to_script(s)?;
210        let output_type = OutputType::from(&script);
211        Self::try_from((&script, output_type))
212    }
213}