kctf/
lib.rs

1//! Asynchronous library and cli to solve proof-of-work challenges generated with the kctf scheme.
2//!
3//! ```rust
4//! use kctf::KctfPow;
5//! let challenge = KctfPow::from_challenge("s.AAU5.AACV7mM375HM8wElUbxsknqD").unwrap();
6//!
7//! // Solve one
8//! let solution = challenge.clone().solve();
9//! println!("{}", solution);
10//! assert_eq!(solution, "s.LR15WHZE5YO/8EEY9BF7pdvxiJxwkDi7mdS52bg7eVUdHbAwBVxfahl/qxceccZV2PHkj4wQTQ9Ng837/KD9IWQL4v2GmRyjc5O9MxiAXBtxn7FYjjA2as/17lF2lEtQtABbSEUgxam+sIsdfDJMAUzn4fYsS7vOarXh7iY6ZYknrwt1S8EHyQeYkoTUzkpUIVAuSvl8jExcPzvmuaoM6A==");
11//!
12//! // Verify solution
13//! assert_eq!(challenge.verify(&solution), Ok(true));
14//! assert_eq!(challenge.verify("s.invalid"), Ok(false));
15//!
16//! // Generate a challenge
17//! let challenge = KctfPow::gen_challenge(50);
18//! println!("{}", challenge.serialize_challenge());
19//! ```
20
21use base64::prelude::*;
22use once_cell::sync::Lazy;
23use rand::prelude::*;
24use rug::integer::Order;
25use rug::ops::Pow;
26use rug::Integer;
27use std::str;
28
29const VERSION: &str = "s";
30struct KctfParams {
31    /// The modulus as defined by kCTF, which is 2 ** 1279 - 1
32    pub modulus: Integer,
33    /// The exponent as defined by kCTF, which is (modulus + 1) / 4
34    pub exponent: Integer,
35}
36
37impl KctfParams {
38    fn new() -> Self {
39        let big_num = Integer::from(2).pow(1279);
40        KctfParams {
41            modulus: big_num.clone() - 1,
42            exponent: big_num / 4,
43        }
44    }
45}
46
47static CUSTOM_BASE64: base64::engine::GeneralPurpose = base64::engine::GeneralPurpose::new(
48    &base64::alphabet::STANDARD,
49    base64::engine::GeneralPurposeConfig::new()
50        .with_decode_allow_trailing_bits(true)
51        .with_encode_padding(true)
52        .with_decode_padding_mode(base64::engine::DecodePaddingMode::Indifferent),
53);
54
55static KCTF_PARAMS: Lazy<KctfParams> = Lazy::new(|| KctfParams::new());
56
57#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
58pub enum KctfErrors {
59    /// The given challenge/solution has an unknown version.
60    ///
61    /// Challenges should start with an `s.`
62    UnknownVersion,
63    /// The given challenge/solution does not follow the kCTF format.
64    ///
65    /// Challenges should follow the format `s.<Base64 encoded difficulty>.<Base64 encoded value>`,
66    /// while solutions should follow the format `s.<Base64 encoded value>`
67    FormatError,
68    /// The given challenge/solution parts cannot be properly decoded from base64.
69    ///
70    /// Note that kctf does not require strict following of the base64 rules, like
71    /// padding, but will fail to decode if the parts does not follow the rules
72    /// `[a-zA-Z0-9]+`
73    DecodeError,
74    /// The give challenge has a large difficulty that can't fit in a u32 variable.
75    LargeDifficulty,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
79pub struct KctfPow {
80    /// The difficulty of the challenge
81    pub difficulty: u32,
82    /// The starting value of the challenge
83    pub value: Integer,
84}
85
86impl KctfPow {
87    /// Used to set the challenge difficulty and starting value manually.
88    ///
89    /// This is not recommended for direct use unless the remote instances
90    /// uses a different format that kctf could decode.
91    pub fn from_difficulty_and_value(difficulty: u32, value: Integer) -> Self {
92        KctfPow { difficulty, value }
93    }
94
95    /// Decodes a challenge and returns an instance of `KctfPow` on success
96    /// and a variant of `KctfErrors` if not.
97    pub fn from_challenge(challenge: &str) -> Result<Self, KctfErrors> {
98        let mut challenge_parts = challenge.split('.');
99
100        // Check if kctf version is known
101        if challenge_parts.next() != Some(VERSION) {
102            return Err(KctfErrors::UnknownVersion);
103        }
104
105        let challenge_parts: Vec<&str> = challenge_parts.collect();
106
107        // Check if the number of remaining parts are expected
108        if challenge_parts.len() != 2 {
109            return Err(KctfErrors::FormatError);
110        }
111
112        // Decode all parts
113        let decoded_parts: Vec<Vec<u8>> = challenge_parts
114            .into_iter()
115            .map(|x| CUSTOM_BASE64.decode(x).map_err(|_| KctfErrors::DecodeError))
116            .collect::<Result<_, KctfErrors>>()?;
117
118        let decoded_difficulty = &decoded_parts[0];
119        let decoded_value = &decoded_parts[1];
120
121        let difficulty: u32 = if decoded_difficulty.len() > 4 {
122            if (&decoded_difficulty[..decoded_difficulty.len() - 4])
123                .iter()
124                .any(|&x| x != 0)
125            {
126                return Err(KctfErrors::LargeDifficulty);
127            }
128            u32::from_be_bytes((&decoded_difficulty[..4]).try_into().unwrap())
129        } else {
130            let mut difficulty_array = [0; 4];
131            difficulty_array[4 - decoded_difficulty.len()..].copy_from_slice(&decoded_difficulty);
132            u32::from_be_bytes(difficulty_array)
133        };
134
135        Ok(KctfPow {
136            difficulty,
137            value: Integer::from_digits(decoded_value, Order::Msf),
138        })
139    }
140
141    /// Solves a challenge. This must be called after initialization.
142    pub fn solve(mut self) -> String {
143        for _ in 0..self.difficulty {
144            let _ = self
145                .value
146                .pow_mod_mut(&KCTF_PARAMS.exponent, &KCTF_PARAMS.modulus);
147            self.value ^= 1;
148        }
149
150        format!(
151            "{}.{}",
152            VERSION,
153            CUSTOM_BASE64.encode(self.value.to_digits(Order::Msf))
154        )
155    }
156
157    /// Generates a challenge.
158    pub fn gen_challenge(difficulty: u32) -> Self {
159        let mut bytes: [u8; 16] = [0; 16];
160        thread_rng().fill(&mut bytes[..]);
161        KctfPow {
162            difficulty: difficulty,
163            value: Integer::from_digits(&bytes, Order::Msf),
164        }
165    }
166
167    /// Serializes a challenge after it is generated.
168    pub fn serialize_challenge(&self) -> String {
169        format!(
170            "{}.{}.{}",
171            VERSION,
172            CUSTOM_BASE64.encode(self.difficulty.to_be_bytes()),
173            CUSTOM_BASE64.encode(self.value.to_digits(Order::Msf))
174        )
175    }
176
177    /// Verifies a solution.
178    pub fn verify(&self, solution: &str) -> Result<bool, KctfErrors> {
179        let mut decoded_solution = decode_solution(solution)?;
180
181        for _ in 0..self.difficulty {
182            decoded_solution ^= 1;
183            let _ = decoded_solution.pow_mod_mut(&2.into(), &KCTF_PARAMS.modulus);
184        }
185
186        Ok(self.value == decoded_solution
187            || Integer::from(&KCTF_PARAMS.modulus - &self.value) == decoded_solution)
188    }
189    /// Solves a challenge asynchronously.
190    pub async fn async_solve(mut self) -> String {
191        for _ in 0..self.difficulty {
192            async {
193                let _ = self
194                    .value
195                    .pow_mod_mut(&KCTF_PARAMS.exponent, &KCTF_PARAMS.modulus);
196                self.value ^= 1;
197            }
198            .await
199        }
200
201        format!(
202            "{}.{}",
203            VERSION,
204            CUSTOM_BASE64.encode(self.value.to_digits(Order::Msf))
205        )
206    }
207
208    /// Verifies a solution asynchronously.
209    pub async fn async_verify(&self, solution: &str) -> Result<bool, KctfErrors> {
210        let mut decoded_solution = decode_solution(solution)?;
211
212        for _ in 0..self.difficulty {
213            async {
214                decoded_solution ^= 1;
215                let _ = decoded_solution.pow_mod_mut(&2.into(), &KCTF_PARAMS.modulus);
216            }
217            .await
218        }
219        Ok(self.value == decoded_solution
220            || Integer::from(&KCTF_PARAMS.modulus - &self.value) == decoded_solution)
221    }
222}
223
224/// Decodes a solution.
225///
226/// Not really meant for public consumption, but it could be useful.
227pub fn decode_solution(solution: &str) -> Result<Integer, KctfErrors> {
228    let mut solution_parts = solution.split('.');
229    if solution_parts.next() != Some(VERSION) {
230        return Err(KctfErrors::UnknownVersion);
231    }
232
233    let decoded_solution = match solution_parts.next() {
234        Some(v) => Integer::from_digits(
235            &CUSTOM_BASE64
236                .decode(v)
237                .map_err(|_| KctfErrors::DecodeError)?,
238            Order::Msf,
239        ),
240
241        None => {
242            return Err(KctfErrors::FormatError);
243        }
244    };
245
246    if solution_parts.next().is_some() {
247        return Err(KctfErrors::FormatError);
248    }
249
250    Ok(decoded_solution)
251}