1use crate::serde_utils::{
2 deserialize_32_bytes,
3 deserialize_signature,
4 serialize_32_bytes,
5 serialize_signature
6};
7
8use chrono::Utc;
9use ed25519_dalek::SigningKey;
10use hex;
11use rand;
12use serde::{
13 Deserialize,
14 Serialize
15};
16
17pub const CHALLENGE_DIFFICULTY: u64 = 60_000_000u64;
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct IronShieldChallenge {
31 pub random_nonce: String,
32 pub created_time: i64,
33 pub expiration_time: i64,
34 pub website_id: String,
35 #[serde(
36 serialize_with = "serialize_32_bytes",
37 deserialize_with = "deserialize_32_bytes"
38 )]
39 pub challenge_param: [u8; 32],
40 pub recommended_attempts: u64,
41 #[serde(
42 serialize_with = "serialize_32_bytes",
43 deserialize_with = "deserialize_32_bytes"
44 )]
45 pub public_key: [u8; 32],
46 #[serde(
47 serialize_with = "serialize_signature",
48 deserialize_with = "deserialize_signature"
49 )]
50 pub challenge_signature: [u8; 64],
51}
52
53impl IronShieldChallenge {
54 pub fn new(
69 website_id: String,
70 challenge_param: [u8; 32],
71 private_key: SigningKey,
72 public_key: [u8; 32],
73 ) -> Self {
74 let random_nonce: String = Self::generate_random_nonce();
76
77 let created_time: i64 = Self::generate_created_time();
79 let expiration_time: i64 = created_time + 30_000; let signing_message = crate::crypto::create_signing_message(
83 &random_nonce,
84 created_time,
85 expiration_time,
86 &website_id,
87 &challenge_param,
88 &public_key
89 );
90
91 let challenge_signature: [u8; 64] = crate::crypto::generate_signature(&private_key, &signing_message)
93 .unwrap_or([0u8; 64]);
94
95 Self {
96 random_nonce,
97 created_time,
98 website_id,
99 expiration_time,
100 challenge_param,
101 recommended_attempts: 0, public_key,
103 challenge_signature,
104 }
105 }
106
107 pub fn difficulty_to_challenge_param(difficulty: u64) -> [u8; 32] {
132 if difficulty == 0 {
133 panic!("Difficulty cannot be zero");
134 }
135
136 if difficulty == 1 {
138 return [0xFF; 32];
139 }
140
141 let difficulty_f64: f64 = difficulty as f64;
143 let log2_difficulty: f64 = difficulty_f64.log2();
144
145 let target_exponent: f64 = 256.0 - log2_difficulty;
148
149 if target_exponent <= 0.0 {
150 let mut result: [u8; 32] = [0u8; 32];
152 result[31] = 1;
153 return result;
154 }
155
156 if target_exponent >= 256.0 {
157 return [0xFF; 32];
159 }
160
161 let exponent: usize = target_exponent.round() as usize;
163
164 if exponent >= 256 {
165 return [0xFF; 32];
166 }
167
168 let mut result: [u8; 32] = [0u8; 32];
169
170 let byte_index: usize = (255 - exponent) / 8;
173 let bit_index: usize = 7 - ((255 - exponent) % 8);
174
175 if byte_index < 32 {
176 result[byte_index] = 1u8 << bit_index;
177 } else {
178 result[31] = 1;
180 }
181
182 result
183 }
184
185 pub fn is_expired(&self) -> bool {
188 Utc::now().timestamp_millis() > self.expiration_time
189 }
190
191 pub fn time_until_expiration(&self) -> i64 {
194 self.expiration_time - Utc::now().timestamp_millis()
195 }
196
197 pub fn generate_created_time() -> i64 {
200 Utc::now().timestamp_millis()
201 }
202
203 pub fn generate_random_nonce() -> String {
206 hex::encode(&rand::random::<[u8; 16]>())
207 }
208
209 pub fn recommended_attempts(difficulty: u64) -> u64 {
225 difficulty.saturating_mul(3)
226 }
227
228 pub fn set_recommended_attempts(&mut self, difficulty: u64) {
233 self.recommended_attempts = Self::recommended_attempts(difficulty);
234 }
235
236 pub fn concat_struct(&self) -> String {
246 format!(
247 "{}|{}|{}|{}|{}|{}|{}",
248 self.random_nonce,
249 self.created_time,
250 self.expiration_time,
251 self.website_id,
252 hex::encode(self.challenge_param),
254 hex::encode(self.public_key),
255 hex::encode(self.challenge_signature)
256 )
257 }
258
259 pub fn from_concat_struct(concat_str: &str) -> Result<Self, String> {
277 let parts: Vec<&str> = concat_str.split('|').collect();
278
279 if parts.len() != 7 {
280 return Err(format!("Expected 7 parts, got {}", parts.len()));
281 }
282
283 let random_nonce: String = parts[0].to_string();
284
285 let created_time: i64 = parts[1].parse::<i64>()
286 .map_err(|_| "Failed to parse created_time as i64")?;
287
288 let expiration_time: i64 = parts[2].parse::<i64>()
289 .map_err(|_| "Failed to parse expiration_time as i64")?;
290
291 let website_id: String = parts[3].to_string();
292
293 let challenge_param_bytes: Vec<u8> = hex::decode(parts[4])
294 .map_err(|_| "Failed to decode challenge_params hex string")?;
295 let challenge_param: [u8; 32] = challenge_param_bytes
296 .try_into()
297 .map_err(|_| "Challenge params must be exactly 32 bytes")?;
298
299 let public_key_bytes: Vec<u8> = hex::decode(parts[5])
300 .map_err(|_| "Failed to decode public_key hex string")?;
301 let public_key: [u8; 32] = public_key_bytes.try_into()
302 .map_err(|_| "Public key must be exactly 32 bytes")?;
303
304 let signature_bytes: Vec<u8> = hex::decode(parts[6])
305 .map_err(|_| "Failed to decode challenge_signature hex string")?;
306 let challenge_signature: [u8; 64] = signature_bytes
307 .try_into()
308 .map_err(|_| "Signature must be exactly 64 bytes")?;
309
310 Ok(Self {
311 random_nonce,
312 created_time,
313 expiration_time,
314 website_id,
315 challenge_param,
316 recommended_attempts: 0, public_key,
318 challenge_signature,
319 })
320 }
321
322 pub fn to_base64url_header(&self) -> String {
345 crate::serde_utils::concat_struct_base64url_encode(&self.concat_struct())
346 }
347
348 pub fn from_base64url_header(encoded_header: &str) -> Result<Self, String> {
377 let concat_str: String = crate::serde_utils::concat_struct_base64url_decode(encoded_header.to_string())?;
379
380 Self::from_concat_struct(&concat_str)
382 }
383}
384
385#[cfg(test)]
386mod tests {
387 use super::*;
388
389 #[test]
390 fn test_difficulty_to_challenge_param_basic_cases() {
391 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(1);
393 assert_eq!(challenge_param, [0xFF; 32]);
394
395 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(2);
397 let expected: [u8; 32] = {
398 let mut arr: [u8; 32] = [0x00; 32];
399 arr[0] = 0x80; arr
401 };
402 assert_eq!(challenge_param, expected);
403
404 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(4);
405 let expected: [u8; 32] = {
406 let mut arr: [u8; 32] = [0x00; 32];
407 arr[0] = 0x40; arr
409 };
410 assert_eq!(challenge_param, expected);
411
412 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(256);
413 let expected: [u8; 32] = {
414 let mut arr: [u8; 32] = [0x00; 32];
415 arr[0] = 0x01; arr
417 };
418 assert_eq!(challenge_param, expected);
419 }
420
421 #[test]
422 fn test_difficulty_to_challenge_param_realistic_range() {
423 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(10_000);
427 assert_eq!(challenge_param[0], 0x00);
429 assert_eq!(challenge_param[1], 0x08); let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(50_000);
433 assert_eq!(challenge_param[0], 0x00);
434 assert_eq!(challenge_param[1], 0x01); let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(100_000);
438 assert_eq!(challenge_param[0], 0x00);
439 assert_eq!(challenge_param[1], 0x00);
440 assert_eq!(challenge_param[2], 0x80); let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(1_000_000);
444 assert_eq!(challenge_param[0], 0x00);
445 assert_eq!(challenge_param[1], 0x00);
446 assert_eq!(challenge_param[2], 0x10); let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(10_000_000);
450 assert_eq!(challenge_param[0], 0x00);
451 assert_eq!(challenge_param[1], 0x00);
452 assert_eq!(challenge_param[2], 0x02); }
454
455 #[test]
456 fn test_difficulty_to_challenge_param_ordering() {
457 let difficulties: [u64; 9] = [1000, 5000, 10_000, 50_000, 100_000, 500_000, 1_000_000, 5_000_000, 10_000_000];
459 let mut challenge_params = Vec::new();
460
461 for &difficulty in &difficulties {
462 challenge_params.push(IronShieldChallenge::difficulty_to_challenge_param(difficulty));
463 }
464
465 for i in 1..challenge_params.len() {
467 assert!(
468 challenge_params[i-1] > challenge_params[i],
469 "Challenge param for difficulty {} should be larger than for difficulty {}",
470 difficulties[i-1], difficulties[i]
471 );
472 }
473 }
474
475 #[test]
476 fn test_difficulty_to_challenge_param_precision() {
477 let base_difficulty: u64 = 100_000;
479 let base_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(base_difficulty);
480
481 let similar_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(100_001);
483
484 assert!(base_param >= similar_param); let much_different_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(200_000);
490 assert!(base_param > much_different_param);
491
492 let big_different_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(400_000);
494 assert!(much_different_param > big_different_param);
495 }
496
497 #[test]
498 fn test_difficulty_to_challenge_param_powers_of_10() {
499 let difficulties: [u64; 6] = [10, 100, 1_000, 10_000, 100_000, 1_000_000];
501
502 for &difficulty in &difficulties {
503 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(difficulty);
504
505 assert_ne!(challenge_param, [0x00; 32]);
507 assert_ne!(challenge_param, [0xFF; 32]);
508
509 let leading_zero_bytes: usize = challenge_param.iter().take_while(|&&b| b == 0).count();
511 assert!(leading_zero_bytes < 32, "Too many leading zero bytes for difficulty {}", difficulty);
512
513 assert!(leading_zero_bytes < 28, "Challenge param too small for difficulty {}", difficulty);
515 }
516 }
517
518 #[test]
519 fn test_difficulty_to_challenge_param_mathematical_properties() {
520 let d1: u64 = 50_000;
525 let d2: u64 = 100_000; let param1: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(d1);
528 let param2: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(d2);
529
530 let val1: u128 = u128::from_be_bytes(param1[0..16].try_into().unwrap());
532 let val2: u128 = u128::from_be_bytes(param2[0..16].try_into().unwrap());
533
534 let ratio: f64 = val1 as f64 / val2 as f64;
536 assert!(ratio > 1.8 && ratio < 2.2, "Ratio should be close to 2.0, got {}", ratio);
537 }
538
539 #[test]
540 fn test_difficulty_to_challenge_param_edge_cases() {
541 let result = std::panic::catch_unwind(|| {
543 IronShieldChallenge::difficulty_to_challenge_param(0);
544 });
545 assert!(result.is_err());
546
547 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(u64::MAX);
549 assert_ne!(challenge_param, [0xFF; 32]);
550 assert_ne!(challenge_param, [0; 32]);
551
552 let high_difficulty: u64 = 1u64 << 40; let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(high_difficulty);
555 assert_ne!(challenge_param, [0; 32]);
556 assert_ne!(challenge_param, [0xFF; 32]);
557 }
558
559 #[test]
560 fn test_difficulty_to_challenge_param_consistency() {
561 let test_difficulties: [u64; 13] = [
563 10_000, 25_000, 50_000, 75_000, 100_000,
564 250_000, 500_000, 750_000, 1_000_000,
565 2_500_000, 5_000_000, 7_500_000, 10_000_000
566 ];
567
568 for &difficulty in &test_difficulties {
569 let param1: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(difficulty);
570 let param2: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(difficulty);
571 assert_eq!(param1, param2, "Function should be deterministic for difficulty {}", difficulty);
572
573 assert_ne!(param1, [0x00; 32]);
575 assert_ne!(param1, [0xFF; 32]);
576 }
577 }
578
579 #[test]
580 fn test_recommended_attempts() {
581 assert_eq!(IronShieldChallenge::recommended_attempts(1000), 3000);
583 assert_eq!(IronShieldChallenge::recommended_attempts(50000), 150000);
584 assert_eq!(IronShieldChallenge::recommended_attempts(0), 0);
585
586 assert_eq!(IronShieldChallenge::recommended_attempts(u64::MAX), u64::MAX);
588
589 assert_eq!(IronShieldChallenge::recommended_attempts(10_000), 30_000);
591 assert_eq!(IronShieldChallenge::recommended_attempts(1_000_000), 3_000_000);
592 }
593
594 #[test]
595 fn test_base64url_header_encoding_roundtrip() {
596 let dummy_key = SigningKey::from_bytes(&[0u8; 32]);
598 let challenge = IronShieldChallenge::new(
599 "test_website".to_string(),
600 [0x12; 32],
601 dummy_key,
602 [0x34; 32],
603 );
604
605 let encoded: String = challenge.to_base64url_header();
607 let decoded: IronShieldChallenge = IronShieldChallenge::from_base64url_header(&encoded).unwrap();
608
609 assert_eq!(challenge.random_nonce, decoded.random_nonce);
611 assert_eq!(challenge.created_time, decoded.created_time);
612 assert_eq!(challenge.expiration_time, decoded.expiration_time);
613 assert_eq!(challenge.website_id, decoded.website_id);
614 assert_eq!(challenge.challenge_param, decoded.challenge_param);
615 assert_eq!(challenge.public_key, decoded.public_key);
616 assert_eq!(challenge.challenge_signature, decoded.challenge_signature);
617 }
618
619 #[test]
620 fn test_base64url_header_invalid_data() {
621 let result: Result<IronShieldChallenge, String> = IronShieldChallenge::from_base64url_header("invalid-base64!");
623 assert!(result.is_err());
624 assert!(result.unwrap_err().contains("Base64 decode error"));
625
626 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
628 let invalid_format: String = URL_SAFE_NO_PAD.encode(b"not|enough|parts");
629 let result: Result<IronShieldChallenge, String> = IronShieldChallenge::from_base64url_header(&invalid_format);
630 assert!(result.is_err());
631 assert!(result.unwrap_err().contains("Expected 7 parts"));
632 }
633
634 #[test]
635 fn test_difficulty_range_boundaries() {
636 let min_difficulty = 10_000;
638 let max_difficulty = 10_000_000;
639
640 let min_param = IronShieldChallenge::difficulty_to_challenge_param(min_difficulty);
641 let max_param = IronShieldChallenge::difficulty_to_challenge_param(max_difficulty);
642
643 assert!(min_param > max_param);
645
646 assert_ne!(min_param, [0x00; 32]);
648 assert_ne!(min_param, [0xFF; 32]);
649 assert_ne!(max_param, [0x00; 32]);
650 assert_ne!(max_param, [0xFF; 32]);
651
652 let below_min = IronShieldChallenge::difficulty_to_challenge_param(9_999);
654 let above_max = IronShieldChallenge::difficulty_to_challenge_param(10_000_001);
655
656 assert!(below_min >= min_param); assert!(above_max <= max_param); }
660
661 #[test]
662 fn test_from_concat_struct_edge_cases() {
663 let valid_32_byte_hex = "0".repeat(64); let valid_64_byte_hex = "0".repeat(128); assert_eq!(valid_32_byte_hex.len(), 64, "32-byte hex string should be exactly 64 characters");
670 assert_eq!(valid_64_byte_hex.len(), 128, "64-byte hex string should be exactly 128 characters");
671
672 let input = format!("test_nonce|1000000|1030000|test_website|{}|{}|{}",
673 valid_32_byte_hex, valid_32_byte_hex, valid_64_byte_hex);
674 let result = IronShieldChallenge::from_concat_struct(&input);
675
676 if result.is_err() {
677 panic!("Expected success but got error: {}", result.unwrap_err());
678 }
679
680 let parsed = result.unwrap();
681 assert_eq!(parsed.random_nonce, "test_nonce");
682 assert_eq!(parsed.created_time, 1000000);
683 assert_eq!(parsed.expiration_time, 1030000);
684 assert_eq!(parsed.website_id, "test_website");
685 assert_eq!(parsed.challenge_param, [0u8; 32]);
686 assert_eq!(parsed.public_key, [0u8; 32]);
687 assert_eq!(parsed.challenge_signature, [0u8; 64]);
688
689 let all_f_32_hex = "f".repeat(64); let all_f_64_hex = "f".repeat(128); assert_eq!(all_f_32_hex.len(), 64, "All F's 32-byte hex string should be exactly 64 characters");
693 assert_eq!(all_f_64_hex.len(), 128, "All F's 64-byte hex string should be exactly 128 characters");
694
695 let input = format!("max_nonce|9999999|9999999|max_website|{}|{}|{}",
696 all_f_32_hex, all_f_32_hex, all_f_64_hex);
697 let result = IronShieldChallenge::from_concat_struct(&input);
698
699 if result.is_err() {
700 panic!("Expected success but got error: {}", result.unwrap_err());
701 }
702
703 let parsed = result.unwrap();
704 assert_eq!(parsed.random_nonce, "max_nonce");
705 assert_eq!(parsed.created_time, 9999999);
706 assert_eq!(parsed.expiration_time, 9999999);
707 assert_eq!(parsed.website_id, "max_website");
708 assert_eq!(parsed.challenge_param, [0xffu8; 32]);
709 assert_eq!(parsed.public_key, [0xffu8; 32]);
710 assert_eq!(parsed.challenge_signature, [0xffu8; 64]);
711 }
712}