Skip to main content

pow_captcha_core/
lib.rs

1use rand::Rng;
2use sha2::{Digest, Sha256, Sha512};
3
4const DEFAULT_DIFFICULTY: usize = 4;
5
6#[derive(Clone,Debug,Default)]
7pub enum HashAlgorithm {
8    #[default]
9    Sha256,
10    Sha512
11}
12
13/// Represents a PoW challenge (the "puzzle").
14/// This is sent to the client.
15#[derive(Debug, Clone)]
16pub struct Challenge {
17    pub hash_algorithm: HashAlgorithm,
18    /// The random string that needs to be hashed.
19    pub puzzle: String,
20    /// The number of leading zeros required in the hash.
21    pub difficulty: usize,
22}
23
24/// Represents the solution to a PoW challenge, provided by the client.
25#[derive(Debug, Clone)]
26pub struct Solution {
27    /// The nonce found by the client.
28    pub nonce: u64,
29}
30
31/// Generates a new PoW challenge to be sent to the client.
32///
33/// # Arguments
34///
35/// * `difficulty` - An optional `usize` that specifies the number of leading zeros required in the hash.
36///   If `None`, `DEFAULT_DIFFICULTY` (4) is used.
37///
38/// # Returns
39///
40/// A `Challenge` struct containing the puzzle and difficulty.
41pub fn generate_challenge( difficulty: Option<usize>, hash_algorithm: Option<HashAlgorithm>) -> Challenge {
42    let hash_algorithm = hash_algorithm.unwrap_or_default();
43    let difficulty = difficulty.unwrap_or(DEFAULT_DIFFICULTY);
44    let puzzle = generate_random_string(16);
45
46    Challenge {
47        hash_algorithm,
48        puzzle,
49        difficulty,
50    }
51}
52
53/// Verifies a solution provided by a client.
54///
55/// # Arguments
56///
57/// * `challenge` - The original `Challenge` that was sent to the client.
58/// * `solution` - The `Solution` submitted by the client.
59///
60/// # Returns
61///
62/// `true` if the solution is valid, `false` otherwise.
63pub fn verify_solution(challenge: &Challenge, solution: &Solution) -> bool {
64    let hash_hex = compute_hash(&challenge.hash_algorithm, &challenge.puzzle, solution.nonce);
65    hash_hex.starts_with(&"0".repeat(challenge.difficulty))
66}
67
68
69/// Generates a random alphanumeric string of a given length.
70fn generate_random_string(len: usize) -> String {
71    const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
72    let mut rng = rand::thread_rng();
73    (0..len)
74        .map(|_| {
75            let idx = rng.gen_range(0..CHARSET.len());
76            CHARSET[idx] as char
77        })
78        .collect()
79}
80
81/// Computes the hash for a given puzzle and nonce using the specified algorithm.
82fn compute_hash(algorithm: &HashAlgorithm, puzzle: &str, nonce: u64) -> String {
83    match algorithm {
84        HashAlgorithm::Sha256 => {
85            let mut hasher = Sha256::new();
86            hasher.update(puzzle.as_bytes());
87            hasher.update(nonce.to_string().as_bytes());
88            let result = hasher.finalize();
89            hex::encode(result)
90        }
91        HashAlgorithm::Sha512 => {
92            let mut hasher = Sha512::new();
93            hasher.update(puzzle.as_bytes());
94            hasher.update(nonce.to_string().as_bytes());
95            let result = hasher.finalize();
96            hex::encode(result)
97        }
98    }
99}
100
101/// Solves a given PoW challenge.
102/// This function is computationally intensive and is intended to be run on the client side.
103///
104/// # Arguments
105///
106/// * `challenge` - The `Challenge` to solve.
107///
108/// # Returns
109///
110/// The `Solution` containing the correct nonce.
111pub fn solve_challenge(challenge: &Challenge) -> Solution {
112    let mut nonce = 0;
113    loop {
114        let hash_hex = compute_hash(&challenge.hash_algorithm, &challenge.puzzle, nonce);
115        if hash_hex.starts_with(&"0".repeat(challenge.difficulty)) {
116            return Solution { nonce };
117        }
118        nonce += 1;
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn it_generates_challenge() {
128        let challenge = generate_challenge(Some(5), None);
129        assert_eq!(challenge.difficulty, 5);
130        assert_eq!(challenge.puzzle.len(), 16);
131    }
132
133    #[test]
134    fn it_uses_default_difficulty() {
135        let challenge = generate_challenge(None, None);
136        assert_eq!(challenge.difficulty, DEFAULT_DIFFICULTY);
137    }
138
139    #[test]
140    fn it_verifies_a_correct_solution() {
141        let challenge = generate_challenge(Some(4), None);
142        // Simulate a client solving the challenge
143        let solution = solve_challenge(&challenge);
144        
145        assert!(verify_solution(&challenge, &solution));
146    }
147
148    #[test]
149    fn it_rejects_an_incorrect_solution() {
150        let challenge = generate_challenge(Some(4), None);
151        // Simulate a client providing a wrong nonce
152        let incorrect_solution = Solution { nonce: 12345 }; // An arbitrary, likely incorrect nonce
153        
154        // We need to ensure the incorrect nonce is actually incorrect
155        let mut hasher = Sha256::new();
156        hasher.update(challenge.puzzle.as_bytes());
157        hasher.update(incorrect_solution.nonce.to_string().as_bytes());
158        let result = hasher.finalize();
159        let hash_hex = hex::encode(result);
160
161        if !hash_hex.starts_with(&"0".repeat(challenge.difficulty)) {
162             assert!(!verify_solution(&challenge, &incorrect_solution));
163        } else {
164            // In the very unlikely event that our random incorrect nonce was correct,
165            // we just pass the test. This is statistically negligible.
166            println!("Warning: Random incorrect nonce happened to be correct.");
167        }
168    }
169}