kctf_pow/
lib.rs

1//! A library to solve, check, and generate proof-of-work challenges using [kCTF](https://google.github.io/kctf/)'s scheme.
2//!
3//! ```rust
4//! use kctf_pow::KctfPow;
5//!
6//! let pow = KctfPow::new();
7//! // decoding then solving a challenge
8//! let chall = pow.decode_challenge("s.AAAAMg==.H+fPiuL32DPbfN97cpd0nA==").unwrap();
9//! println!("{}", chall.solve());
10//! // decoding then checking a challenge
11//! let chall = pow.decode_challenge("s.AAAAMg==.NDtqORW1uZlIgzszbdMGZA==").unwrap();
12//! let sol = "s.NUH3arymnKB+ysUGdv+67ypDamn4wOKCPORB2ivWE1Yhinam2v4S6q4nAoC5LP97LScdVoq+NuFVF++Win5mNRYZS6bJAs8fk0h8XgvfcC/7JfmFISqeCIo/CIUgIucVAM+eGDjqitRULGXqIOyviJoJjW8DMouMRuJM/3eg/z18kutQHkX0N3sqPeF7Nzkk8S3Bs6aiHUORM30syUKYug==";
13//! assert_eq!(chall.check(sol), Ok(true));
14//! assert_eq!(chall.check("s.asdf"), Ok(false));
15//! // generating a random challenge of difficulty 50
16//! let chall = pow.generate_challenge(50);
17//! println!("{}", chall);
18//! ```
19
20use base64::prelude::*;
21use rand::prelude::*;
22use rug::integer::Order;
23use rug::ops::Pow;
24use rug::Integer;
25use std::convert::TryInto;
26use std::fmt;
27
28const VERSION: &str = "s";
29
30/// The parameters for a proof-of-work challenge.
31///
32/// This contains most of the logic, however [`KctfPow`] and [`Challenge`] should be used instead as they provide a nicer API.
33/// If you want to serialize it to a string, use the [`Display`](std::fmt::Display) implementation.
34#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
35pub struct ChallengeParams {
36    /// The difficulty of the challenge.
37    pub difficulty: u32,
38    /// The starting value of the challenge.
39    pub val: Integer,
40}
41
42/// Squares an integer in-place modulo 2^1279-1
43fn square_mod(n: &mut Integer) {
44    n.square_mut();
45    let high = Integer::from(&*n >> 1279);
46    n.keep_bits_mut(1279);
47    *n += high;
48    if n.get_bit(1279) {
49        n.set_bit(1279, false);
50        *n += 1;
51    }
52}
53
54impl ChallengeParams {
55    /// Decodes a challenge from a string and returns it.
56    ///
57    /// For optimization purposes, the difficulty of the challenge must be able to fit in a [`u32`].
58    /// This shouldn't be an issue, since difficulties that can't fit into a [`u32`] will probably take too long anyways.
59    pub fn decode_challenge(chall_string: &str) -> Result<ChallengeParams, &'static str> {
60        let mut parts = chall_string.split('.');
61        if parts.next() != Some(VERSION) {
62            return Err("Incorrect version");
63        }
64        let data: Vec<_> = parts.collect();
65        if data.len() != 2 {
66            return Err("Incorrect number of parts");
67        }
68        let decoded_data: Vec<_> = data
69            .into_iter()
70            .map(|x| {
71                BASE64_STANDARD
72                    .decode(x)
73                    .map_err(|_| "Parts aren't valid base64")
74            })
75            .collect::<Result<_, _>>()?;
76        let difficulty_bytes = &decoded_data[0];
77        let difficulty: u32 = if difficulty_bytes.len() > 4 {
78            let (first, last) = difficulty_bytes.split_at(difficulty_bytes.len() - 4);
79            // if difficulty is 0-padded to longer than 4 bytes it should still work
80            if first.iter().any(|&x| x != 0) {
81                return Err("Difficulty is too large");
82            }
83            u32::from_be_bytes(last.try_into().unwrap())
84        } else {
85            let mut difficulty_array = [0; 4];
86            difficulty_array[4 - difficulty_bytes.len()..].copy_from_slice(difficulty_bytes);
87            u32::from_be_bytes(difficulty_array)
88        };
89        Ok(Self {
90            val: Integer::from_digits(&decoded_data[1], Order::Msf),
91            difficulty,
92        })
93    }
94
95    /// Generates a random challenge given a difficulty.
96    pub fn generate_challenge(difficulty: u32) -> ChallengeParams {
97        let mut bytes: [u8; 16] = [0; 16];
98        rand::rng().fill(&mut bytes[..]);
99        Self {
100            val: Integer::from_digits(&bytes, Order::Msf),
101            difficulty,
102        }
103    }
104
105    /// Solves a challenge given a proof-of-work system and returns the solution.
106    pub fn solve(mut self) -> String {
107        for _ in 0..self.difficulty {
108            // guaranteed to succeed so ignore the result
109            for _ in 0..1277 {
110                square_mod(&mut self.val);
111            }
112            self.val ^= 1;
113        }
114        format!(
115            "{}.{}",
116            VERSION,
117            BASE64_STANDARD.encode(self.val.to_digits(Order::Msf))
118        )
119    }
120
121    /// Checks a solution to see if it satisfies the challenge under a given proof-of-work system.
122    pub fn check(&self, sol: &str) -> Result<bool, &'static str> {
123        let mut parts = sol.split('.');
124        if parts.next() != Some(VERSION) {
125            return Err("Incorrect version");
126        }
127        let Some(data) = parts.next() else {
128            return Err("Incorrect number of parts");
129        };
130        if parts.next().is_some() {
131            return Err("Incorrect number of parts");
132        }
133        let decoded_data = BASE64_STANDARD
134            .decode(data)
135            .map_err(|_| "Parts aren't valid base64")?;
136        let mut sol_val = Integer::from_digits(&decoded_data, Order::Msf);
137        for _ in 0..self.difficulty {
138            sol_val ^= 1;
139            square_mod(&mut sol_val);
140        }
141        Ok(self.val == sol_val || Integer::from(2).pow(1279) - 1 - &self.val == sol_val)
142    }
143}
144
145impl fmt::Display for ChallengeParams {
146    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
147        write!(
148            fmt,
149            "{}.{}.{}",
150            VERSION,
151            BASE64_STANDARD.encode(self.difficulty.to_be_bytes()),
152            BASE64_STANDARD.encode(self.val.to_digits(Order::Msf))
153        )
154    }
155}
kctf_pow/lib.rs

kctf_pow/
lib.rs
