Skip to main content

bsv/primitives/
utils.rs

1//! Utility functions shared across primitive modules.
2//!
3//! Includes hex encoding/decoding, Base58 encoding/decoding,
4//! and Base58Check encoding/decoding with checksum verification.
5
6use super::error::PrimitivesError;
7use super::hash::hash256;
8
9/// Base58 alphabet used by Bitcoin.
10const BASE58_ALPHABET: &[u8; 58] = b"123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";
11
12/// Encode bytes as a hexadecimal string.
13pub fn to_hex(bytes: &[u8]) -> String {
14    bytes.iter().map(|b| format!("{:02x}", b)).collect()
15}
16
17/// Decode a hexadecimal string into bytes.
18pub fn from_hex(hex: &str) -> Result<Vec<u8>, PrimitivesError> {
19    if !hex.len().is_multiple_of(2) {
20        return Err(PrimitivesError::InvalidHex(
21            "odd length hex string".to_string(),
22        ));
23    }
24    (0..hex.len())
25        .step_by(2)
26        .map(|i| {
27            u8::from_str_radix(&hex[i..i + 2], 16).map_err(|e| {
28                PrimitivesError::InvalidHex(format!("invalid hex char at position {}: {}", i, e))
29            })
30        })
31        .collect()
32}
33
34/// Encode bytes to a Base58 string.
35///
36/// Leading zero bytes in the input are preserved as '1' characters
37/// in the output, following the Bitcoin Base58 convention.
38pub fn base58_encode(data: &[u8]) -> String {
39    // Count leading zeros
40    let leading_zeros = data.iter().take_while(|&&b| b == 0).count();
41
42    // Convert bytes to base58 using repeated division
43    // We work with a mutable copy of the data as big-endian number
44    let mut result: Vec<u8> = Vec::new();
45
46    for &byte in data.iter() {
47        let mut carry = byte as u32;
48        for digit in result.iter_mut() {
49            let x = (*digit as u32) * 256 + carry;
50            *digit = (x % 58) as u8;
51            carry = x / 58;
52        }
53        while carry > 0 {
54            result.push((carry % 58) as u8);
55            carry /= 58;
56        }
57    }
58
59    // Build result string
60    let mut s = String::with_capacity(leading_zeros + result.len());
61
62    // Add '1' for each leading zero byte
63    for _ in 0..leading_zeros {
64        s.push('1');
65    }
66
67    // Convert base58 digits to characters (result is in reverse order)
68    for &digit in result.iter().rev() {
69        s.push(BASE58_ALPHABET[digit as usize] as char);
70    }
71
72    s
73}
74
75/// Decode a Base58 string to bytes.
76///
77/// Leading '1' characters in the input are converted back to zero bytes.
78pub fn base58_decode(s: &str) -> Result<Vec<u8>, PrimitivesError> {
79    if s.is_empty() {
80        return Err(PrimitivesError::InvalidFormat(
81            "empty base58 string".to_string(),
82        ));
83    }
84
85    // Build reverse lookup table
86    let mut alphabet_map = [255u8; 128];
87    for (i, &ch) in BASE58_ALPHABET.iter().enumerate() {
88        alphabet_map[ch as usize] = i as u8;
89    }
90
91    // Count leading '1' characters (zero bytes)
92    let leading_ones = s.chars().take_while(|&c| c == '1').count();
93
94    // Estimate output size
95    let size = ((s.len() - leading_ones) as f64 * (58.0_f64.ln() / 256.0_f64.ln()) + 1.0) as usize;
96    let mut result = vec![0u8; size];
97
98    for ch in s.chars() {
99        let ch_val = ch as usize;
100        if ch_val >= 128 || alphabet_map[ch_val] == 255 {
101            return Err(PrimitivesError::InvalidFormat(format!(
102                "invalid base58 character: {}",
103                ch
104            )));
105        }
106        let mut carry = alphabet_map[ch_val] as u32;
107        for byte in result.iter_mut() {
108            let x = (*byte as u32) * 58 + carry;
109            *byte = (x & 0xff) as u8;
110            carry = x >> 8;
111        }
112    }
113
114    // Remove leading zeros from the result (which is in reverse order)
115    result.reverse();
116    let skip = result.iter().take_while(|&&b| b == 0).count();
117    let result = &result[skip..];
118
119    // Prepend leading zero bytes
120    let mut output = vec![0u8; leading_ones];
121    output.extend_from_slice(result);
122
123    Ok(output)
124}
125
126/// Encode data with a prefix using Base58Check (includes 4-byte checksum).
127///
128/// Format: Base58(prefix || payload || checksum)
129/// where checksum = hash256(prefix || payload)[0..4]
130pub fn base58_check_encode(payload: &[u8], prefix: &[u8]) -> String {
131    let mut data = Vec::with_capacity(prefix.len() + payload.len() + 4);
132    data.extend_from_slice(prefix);
133    data.extend_from_slice(payload);
134
135    let checksum = hash256(&data);
136    data.extend_from_slice(&checksum[..4]);
137
138    base58_encode(&data)
139}
140
141/// Decode a Base58Check string, verifying the checksum.
142///
143/// Returns (prefix, payload) on success.
144/// The prefix_length parameter specifies how many bytes of prefix to expect.
145pub fn base58_check_decode(
146    s: &str,
147    prefix_length: usize,
148) -> Result<(Vec<u8>, Vec<u8>), PrimitivesError> {
149    let bin = base58_decode(s)?;
150
151    if bin.len() < prefix_length + 4 {
152        return Err(PrimitivesError::InvalidFormat(
153            "base58check data too short".to_string(),
154        ));
155    }
156
157    let prefix = bin[..prefix_length].to_vec();
158    let payload = bin[prefix_length..bin.len() - 4].to_vec();
159    let checksum = &bin[bin.len() - 4..];
160
161    // Verify checksum
162    let mut hash_input = Vec::with_capacity(prefix.len() + payload.len());
163    hash_input.extend_from_slice(&prefix);
164    hash_input.extend_from_slice(&payload);
165    let expected_checksum = hash256(&hash_input);
166
167    if checksum != &expected_checksum[..4] {
168        return Err(PrimitivesError::ChecksumMismatch);
169    }
170
171    Ok((prefix, payload))
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    // -----------------------------------------------------------------------
179    // Hex utilities
180    // -----------------------------------------------------------------------
181
182    #[test]
183    fn test_hex_encode_decode_roundtrip() {
184        let data = vec![0x00, 0x01, 0xff, 0xab, 0xcd];
185        let hex = to_hex(&data);
186        assert_eq!(hex, "0001ffabcd");
187        let decoded = from_hex(&hex).unwrap();
188        assert_eq!(decoded, data);
189    }
190
191    #[test]
192    fn test_hex_empty() {
193        assert_eq!(to_hex(&[]), "");
194        assert_eq!(from_hex("").unwrap(), Vec::<u8>::new());
195    }
196
197    #[test]
198    fn test_hex_odd_length() {
199        assert!(from_hex("abc").is_err());
200    }
201
202    #[test]
203    fn test_hex_invalid_char() {
204        assert!(from_hex("gg").is_err());
205    }
206
207    // -----------------------------------------------------------------------
208    // Base58
209    // -----------------------------------------------------------------------
210
211    #[test]
212    fn test_base58_encode_known_vector() {
213        // "Hello World" in base58
214        let data = b"Hello World";
215        let encoded = base58_encode(data);
216        assert_eq!(encoded, "JxF12TrwUP45BMd");
217    }
218
219    #[test]
220    fn test_base58_decode_known_vector() {
221        let decoded = base58_decode("JxF12TrwUP45BMd").unwrap();
222        assert_eq!(decoded, b"Hello World");
223    }
224
225    #[test]
226    fn test_base58_roundtrip() {
227        let test_cases: Vec<&[u8]> =
228            vec![b"", &[0], &[0, 0, 0], &[0, 0, 0, 1], b"test", &[0xff; 32]];
229        // Note: empty string decodes/encodes specially
230        for data in test_cases.iter().skip(1) {
231            let encoded = base58_encode(data);
232            let decoded = base58_decode(&encoded).unwrap();
233            assert_eq!(&decoded, data, "Base58 roundtrip failed for {:?}", data);
234        }
235    }
236
237    #[test]
238    fn test_base58_leading_zeros() {
239        // Leading zero bytes should map to '1' characters
240        let data = vec![0, 0, 0, 1];
241        let encoded = base58_encode(&data);
242        assert!(
243            encoded.starts_with("111"),
244            "Expected 3 leading '1's for 3 leading zero bytes, got: {}",
245            encoded
246        );
247        let decoded = base58_decode(&encoded).unwrap();
248        assert_eq!(decoded, data);
249    }
250
251    #[test]
252    fn test_base58_invalid_char() {
253        // '0', 'O', 'I', 'l' are not in Base58 alphabet
254        assert!(base58_decode("0abc").is_err());
255        assert!(base58_decode("Oabc").is_err());
256        assert!(base58_decode("Iabc").is_err());
257        assert!(base58_decode("labc").is_err());
258    }
259
260    // -----------------------------------------------------------------------
261    // Base58Check
262    // -----------------------------------------------------------------------
263
264    #[test]
265    fn test_base58_check_encode_decode_roundtrip() {
266        let payload = vec![0x01, 0x02, 0x03, 0x04];
267        let prefix = vec![0x00];
268
269        let encoded = base58_check_encode(&payload, &prefix);
270        let (dec_prefix, dec_payload) = base58_check_decode(&encoded, 1).unwrap();
271
272        assert_eq!(dec_prefix, prefix);
273        assert_eq!(dec_payload, payload);
274    }
275
276    #[test]
277    fn test_base58_check_bad_checksum() {
278        let payload = vec![0x01, 0x02, 0x03, 0x04];
279        let prefix = vec![0x00];
280
281        let encoded = base58_check_encode(&payload, &prefix);
282
283        // Tamper with the encoded string by changing last character
284        let mut chars: Vec<char> = encoded.chars().collect();
285        let last = chars.len() - 1;
286        chars[last] = if chars[last] == '1' { '2' } else { '1' };
287        let tampered: String = chars.into_iter().collect();
288
289        assert!(
290            base58_check_decode(&tampered, 1).is_err(),
291            "Should fail with tampered checksum"
292        );
293    }
294
295    #[test]
296    fn test_base58_check_wif_known_vector() {
297        // Known WIF: private key = 1
298        // hex private key: 0000000000000000000000000000000000000000000000000000000000000001
299        // mainnet WIF (compressed): KwDiBf89QgGbjEhKnhXJuH7LrciVrZi3qYjgd9M7rFU73sVHnoWn
300        let wif = "KwDiBf89QgGbjEhKnhXJuH7LrciVrZi3qYjgd9M7rFU73sVHnoWn";
301        let result = base58_check_decode(wif, 1);
302        assert!(result.is_ok(), "Known WIF should decode successfully");
303        let (prefix, payload) = result.unwrap();
304        assert_eq!(prefix, vec![0x80]);
305        // payload = 32-byte key + 0x01 compression flag = 33 bytes
306        assert_eq!(payload.len(), 33);
307        // Key should be 0x01 (32 bytes zero-padded)
308        assert_eq!(payload[..31], vec![0u8; 31]);
309        assert_eq!(payload[31], 1);
310        // Compression flag
311        assert_eq!(payload[32], 1);
312    }
313
314    #[test]
315    fn test_base58_check_encode_wif() {
316        // Encode private key 1 as WIF
317        let mut key_data = vec![0u8; 32];
318        key_data[31] = 1;
319        key_data.push(0x01); // compression flag
320        let prefix = vec![0x80];
321
322        let wif = base58_check_encode(&key_data, &prefix);
323        assert_eq!(wif, "KwDiBf89QgGbjEhKnhXJuH7LrciVrZi3qYjgd9M7rFU73sVHnoWn");
324    }
325}