1use crate::encryption::{decrypt, encrypt};
25use blake3;
26use rand::RngCore;
27use serde::{Deserialize, Serialize};
28use thiserror::Error;
29use zeroize::ZeroizeOnDrop;
30
31#[derive(Debug, Error)]
33pub enum TimeLockError {
34 #[error("Invalid time parameter: must be > 0")]
35 InvalidTimeParameter,
36
37 #[error("Decryption failed")]
38 DecryptionFailed,
39
40 #[error("Serialization error: {0}")]
41 SerializationError(String),
42
43 #[error("Invalid ciphertext")]
44 InvalidCiphertext,
45}
46
47pub type TimeLockResult<T> = Result<T, TimeLockError>;
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct TimeParams {
56 pub iterations: u64,
58}
59
60impl TimeParams {
61 pub fn new(iterations: u64) -> Self {
69 Self { iterations }
70 }
71
72 pub fn from_duration_ms(duration_ms: u64) -> Self {
77 Self {
78 iterations: duration_ms * 10_000,
79 }
80 }
81
82 pub fn estimated_delay_ms(&self) -> u64 {
86 self.iterations / 10_000
87 }
88}
89
90#[derive(Debug, Clone, Serialize, Deserialize)]
93pub struct TimeLockCiphertext {
94 ciphertext: Vec<u8>,
96 puzzle_start: [u8; 32],
98 params: TimeParams,
100 nonce: [u8; 12],
102}
103
104impl TimeLockCiphertext {
105 pub fn to_bytes(&self) -> TimeLockResult<Vec<u8>> {
107 crate::codec::encode(self).map_err(|e| TimeLockError::SerializationError(e.to_string()))
108 }
109
110 pub fn from_bytes(bytes: &[u8]) -> TimeLockResult<Self> {
112 crate::codec::decode(bytes).map_err(|e| TimeLockError::SerializationError(e.to_string()))
113 }
114
115 pub fn iterations(&self) -> u64 {
117 self.params.iterations
118 }
119
120 pub fn estimated_time_ms(&self) -> u64 {
122 self.params.estimated_delay_ms()
123 }
124}
125
126pub fn timelock_encrypt(data: &[u8], params: &TimeParams) -> TimeLockResult<TimeLockCiphertext> {
132 if params.iterations == 0 {
133 return Err(TimeLockError::InvalidTimeParameter);
134 }
135
136 let mut puzzle_start = [0u8; 32];
138 rand::thread_rng().fill_bytes(&mut puzzle_start);
139
140 let key = solve_time_lock_puzzle(&puzzle_start, params.iterations);
142
143 let mut nonce = [0u8; 12];
145 rand::thread_rng().fill_bytes(&mut nonce);
146
147 let ciphertext = encrypt(data, &key, &nonce).map_err(|_| TimeLockError::DecryptionFailed)?;
149
150 Ok(TimeLockCiphertext {
151 ciphertext,
152 puzzle_start,
153 params: params.clone(),
154 nonce,
155 })
156}
157
158pub fn timelock_decrypt(ciphertext: &TimeLockCiphertext) -> TimeLockResult<Vec<u8>> {
163 let key = solve_time_lock_puzzle(&ciphertext.puzzle_start, ciphertext.params.iterations);
165
166 decrypt(&ciphertext.ciphertext, &key, &ciphertext.nonce)
168 .map_err(|_| TimeLockError::DecryptionFailed)
169}
170
171fn solve_time_lock_puzzle(start: &[u8; 32], iterations: u64) -> [u8; 32] {
176 let mut current = *start;
177
178 for _ in 0..iterations {
179 current = *blake3::hash(¤t).as_bytes();
180 }
181
182 current
183}
184
185#[derive(Debug, Clone, ZeroizeOnDrop)]
190pub struct TimeLockPuzzle {
191 start: [u8; 32],
193 iterations: u64,
195 #[zeroize(skip)]
197 solution: Option<[u8; 32]>,
198}
199
200impl TimeLockPuzzle {
201 pub fn new(params: &TimeParams) -> TimeLockResult<Self> {
203 if params.iterations == 0 {
204 return Err(TimeLockError::InvalidTimeParameter);
205 }
206
207 let mut start = [0u8; 32];
208 rand::thread_rng().fill_bytes(&mut start);
209
210 Ok(Self {
211 start,
212 iterations: params.iterations,
213 solution: None,
214 })
215 }
216
217 pub fn from_start(start: [u8; 32], iterations: u64) -> TimeLockResult<Self> {
219 if iterations == 0 {
220 return Err(TimeLockError::InvalidTimeParameter);
221 }
222
223 Ok(Self {
224 start,
225 iterations,
226 solution: None,
227 })
228 }
229
230 pub fn solve(&mut self) -> [u8; 32] {
232 if let Some(solution) = self.solution {
233 return solution;
234 }
235
236 let solution = solve_time_lock_puzzle(&self.start, self.iterations);
237 self.solution = Some(solution);
238 solution
239 }
240
241 pub fn start(&self) -> &[u8; 32] {
243 &self.start
244 }
245
246 pub fn iterations(&self) -> u64 {
248 self.iterations
249 }
250
251 pub fn is_solved(&self) -> bool {
253 self.solution.is_some()
254 }
255
256 pub fn solution(&self) -> Option<[u8; 32]> {
258 self.solution
259 }
260}
261
262pub fn timelock_encrypt_with_puzzle(
264 data: &[u8],
265 puzzle: &TimeLockPuzzle,
266) -> TimeLockResult<TimeLockCiphertext> {
267 let key = solve_time_lock_puzzle(puzzle.start(), puzzle.iterations());
269
270 let mut nonce = [0u8; 12];
272 rand::thread_rng().fill_bytes(&mut nonce);
273
274 let ciphertext = encrypt(data, &key, &nonce).map_err(|_| TimeLockError::DecryptionFailed)?;
276
277 Ok(TimeLockCiphertext {
278 ciphertext,
279 puzzle_start: *puzzle.start(),
280 params: TimeParams::new(puzzle.iterations()),
281 nonce,
282 })
283}
284
285#[cfg(test)]
286mod tests {
287 use super::*;
288
289 #[test]
290 fn test_timelock_basic() {
291 let data = b"Time-locked secret message";
292 let params = TimeParams::new(1000);
293
294 let locked = timelock_encrypt(data, ¶ms).unwrap();
295 let decrypted = timelock_decrypt(&locked).unwrap();
296
297 assert_eq!(data, &decrypted[..]);
298 }
299
300 #[test]
301 fn test_timelock_different_iterations() {
302 let data = b"Secret";
303
304 for iterations in [100, 1_000, 10_000] {
305 let params = TimeParams::new(iterations);
306 let locked = timelock_encrypt(data, ¶ms).unwrap();
307
308 assert_eq!(locked.iterations(), iterations);
309
310 let decrypted = timelock_decrypt(&locked).unwrap();
311 assert_eq!(data, &decrypted[..]);
312 }
313 }
314
315 #[test]
316 fn test_timelock_serialization() {
317 let data = b"Serialization test";
318 let params = TimeParams::new(500);
319
320 let locked = timelock_encrypt(data, ¶ms).unwrap();
321
322 let bytes = locked.to_bytes().unwrap();
324 let deserialized = TimeLockCiphertext::from_bytes(&bytes).unwrap();
325
326 let decrypted = timelock_decrypt(&deserialized).unwrap();
328 assert_eq!(data, &decrypted[..]);
329 }
330
331 #[test]
332 fn test_invalid_time_parameter() {
333 let data = b"Test";
334 let params = TimeParams::new(0);
335
336 let result = timelock_encrypt(data, ¶ms);
337 assert!(matches!(result, Err(TimeLockError::InvalidTimeParameter)));
338 }
339
340 #[test]
341 fn test_time_params_from_duration() {
342 let params = TimeParams::from_duration_ms(100);
343 assert_eq!(params.iterations, 1_000_000);
344 assert_eq!(params.estimated_delay_ms(), 100);
345 }
346
347 #[test]
348 fn test_puzzle_basic() {
349 let params = TimeParams::new(1000);
350 let mut puzzle = TimeLockPuzzle::new(¶ms).unwrap();
351
352 assert!(!puzzle.is_solved());
353 assert_eq!(puzzle.solution(), None);
354
355 let solution1 = puzzle.solve();
356 assert!(puzzle.is_solved());
357 assert_eq!(puzzle.solution(), Some(solution1));
358
359 let solution2 = puzzle.solve();
361 assert_eq!(solution1, solution2);
362 }
363
364 #[test]
365 fn test_puzzle_deterministic() {
366 let start = [42u8; 32];
367 let iterations = 1000;
368
369 let mut puzzle1 = TimeLockPuzzle::from_start(start, iterations).unwrap();
370 let mut puzzle2 = TimeLockPuzzle::from_start(start, iterations).unwrap();
371
372 let solution1 = puzzle1.solve();
373 let solution2 = puzzle2.solve();
374
375 assert_eq!(solution1, solution2);
376 }
377
378 #[test]
379 fn test_timelock_with_puzzle() {
380 let data = b"Test data";
381 let params = TimeParams::new(500);
382 let puzzle = TimeLockPuzzle::new(¶ms).unwrap();
383
384 let locked = timelock_encrypt_with_puzzle(data, &puzzle).unwrap();
385 let decrypted = timelock_decrypt(&locked).unwrap();
386
387 assert_eq!(data, &decrypted[..]);
388 }
389
390 #[test]
391 fn test_large_data() {
392 let data = vec![0x42u8; 10_000]; let params = TimeParams::new(1000);
394
395 let locked = timelock_encrypt(&data, ¶ms).unwrap();
396 let decrypted = timelock_decrypt(&locked).unwrap();
397
398 assert_eq!(data, decrypted);
399 }
400
401 #[test]
402 fn test_puzzle_different_iterations_different_solutions() {
403 let start = [1u8; 32];
404
405 let mut puzzle1 = TimeLockPuzzle::from_start(start, 100).unwrap();
406 let mut puzzle2 = TimeLockPuzzle::from_start(start, 200).unwrap();
407
408 let solution1 = puzzle1.solve();
409 let solution2 = puzzle2.solve();
410
411 assert_ne!(solution1, solution2);
412 }
413}