1use alloc::vec::Vec;
9
10use rns_crypto::hkdf::hkdf;
11use rns_crypto::sha256::sha256;
12
13extern crate alloc;
14
15pub fn stamp_workblock(material: &[u8], expand_rounds: u32) -> Vec<u8> {
20 use crate::msgpack::{self, Value};
21
22 let mut workblock = Vec::with_capacity(expand_rounds as usize * 256);
23 for n in 0..expand_rounds {
24 let packed_n = msgpack::pack(&Value::UInt(n as u64));
25 let mut salt_input = Vec::with_capacity(material.len() + packed_n.len());
26 salt_input.extend_from_slice(material);
27 salt_input.extend_from_slice(&packed_n);
28 let salt = sha256(&salt_input);
29
30 let Ok(expanded) = hkdf(256, material, Some(&salt), None) else {
31 break;
32 };
33 workblock.extend_from_slice(&expanded);
34 }
35 workblock
36}
37
38pub fn leading_zeros(hash: &[u8; 32]) -> u32 {
40 let mut count = 0u32;
41 for &byte in hash.iter() {
42 if byte == 0 {
43 count += 8;
44 } else {
45 count += byte.leading_zeros();
46 break;
47 }
48 }
49 count
50}
51
52pub fn stamp_value(workblock: &[u8], stamp: &[u8]) -> u32 {
54 let mut material = Vec::with_capacity(workblock.len() + stamp.len());
55 material.extend_from_slice(workblock);
56 material.extend_from_slice(stamp);
57 let hash = sha256(&material);
58 leading_zeros(&hash)
59}
60
61pub fn stamp_valid(stamp: &[u8], target_cost: u8, workblock: &[u8]) -> bool {
65 let mut material = Vec::with_capacity(workblock.len() + stamp.len());
66 material.extend_from_slice(workblock);
67 material.extend_from_slice(stamp);
68 let result = sha256(&material);
69
70 leading_zeros(&result) >= target_cost as u32
78}
79
80#[cfg(test)]
81mod tests {
82 use super::*;
83
84 #[test]
85 fn test_leading_zeros_all_zero() {
86 assert_eq!(leading_zeros(&[0u8; 32]), 256);
87 }
88
89 #[test]
90 fn test_leading_zeros_first_byte_nonzero() {
91 let mut hash = [0u8; 32];
92 hash[0] = 0x80; assert_eq!(leading_zeros(&hash), 0);
94
95 hash[0] = 0x40; assert_eq!(leading_zeros(&hash), 1);
97
98 hash[0] = 0x01; assert_eq!(leading_zeros(&hash), 7);
100
101 hash[0] = 0xFF; assert_eq!(leading_zeros(&hash), 0);
103 }
104
105 #[test]
106 fn test_leading_zeros_multiple_bytes() {
107 let mut hash = [0u8; 32];
108 hash[0] = 0;
109 hash[1] = 0x80; assert_eq!(leading_zeros(&hash), 8);
111
112 hash[1] = 0x01; assert_eq!(leading_zeros(&hash), 15);
114 }
115
116 #[test]
117 fn test_stamp_workblock_size() {
118 let material = b"test material";
119 let wb = stamp_workblock(material, 20);
120 assert_eq!(wb.len(), 20 * 256);
121 }
122
123 #[test]
124 fn test_stamp_workblock_deterministic() {
125 let material = b"test material";
126 let wb1 = stamp_workblock(material, 5);
127 let wb2 = stamp_workblock(material, 5);
128 assert_eq!(wb1, wb2);
129 }
130
131 #[test]
132 fn test_python_interop_workblock_and_stamp() {
133 let infohash =
139 hex_to_bytes("916f0027a575074ce72a331777c3478d6513f786a591bd892da1a577bf2335f9");
140 let expected_wb_prefix =
141 hex_to_bytes("9e36b853221f04ca1cf54447abce3e9eb47d01d55215414ee5b540eaa796caf2");
142 let stamp =
143 hex_to_bytes("4a1aa3a295482fa9a340b05f2c4779e701b53cd0f158c1bbe559730ae5ff6d17");
144
145 let wb = stamp_workblock(&infohash, 20);
146 assert_eq!(wb.len(), 5120);
147 assert_eq!(&wb[..32], &expected_wb_prefix[..]);
148
149 let value = stamp_value(&wb, &stamp);
150 assert_eq!(value, 8);
151 assert!(stamp_valid(&stamp, 8, &wb));
152 }
153
154 fn hex_to_bytes(s: &str) -> Vec<u8> {
155 (0..s.len())
156 .step_by(2)
157 .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
158 .collect()
159 }
160}