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), P2PK33(P2PK33Bytes), P2PKH(P2PKHBytes), P2SH(P2SHBytes), P2WPKH(P2WPKHBytes), P2WSH(P2WSHBytes), P2TR(P2TRBytes), P2A(P2ABytes), }
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 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 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}