enigma_node_registry/
pow.rs

1#[cfg(feature = "pow")]
2use std::collections::HashMap;
3#[cfg(feature = "pow")]
4use std::sync::Arc;
5#[cfg(feature = "pow")]
6use std::time::{Duration, Instant};
7
8use serde::{Deserialize, Serialize};
9
10#[cfg(feature = "pow")]
11use rand::RngCore;
12#[cfg(feature = "pow")]
13use sha2::{Digest, Sha256};
14#[cfg(feature = "pow")]
15use tokio::sync::Mutex;
16
17use crate::config::PowConfig;
18use crate::error::{RegistryError, RegistryResult};
19
20#[cfg(feature = "pow")]
21#[derive(Clone, Debug, Serialize, Deserialize)]
22pub struct PowChallenge {
23    pub challenge: String,
24    pub difficulty: u8,
25    pub expires_ms: u64,
26}
27
28#[cfg(not(feature = "pow"))]
29#[derive(Clone, Debug, Serialize, Deserialize)]
30pub struct PowChallenge {
31    pub challenge: String,
32    pub difficulty: u8,
33    pub expires_ms: u64,
34}
35
36#[cfg(feature = "pow")]
37struct StoredChallenge {
38    challenge: String,
39    difficulty: u8,
40    expires_at: Instant,
41}
42
43#[cfg(feature = "pow")]
44#[derive(Clone)]
45pub struct PowManager {
46    enabled: bool,
47    difficulty: u8,
48    ttl: Duration,
49    inner: Arc<Mutex<HashMap<String, StoredChallenge>>>,
50}
51
52#[cfg(feature = "pow")]
53impl PowManager {
54    pub fn new(cfg: PowConfig) -> Self {
55        PowManager {
56            enabled: cfg.enabled,
57            difficulty: cfg.difficulty,
58            ttl: Duration::from_secs(cfg.ttl_seconds),
59            inner: Arc::new(Mutex::new(HashMap::new())),
60        }
61    }
62
63    pub fn enabled(&self) -> bool {
64        self.enabled
65    }
66
67    pub async fn issue(&self) -> RegistryResult<PowChallenge> {
68        if !self.enabled {
69            return Err(RegistryError::FeatureDisabled("pow".to_string()));
70        }
71        let mut bytes = [0u8; 16];
72        rand::thread_rng().fill_bytes(&mut bytes);
73        let challenge = hex::encode(bytes);
74        let expires_at = Instant::now() + self.ttl;
75        let expires_ms = expires_at
76            .checked_duration_since(Instant::now())
77            .map(|d| d.as_millis() as u64 + current_time_ms())
78            .unwrap_or_else(current_time_ms);
79        let stored = StoredChallenge {
80            challenge: challenge.clone(),
81            difficulty: self.difficulty,
82            expires_at,
83        };
84        let mut guard = self.inner.lock().await;
85        guard.insert(challenge.clone(), stored);
86        Ok(PowChallenge {
87            challenge,
88            difficulty: self.difficulty,
89            expires_ms,
90        })
91    }
92
93    pub async fn verify_header(&self, value: Option<&str>) -> RegistryResult<()> {
94        if !self.enabled {
95            return Ok(());
96        }
97        let Some(v) = value else {
98            return Err(RegistryError::PowRequired);
99        };
100        let mut parts = v.splitn(2, ':');
101        let challenge = parts.next().unwrap_or_default();
102        let solution = parts.next().unwrap_or_default();
103        if challenge.is_empty() || solution.is_empty() {
104            return Err(RegistryError::PowRequired);
105        }
106        let mut guard = self.inner.lock().await;
107        let Some(stored) = guard.remove(challenge) else {
108            return Err(RegistryError::PowRequired);
109        };
110        if Instant::now() > stored.expires_at {
111            return Err(RegistryError::PowRequired);
112        }
113        let mut hasher = Sha256::new();
114        hasher.update(stored.challenge.as_bytes());
115        hasher.update(solution.as_bytes());
116        let digest = hasher.finalize();
117        if !meets_difficulty(&digest, stored.difficulty) {
118            return Err(RegistryError::PowRequired);
119        }
120        Ok(())
121    }
122}
123
124#[cfg(feature = "pow")]
125fn meets_difficulty(digest: &[u8], difficulty: u8) -> bool {
126    let zero_bytes = (difficulty / 8) as usize;
127    let zero_bits = (difficulty % 8) as usize;
128    if digest.len() <= zero_bytes {
129        return false;
130    }
131    if digest.iter().take(zero_bytes).any(|b| *b != 0) {
132        return false;
133    }
134    if zero_bits == 0 {
135        return true;
136    }
137    let next = digest[zero_bytes];
138    next.leading_zeros() as usize >= zero_bits
139}
140
141#[cfg(feature = "pow")]
142fn current_time_ms() -> u64 {
143    use std::time::{SystemTime, UNIX_EPOCH};
144    match SystemTime::now().duration_since(UNIX_EPOCH) {
145        Ok(duration) => duration.as_millis() as u64,
146        Err(_) => 0,
147    }
148}
149
150#[cfg(not(feature = "pow"))]
151#[derive(Clone)]
152pub struct PowManager;
153
154#[cfg(not(feature = "pow"))]
155impl PowManager {
156    pub fn new(_cfg: PowConfig) -> Self {
157        PowManager
158    }
159
160    pub fn enabled(&self) -> bool {
161        false
162    }
163
164    pub async fn issue(&self) -> RegistryResult<PowChallenge> {
165        Err(RegistryError::FeatureDisabled("pow".to_string()))
166    }
167
168    pub async fn verify_header(&self, _value: Option<&str>) -> RegistryResult<()> {
169        Ok(())
170    }
171}
enigma_node_registry/pow.rs

enigma_node_registry/
pow.rs
