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