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

kctf_pow/
lib.rs
