Skip to main content

rns_core/
stamp.rs

1//! Stamp validation functions for proof-of-work verification.
2//!
3//! These functions implement the stamp/workblock algorithm used for:
4//! - LXMF message stamps
5//! - Interface discovery verification
6//! - Peering key validation
7
8use alloc::vec::Vec;
9
10use rns_crypto::hkdf::hkdf;
11use rns_crypto::sha256::sha256;
12
13extern crate alloc;
14
15/// Generate a workblock from material with the specified number of HKDF expansion rounds.
16///
17/// Each round produces 256 bytes via HKDF with a salt derived from SHA256(material + msgpack(n)).
18/// Total workblock size = rounds * 256 bytes.
19pub 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
38/// Count leading zero bits in a 32-byte hash.
39pub 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
52/// Calculate the stamp value (number of leading zero bits in SHA256(workblock + stamp)).
53pub 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
61/// Check if a stamp meets the target cost.
62///
63/// Returns true if SHA256(workblock + stamp) has >= target_cost leading zero bits.
64pub 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    // Check: int.from_bytes(result, "big") <= (1 << (256 - target_cost))
71    // Equivalent to: leading_zeros(result) >= target_cost
72    // But Python uses `>` not `>=` for the comparison with target:
73    //   target = 1 << (256 - target_cost)
74    //   int.from_bytes(result) > target -> invalid
75    // So valid means: int.from_bytes(result) <= target
76    // Which is: leading_zeros >= target_cost
77    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; // 10000000 - 0 leading zero bits
93        assert_eq!(leading_zeros(&hash), 0);
94
95        hash[0] = 0x40; // 01000000 - 1 leading zero bit
96        assert_eq!(leading_zeros(&hash), 1);
97
98        hash[0] = 0x01; // 00000001 - 7 leading zero bits
99        assert_eq!(leading_zeros(&hash), 7);
100
101        hash[0] = 0xFF; // 11111111 - 0 leading zero bits
102        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; // 8 (from 0x00) + 0 (from 0x80) = 8 leading zero bits
110        assert_eq!(leading_zeros(&hash), 8);
111
112        hash[1] = 0x01; // 8 (from 0x00) + 7 (from 0x01) = 15 leading zero bits
113        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        // Values from Python:
134        //   packed = b"test data"
135        //   infohash = RNS.Identity.full_hash(packed)
136        //   wb = LXStamper.stamp_workblock(infohash, expand_rounds=20)
137        //   stamp = LXStamper.generate_stamp(infohash, stamp_cost=8, expand_rounds=20)[0]
138        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}