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 = 200_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 difficulty: u64,
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 challenge_param: [u8; 32] = Self::difficulty_to_challenge_param(difficulty);
82
83 let signing_message = crate::crypto::create_signing_message(
85 &random_nonce,
86 created_time,
87 expiration_time,
88 &website_id,
89 &challenge_param,
90 &public_key
91 );
92
93 let challenge_signature: [u8; 64] = crate::crypto::generate_signature(&private_key, &signing_message)
95 .unwrap_or([0u8; 64]);
96
97 Self {
98 random_nonce,
99 created_time,
100 website_id,
101 expiration_time,
102 challenge_param,
103 recommended_attempts: Self::recommended_attempts(difficulty),
104 public_key,
105 challenge_signature,
106 }
107 }
108
109 pub fn difficulty_to_challenge_param(difficulty: u64) -> [u8; 32] {
134 if difficulty == 0 {
135 panic!("Difficulty cannot be zero");
136 }
137
138 if difficulty == 1 {
140 return [0xFF; 32];
141 }
142
143 let difficulty_f64: f64 = difficulty as f64;
145 let log2_difficulty: f64 = difficulty_f64.log2();
146
147 let target_exponent: f64 = 256.0 - log2_difficulty;
150
151 if target_exponent <= 0.0 {
152 let mut result: [u8; 32] = [0u8; 32];
154 result[31] = 1;
155 return result;
156 }
157
158 if target_exponent >= 256.0 {
159 return [0xFF; 32];
161 }
162
163 let exponent: usize = target_exponent.round() as usize;
165
166 if exponent >= 256 {
167 return [0xFF; 32];
168 }
169
170 let mut result: [u8; 32] = [0u8; 32];
171
172 let byte_index: usize = (255 - exponent) / 8;
175 let bit_index: usize = 7 - ((255 - exponent) % 8);
176
177 if byte_index < 32 {
178 result[byte_index] = 1u8 << bit_index;
179 } else {
180 result[31] = 1;
182 }
183
184 result
185 }
186
187 pub fn is_expired(&self) -> bool {
191 Utc::now().timestamp_millis() > self.expiration_time
192 }
193
194 pub fn time_until_expiration(&self) -> i64 {
197 self.expiration_time - Utc::now().timestamp_millis()
198 }
199
200 pub fn generate_created_time() -> i64 {
203 Utc::now().timestamp_millis()
204 }
205
206 pub fn generate_random_nonce() -> String {
209 hex::encode(&rand::random::<[u8; 16]>())
210 }
211
212 pub fn recommended_attempts(difficulty: u64) -> u64 {
228 difficulty.saturating_mul(2)
229 }
230
231 pub fn concat_struct(&self) -> String {
241 format!(
242 "{}|{}|{}|{}|{}|{}|{}|{}",
243 self.random_nonce,
244 self.created_time,
245 self.expiration_time,
246 self.website_id,
247 hex::encode(self.challenge_param),
249 self.recommended_attempts,
250 hex::encode(self.public_key),
251 hex::encode(self.challenge_signature)
252 )
253 }
254
255 pub fn from_concat_struct(concat_str: &str) -> Result<Self, String> {
273 let parts: Vec<&str> = concat_str.split('|').collect();
274
275 if parts.len() != 8 {
276 return Err(format!("Expected 8 parts, got {}", parts.len()));
277 }
278
279 let random_nonce: String = parts[0].to_string();
280
281 let created_time: i64 = parts[1].parse::<i64>()
282 .map_err(|_| "Failed to parse created_time as i64")?;
283
284 let expiration_time: i64 = parts[2].parse::<i64>()
285 .map_err(|_| "Failed to parse expiration_time as i64")?;
286
287 let website_id: String = parts[3].to_string();
288
289 let challenge_param_bytes: Vec<u8> = hex::decode(parts[4])
290 .map_err(|_| "Failed to decode challenge_params hex string")?;
291 let challenge_param: [u8; 32] = challenge_param_bytes
292 .try_into()
293 .map_err(|_| "Challenge params must be exactly 32 bytes")?;
294
295 let recommended_attempts: u64 = parts[5].parse::<u64>()
296 .map_err(|_| "Failed to parse recommended_attempts as u64")?;
297
298 let public_key_bytes: Vec<u8> = hex::decode(parts[6])
299 .map_err(|_| "Failed to decode public_key hex string")?;
300 let public_key: [u8; 32] = public_key_bytes.try_into()
301 .map_err(|_| "Public key must be exactly 32 bytes")?;
302
303 let signature_bytes: Vec<u8> = hex::decode(parts[7])
304 .map_err(|_| "Failed to decode challenge_signature hex string")?;
305 let challenge_signature: [u8; 64] = signature_bytes
306 .try_into()
307 .map_err(|_| "Signature must be exactly 64 bytes")?;
308
309 Ok(Self {
310 random_nonce,
311 created_time,
312 expiration_time,
313 website_id,
314 challenge_param,
315 recommended_attempts,
316 public_key,
317 challenge_signature,
318 })
319 }
320
321 pub fn to_base64url_header(&self) -> String {
344 crate::serde_utils::concat_struct_base64url_encode(&self.concat_struct())
345 }
346
347 pub fn from_base64url_header(encoded_header: &str) -> Result<Self, String> {
376 let concat_str: String = crate::serde_utils::concat_struct_base64url_decode(encoded_header.to_string())?;
378
379 Self::from_concat_struct(&concat_str)
381 }
382}
383
384#[cfg(test)]
385mod tests {
386 use super::*;
387
388 #[test]
389 fn test_difficulty_to_challenge_param_basic_cases() {
390 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(1);
392 assert_eq!(challenge_param, [0xFF; 32]);
393
394 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(2);
396 let expected: [u8; 32] = {
397 let mut arr: [u8; 32] = [0x00; 32];
398 arr[0] = 0x80; arr
400 };
401 assert_eq!(challenge_param, expected);
402
403 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(4);
404 let expected: [u8; 32] = {
405 let mut arr: [u8; 32] = [0x00; 32];
406 arr[0] = 0x40; arr
408 };
409 assert_eq!(challenge_param, expected);
410
411 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(256);
412 let expected: [u8; 32] = {
413 let mut arr: [u8; 32] = [0x00; 32];
414 arr[0] = 0x01; arr
416 };
417 assert_eq!(challenge_param, expected);
418 }
419
420 #[test]
421 fn test_difficulty_to_challenge_param_realistic_range() {
422 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(10_000);
426 assert_eq!(challenge_param[0], 0x00);
428 assert_eq!(challenge_param[1], 0x08); let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(50_000);
432 assert_eq!(challenge_param[0], 0x00);
433 assert_eq!(challenge_param[1], 0x01); let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(100_000);
437 assert_eq!(challenge_param[0], 0x00);
438 assert_eq!(challenge_param[1], 0x00);
439 assert_eq!(challenge_param[2], 0x80); let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(1_000_000);
443 assert_eq!(challenge_param[0], 0x00);
444 assert_eq!(challenge_param[1], 0x00);
445 assert_eq!(challenge_param[2], 0x10); let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(10_000_000);
449 assert_eq!(challenge_param[0], 0x00);
450 assert_eq!(challenge_param[1], 0x00);
451 assert_eq!(challenge_param[2], 0x02); }
453
454 #[test]
455 fn test_difficulty_to_challenge_param_ordering() {
456 let difficulties: [u64; 9] = [1000, 5000, 10_000, 50_000, 100_000, 500_000, 1_000_000, 5_000_000, 10_000_000];
458 let mut challenge_params = Vec::new();
459
460 for &difficulty in &difficulties {
461 challenge_params.push(IronShieldChallenge::difficulty_to_challenge_param(difficulty));
462 }
463
464 for i in 1..challenge_params.len() {
466 assert!(
467 challenge_params[i-1] > challenge_params[i],
468 "Challenge param for difficulty {} should be larger than for difficulty {}",
469 difficulties[i-1], difficulties[i]
470 );
471 }
472 }
473
474 #[test]
475 fn test_difficulty_to_challenge_param_precision() {
476 let base_difficulty: u64 = 100_000;
478 let base_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(base_difficulty);
479
480 let similar_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(100_001);
482
483 assert!(base_param >= similar_param); let much_different_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(200_000);
489 assert!(base_param > much_different_param);
490
491 let big_different_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(400_000);
493 assert!(much_different_param > big_different_param);
494 }
495
496 #[test]
497 fn test_difficulty_to_challenge_param_powers_of_10() {
498 let difficulties: [u64; 6] = [10, 100, 1_000, 10_000, 100_000, 1_000_000];
500
501 for &difficulty in &difficulties {
502 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(difficulty);
503
504 assert_ne!(challenge_param, [0x00; 32]);
506 assert_ne!(challenge_param, [0xFF; 32]);
507
508 let leading_zero_bytes: usize = challenge_param.iter().take_while(|&&b| b == 0).count();
510 assert!(leading_zero_bytes < 32, "Too many leading zero bytes for difficulty {}", difficulty);
511
512 assert!(leading_zero_bytes < 28, "Challenge param too small for difficulty {}", difficulty);
514 }
515 }
516
517 #[test]
518 fn test_difficulty_to_challenge_param_mathematical_properties() {
519 let d1: u64 = 50_000;
524 let d2: u64 = 100_000; let param1: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(d1);
527 let param2: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(d2);
528
529 let val1: u128 = u128::from_be_bytes(param1[0..16].try_into().unwrap());
531 let val2: u128 = u128::from_be_bytes(param2[0..16].try_into().unwrap());
532
533 let ratio: f64 = val1 as f64 / val2 as f64;
535 assert!(ratio > 1.8 && ratio < 2.2, "Ratio should be close to 2.0, got {}", ratio);
536 }
537
538 #[test]
539 fn test_difficulty_to_challenge_param_edge_cases() {
540 let result = std::panic::catch_unwind(|| {
542 IronShieldChallenge::difficulty_to_challenge_param(0);
543 });
544 assert!(result.is_err());
545
546 let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(u64::MAX);
548 assert_ne!(challenge_param, [0xFF; 32]);
549 assert_ne!(challenge_param, [0; 32]);
550
551 let high_difficulty: u64 = 1u64 << 40; let challenge_param: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(high_difficulty);
554 assert_ne!(challenge_param, [0; 32]);
555 assert_ne!(challenge_param, [0xFF; 32]);
556 }
557
558 #[test]
559 fn test_difficulty_to_challenge_param_consistency() {
560 let test_difficulties: [u64; 13] = [
562 10_000, 25_000, 50_000, 75_000, 100_000,
563 250_000, 500_000, 750_000, 1_000_000,
564 2_500_000, 5_000_000, 7_500_000, 10_000_000
565 ];
566
567 for &difficulty in &test_difficulties {
568 let param1: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(difficulty);
569 let param2: [u8; 32] = IronShieldChallenge::difficulty_to_challenge_param(difficulty);
570 assert_eq!(param1, param2, "Function should be deterministic for difficulty {}", difficulty);
571
572 assert_ne!(param1, [0x00; 32]);
574 assert_ne!(param1, [0xFF; 32]);
575 }
576 }
577
578 #[test]
579 fn test_recommended_attempts() {
580 assert_eq!(IronShieldChallenge::recommended_attempts(1000), 2000);
582 assert_eq!(IronShieldChallenge::recommended_attempts(50000), 100000);
583 assert_eq!(IronShieldChallenge::recommended_attempts(0), 0);
584
585 assert_eq!(IronShieldChallenge::recommended_attempts(u64::MAX), u64::MAX);
587
588 assert_eq!(IronShieldChallenge::recommended_attempts(10_000), 20_000);
590 assert_eq!(IronShieldChallenge::recommended_attempts(1_000_000), 2_000_000);
591 }
592
593 #[test]
594 fn test_base64url_header_encoding_roundtrip() {
595 let private_key = SigningKey::from_bytes(&[0; 32]);
597 let public_key = private_key.verifying_key().to_bytes();
598 let original_challenge = IronShieldChallenge::new(
599 "test-site".to_string(),
600 100_000,
601 private_key,
602 public_key,
603 );
604
605 let encoded = original_challenge.to_base64url_header();
607 let decoded_challenge = IronShieldChallenge::from_base64url_header(&encoded)
608 .expect("Failed to decode header");
609
610 assert_eq!(original_challenge.random_nonce, decoded_challenge.random_nonce);
612 assert_eq!(original_challenge.created_time, decoded_challenge.created_time);
613 assert_eq!(original_challenge.expiration_time, decoded_challenge.expiration_time);
614 assert_eq!(original_challenge.website_id, decoded_challenge.website_id);
615 assert_eq!(original_challenge.challenge_param, decoded_challenge.challenge_param);
616 assert_eq!(original_challenge.public_key, decoded_challenge.public_key);
617 assert_eq!(original_challenge.challenge_signature, decoded_challenge.challenge_signature);
618 }
619
620 #[test]
621 fn test_base64url_header_invalid_data() {
622 let result: Result<IronShieldChallenge, String> = IronShieldChallenge::from_base64url_header("invalid-base64!");
624 assert!(result.is_err());
625 assert!(result.unwrap_err().contains("Base64 decode error"));
626
627 use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
629 let invalid_format: String = URL_SAFE_NO_PAD.encode(b"not_enough_parts");
630 let result: Result<IronShieldChallenge, String> = IronShieldChallenge::from_base64url_header(&invalid_format);
631 assert!(result.is_err());
632 assert!(result.unwrap_err().contains("Expected 8 parts"));
633 }
634
635 #[test]
636 fn test_difficulty_range_boundaries() {
637 let min_difficulty = 10_000;
639 let max_difficulty = 10_000_000;
640
641 let min_param = IronShieldChallenge::difficulty_to_challenge_param(min_difficulty);
642 let max_param = IronShieldChallenge::difficulty_to_challenge_param(max_difficulty);
643
644 assert!(min_param > max_param);
646
647 assert_ne!(min_param, [0x00; 32]);
649 assert_ne!(min_param, [0xFF; 32]);
650 assert_ne!(max_param, [0x00; 32]);
651 assert_ne!(max_param, [0xFF; 32]);
652
653 let below_min = IronShieldChallenge::difficulty_to_challenge_param(9_999);
655 let above_max = IronShieldChallenge::difficulty_to_challenge_param(10_000_001);
656
657 assert!(below_min >= min_param); assert!(above_max <= max_param); }
661
662 #[test]
663 fn test_from_concat_struct_edge_cases() {
664 let valid_32_byte_hex = "0000000000000000000000000000000000000000000000000000000000000000";
666 assert_eq!(valid_32_byte_hex.len(), 64, "32-byte hex string should be exactly 64 characters");
667 let valid_64_byte_hex = "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000";
668 assert_eq!(valid_64_byte_hex.len(), 128, "64-byte hex string should be exactly 128 characters");
669
670 let input = format!("test_nonce|1000000|1030000|test_website|{}|0|{}|{}",
671 valid_32_byte_hex, valid_32_byte_hex, valid_64_byte_hex);
672 let result = IronShieldChallenge::from_concat_struct(&input);
673
674 assert!(result.is_ok(), "Should parse valid zero-value data");
675 let parsed = result.unwrap();
676 assert_eq!(parsed.random_nonce, "test_nonce");
677 assert_eq!(parsed.created_time, 1000000);
678 assert_eq!(parsed.expiration_time, 1030000);
679 assert_eq!(parsed.website_id, "test_website");
680 assert_eq!(parsed.challenge_param, [0u8; 32]);
681 assert_eq!(parsed.recommended_attempts, 0);
682 assert_eq!(parsed.public_key, [0u8; 32]);
683 assert_eq!(parsed.challenge_signature, [0u8; 64]);
684
685 let all_f_32_hex = "f".repeat(64);
687 assert_eq!(all_f_32_hex.len(), 64, "All F's 32-byte hex string should be exactly 64 characters");
688 let all_f_64_hex = "f".repeat(128);
689 assert_eq!(all_f_64_hex.len(), 128, "All F's 64-byte hex string should be exactly 128 characters");
690
691 let input = format!("max_nonce|{}|{}|max_website|{}|{}|{}|{}",
692 i64::MAX, i64::MAX, all_f_32_hex, u64::MAX, all_f_32_hex, all_f_64_hex);
693 let result = IronShieldChallenge::from_concat_struct(&input);
694
695 assert!(result.is_ok(), "Should parse valid max-value data");
696 let parsed = result.unwrap();
697 assert_eq!(parsed.random_nonce, "max_nonce");
698 assert_eq!(parsed.created_time, i64::MAX);
699 assert_eq!(parsed.expiration_time, i64::MAX);
700 assert_eq!(parsed.website_id, "max_website");
701 assert_eq!(parsed.challenge_param, [0xffu8; 32]);
702 assert_eq!(parsed.recommended_attempts, u64::MAX);
703 assert_eq!(parsed.public_key, [0xffu8; 32]);
704 assert_eq!(parsed.challenge_signature, [0xffu8; 64]);
705 }
706}