1use std::sync::Mutex;
13use std::time::Instant;
14
15use sha2::{Digest, Sha256};
16
17use crate::conditioning::{quick_min_entropy, quick_shannon};
18use crate::source::{EntropySource, SourceState};
19
20pub struct EntropyPool {
22 sources: Vec<Mutex<SourceState>>,
23 buffer: Mutex<Vec<u8>>,
24 state: Mutex<[u8; 32]>,
25 counter: Mutex<u64>,
26 total_output: Mutex<u64>,
27}
28
29impl EntropyPool {
30 pub fn new(seed: Option<&[u8]>) -> Self {
32 let initial_state = {
33 let mut h = Sha256::new();
34 if let Some(s) = seed {
35 h.update(s);
36 } else {
37 let mut os_random = [0u8; 32];
39 getrandom(&mut os_random);
40 h.update(os_random);
41 }
42 let digest: [u8; 32] = h.finalize().into();
43 digest
44 };
45
46 Self {
47 sources: Vec::new(),
48 buffer: Mutex::new(Vec::new()),
49 state: Mutex::new(initial_state),
50 counter: Mutex::new(0),
51 total_output: Mutex::new(0),
52 }
53 }
54
55 pub fn auto() -> Self {
57 let mut pool = Self::new(None);
58 for source in crate::platform::detect_available_sources() {
59 pool.add_source(source, 1.0);
60 }
61 pool
62 }
63
64 pub fn add_source(&mut self, source: Box<dyn EntropySource>, weight: f64) {
66 self.sources
67 .push(Mutex::new(SourceState::new(source, weight)));
68 }
69
70 pub fn source_count(&self) -> usize {
72 self.sources.len()
73 }
74
75 pub fn collect_all(&self) -> usize {
79 self.collect_all_parallel(10.0)
80 }
81
82 pub fn collect_all_parallel(&self, timeout_secs: f64) -> usize {
84 use std::sync::Arc;
85 let results: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
86 let deadline = Instant::now() + std::time::Duration::from_secs_f64(timeout_secs);
87
88 std::thread::scope(|s| {
89 let handles: Vec<_> = self
90 .sources
91 .iter()
92 .map(|ss_mutex| {
93 let results = Arc::clone(&results);
94 s.spawn(move || {
95 let data = Self::collect_one(ss_mutex);
96 if !data.is_empty() {
97 results.lock().unwrap().extend_from_slice(&data);
98 }
99 })
100 })
101 .collect();
102
103 for handle in handles {
104 let remaining = deadline.saturating_duration_since(Instant::now());
105 if remaining.is_zero() {
106 break;
107 }
108 let _ = handle.join();
109 }
110 });
111
112 let results = Arc::try_unwrap(results).unwrap().into_inner().unwrap();
113 let n = results.len();
114 self.buffer.lock().unwrap().extend_from_slice(&results);
115 n
116 }
117
118 pub fn collect_enabled(&self, enabled_names: &[String]) -> usize {
121 self.collect_enabled_n(enabled_names, 1000)
122 }
123
124 pub fn collect_enabled_n(&self, enabled_names: &[String], n_samples: usize) -> usize {
127 use std::sync::Arc;
128 let results: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
129
130 std::thread::scope(|s| {
131 let handles: Vec<_> = self
132 .sources
133 .iter()
134 .filter(|ss_mutex| {
135 let ss = ss_mutex.lock().unwrap();
136 enabled_names.iter().any(|n| n == ss.source.info().name)
137 })
138 .map(|ss_mutex| {
139 let results = Arc::clone(&results);
140 s.spawn(move || {
141 let data = Self::collect_one_n(ss_mutex, n_samples);
142 if !data.is_empty() {
143 results.lock().unwrap().extend_from_slice(&data);
144 }
145 })
146 })
147 .collect();
148
149 for handle in handles {
150 let _ = handle.join();
151 }
152 });
153
154 let results = Arc::try_unwrap(results).unwrap().into_inner().unwrap();
155 let n = results.len();
156 self.buffer.lock().unwrap().extend_from_slice(&results);
157 n
158 }
159
160 fn collect_one(ss_mutex: &Mutex<SourceState>) -> Vec<u8> {
161 Self::collect_one_n(ss_mutex, 1000)
162 }
163
164 fn collect_one_n(ss_mutex: &Mutex<SourceState>, n_samples: usize) -> Vec<u8> {
165 let mut ss = ss_mutex.lock().unwrap();
166 let t0 = Instant::now();
167 match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| ss.source.collect(n_samples))) {
168 Ok(data) if !data.is_empty() => {
169 ss.last_collect_time = t0.elapsed();
170 ss.total_bytes += data.len() as u64;
171 ss.last_entropy = quick_shannon(&data);
172 ss.last_min_entropy = quick_min_entropy(&data);
173 ss.healthy = ss.last_entropy > 1.0;
174 data
175 }
176 Ok(_) => {
177 ss.last_collect_time = t0.elapsed();
178 ss.failures += 1;
179 ss.healthy = false;
180 Vec::new()
181 }
182 Err(_) => {
183 ss.last_collect_time = t0.elapsed();
184 ss.failures += 1;
185 ss.healthy = false;
186 Vec::new()
187 }
188 }
189 }
190
191 pub fn get_raw_bytes(&self, n_bytes: usize) -> Vec<u8> {
196 {
198 let buf = self.buffer.lock().unwrap();
199 if buf.len() < n_bytes {
200 drop(buf);
201 self.collect_all();
202 }
203 }
204
205 let mut buf = self.buffer.lock().unwrap();
206 while buf.len() < n_bytes {
208 drop(buf);
209 self.collect_all();
210 buf = self.buffer.lock().unwrap();
211 }
212
213 let output: Vec<u8> = buf.drain(..n_bytes).collect();
214 drop(buf);
215 *self.total_output.lock().unwrap() += n_bytes as u64;
216 output
217 }
218
219 pub fn get_random_bytes(&self, n_bytes: usize) -> Vec<u8> {
221 {
223 let buf = self.buffer.lock().unwrap();
224 if buf.len() < n_bytes * 2 {
225 drop(buf);
226 self.collect_all();
227 }
228 }
229
230 let mut output = Vec::with_capacity(n_bytes);
231 while output.len() < n_bytes {
232 let mut counter = self.counter.lock().unwrap();
233 *counter += 1;
234 let cnt = *counter;
235 drop(counter);
236
237 let sample = {
239 let mut buf = self.buffer.lock().unwrap();
240 let take = buf.len().min(256);
241 let sample: Vec<u8> = buf.drain(..take).collect();
242 sample
243 };
244
245 let mut h = Sha256::new();
247 let state = self.state.lock().unwrap();
248 h.update(*state);
249 drop(state);
250 h.update(&sample);
251 h.update(cnt.to_le_bytes());
252
253 let ts = std::time::SystemTime::now()
254 .duration_since(std::time::UNIX_EPOCH)
255 .unwrap_or_default();
256 h.update(ts.as_nanos().to_le_bytes());
257
258 let mut os_random = [0u8; 8];
260 getrandom(&mut os_random);
261 h.update(os_random);
262
263 let digest: [u8; 32] = h.finalize().into();
264 *self.state.lock().unwrap() = digest;
265 output.extend_from_slice(&digest);
266 }
267
268 *self.total_output.lock().unwrap() += n_bytes as u64;
269 output.truncate(n_bytes);
270 output
271 }
272
273 pub fn get_bytes(
279 &self,
280 n_bytes: usize,
281 mode: crate::conditioning::ConditioningMode,
282 ) -> Vec<u8> {
283 use crate::conditioning::ConditioningMode;
284 match mode {
285 ConditioningMode::Raw => self.get_raw_bytes(n_bytes),
286 ConditioningMode::VonNeumann => {
287 let raw = self.get_raw_bytes(n_bytes * 6);
289 crate::conditioning::condition(&raw, n_bytes, ConditioningMode::VonNeumann)
290 }
291 ConditioningMode::Sha256 => self.get_random_bytes(n_bytes),
292 }
293 }
294
295 pub fn health_report(&self) -> HealthReport {
297 let mut sources = Vec::new();
298 let mut healthy_count = 0;
299 let mut total_raw = 0u64;
300
301 for ss_mutex in &self.sources {
302 let ss = ss_mutex.lock().unwrap();
303 if ss.healthy {
304 healthy_count += 1;
305 }
306 total_raw += ss.total_bytes;
307 sources.push(SourceHealth {
308 name: ss.source.name().to_string(),
309 healthy: ss.healthy,
310 bytes: ss.total_bytes,
311 entropy: ss.last_entropy,
312 min_entropy: ss.last_min_entropy,
313 time: ss.last_collect_time.as_secs_f64(),
314 failures: ss.failures,
315 });
316 }
317
318 HealthReport {
319 healthy: healthy_count,
320 total: self.sources.len(),
321 raw_bytes: total_raw,
322 output_bytes: *self.total_output.lock().unwrap(),
323 buffer_size: self.buffer.lock().unwrap().len(),
324 sources,
325 }
326 }
327
328 pub fn print_health(&self) {
330 let r = self.health_report();
331 println!("\n{}", "=".repeat(60));
332 println!("ENTROPY POOL HEALTH REPORT");
333 println!("{}", "=".repeat(60));
334 println!("Sources: {}/{} healthy", r.healthy, r.total);
335 println!("Raw collected: {} bytes", r.raw_bytes);
336 println!(
337 "Output: {} bytes | Buffer: {} bytes",
338 r.output_bytes, r.buffer_size
339 );
340 println!(
341 "\n{:<25} {:>4} {:>10} {:>6} {:>6} {:>7} {:>5}",
342 "Source", "OK", "Bytes", "H", "H∞", "Time", "Fail"
343 );
344 println!("{}", "-".repeat(68));
345 for s in &r.sources {
346 let ok = if s.healthy { "✓" } else { "✗" };
347 println!(
348 "{:<25} {:>4} {:>10} {:>5.2} {:>5.2} {:>6.3}s {:>5}",
349 s.name, ok, s.bytes, s.entropy, s.min_entropy, s.time, s.failures
350 );
351 }
352 }
353
354 pub fn source_infos(&self) -> Vec<SourceInfoSnapshot> {
356 self.sources
357 .iter()
358 .map(|ss_mutex| {
359 let ss = ss_mutex.lock().unwrap();
360 let info = ss.source.info();
361 SourceInfoSnapshot {
362 name: info.name.to_string(),
363 description: info.description.to_string(),
364 physics: info.physics.to_string(),
365 category: info.category.to_string(),
366 entropy_rate_estimate: info.entropy_rate_estimate,
367 composite: info.composite,
368 }
369 })
370 .collect()
371 }
372}
373
374fn getrandom(buf: &mut [u8]) {
380 getrandom::fill(buf).expect("OS CSPRNG failed");
381}
382
383#[derive(Debug, Clone)]
385pub struct HealthReport {
386 pub healthy: usize,
388 pub total: usize,
390 pub raw_bytes: u64,
392 pub output_bytes: u64,
394 pub buffer_size: usize,
396 pub sources: Vec<SourceHealth>,
398}
399
400#[derive(Debug, Clone)]
402pub struct SourceHealth {
403 pub name: String,
405 pub healthy: bool,
407 pub bytes: u64,
409 pub entropy: f64,
411 pub min_entropy: f64,
413 pub time: f64,
415 pub failures: u64,
417}
418
419#[derive(Debug, Clone)]
421pub struct SourceInfoSnapshot {
422 pub name: String,
424 pub description: String,
426 pub physics: String,
428 pub category: String,
430 pub entropy_rate_estimate: f64,
432 pub composite: bool,
434}
435
436#[cfg(test)]
437mod tests {
438 use super::*;
439 use crate::source::{SourceCategory, SourceInfo};
440
441 struct MockSource {
447 info: SourceInfo,
448 data: Vec<u8>,
449 }
450
451 impl MockSource {
452 fn new(name: &'static str, data: Vec<u8>) -> Self {
453 Self {
454 info: SourceInfo {
455 name,
456 description: "mock source",
457 physics: "deterministic test data",
458 category: SourceCategory::System,
459 platform_requirements: &[],
460 entropy_rate_estimate: 1.0,
461 composite: false,
462 },
463 data,
464 }
465 }
466 }
467
468 impl EntropySource for MockSource {
469 fn info(&self) -> &SourceInfo {
470 &self.info
471 }
472 fn is_available(&self) -> bool {
473 true
474 }
475 fn collect(&self, n_samples: usize) -> Vec<u8> {
476 self.data.iter().copied().cycle().take(n_samples).collect()
477 }
478 }
479
480 struct FailingSource {
482 info: SourceInfo,
483 }
484
485 impl FailingSource {
486 fn new(name: &'static str) -> Self {
487 Self {
488 info: SourceInfo {
489 name,
490 description: "failing mock",
491 physics: "always fails",
492 category: SourceCategory::System,
493 platform_requirements: &[],
494 entropy_rate_estimate: 0.0,
495 composite: false,
496 },
497 }
498 }
499 }
500
501 impl EntropySource for FailingSource {
502 fn info(&self) -> &SourceInfo {
503 &self.info
504 }
505 fn is_available(&self) -> bool {
506 true
507 }
508 fn collect(&self, _n_samples: usize) -> Vec<u8> {
509 Vec::new()
510 }
511 }
512
513 #[test]
518 fn test_pool_new_empty() {
519 let pool = EntropyPool::new(None);
520 assert_eq!(pool.source_count(), 0);
521 }
522
523 #[test]
524 fn test_pool_new_with_seed() {
525 let pool = EntropyPool::new(Some(b"test seed"));
526 assert_eq!(pool.source_count(), 0);
527 }
528
529 #[test]
530 fn test_pool_add_source() {
531 let mut pool = EntropyPool::new(Some(b"test"));
532 pool.add_source(Box::new(MockSource::new("mock1", vec![42])), 1.0);
533 assert_eq!(pool.source_count(), 1);
534 }
535
536 #[test]
537 fn test_pool_add_multiple_sources() {
538 let mut pool = EntropyPool::new(Some(b"test"));
539 pool.add_source(Box::new(MockSource::new("mock1", vec![1])), 1.0);
540 pool.add_source(Box::new(MockSource::new("mock2", vec![2])), 1.0);
541 pool.add_source(Box::new(MockSource::new("mock3", vec![3])), 0.5);
542 assert_eq!(pool.source_count(), 3);
543 }
544
545 #[test]
550 fn test_collect_all_returns_bytes() {
551 let mut pool = EntropyPool::new(Some(b"test"));
552 pool.add_source(
553 Box::new(MockSource::new("mock1", vec![0xAA, 0xBB, 0xCC])),
554 1.0,
555 );
556 let n = pool.collect_all();
557 assert!(n > 0, "Should have collected some bytes");
558 }
559
560 #[test]
561 fn test_collect_all_parallel_with_timeout() {
562 let mut pool = EntropyPool::new(Some(b"test"));
563 pool.add_source(Box::new(MockSource::new("mock1", vec![1, 2])), 1.0);
564 pool.add_source(Box::new(MockSource::new("mock2", vec![3, 4])), 1.0);
565 let n = pool.collect_all_parallel(5.0);
566 assert!(n > 0);
567 }
568
569 #[test]
570 fn test_collect_enabled_filters_sources() {
571 let mut pool = EntropyPool::new(Some(b"test"));
572 pool.add_source(Box::new(MockSource::new("alpha", vec![1])), 1.0);
573 pool.add_source(Box::new(MockSource::new("beta", vec![2])), 1.0);
574
575 let enabled = vec!["alpha".to_string()];
576 let n = pool.collect_enabled(&enabled);
577 assert!(n > 0, "Should collect from enabled source");
578 }
579
580 #[test]
581 fn test_collect_enabled_no_match() {
582 let mut pool = EntropyPool::new(Some(b"test"));
583 pool.add_source(Box::new(MockSource::new("alpha", vec![1])), 1.0);
584
585 let enabled = vec!["nonexistent".to_string()];
586 let n = pool.collect_enabled(&enabled);
587 assert_eq!(n, 0, "No sources should match");
588 }
589
590 #[test]
595 fn test_get_raw_bytes_length() {
596 let mut pool = EntropyPool::new(Some(b"test"));
597 pool.add_source(Box::new(MockSource::new("mock", (0..=255).collect())), 1.0);
598 let bytes = pool.get_raw_bytes(64);
599 assert_eq!(bytes.len(), 64);
600 }
601
602 #[test]
603 fn test_get_random_bytes_length() {
604 let mut pool = EntropyPool::new(Some(b"test"));
605 pool.add_source(Box::new(MockSource::new("mock", (0..=255).collect())), 1.0);
606 let bytes = pool.get_random_bytes(64);
607 assert_eq!(bytes.len(), 64);
608 }
609
610 #[test]
611 fn test_get_random_bytes_various_sizes() {
612 let mut pool = EntropyPool::new(Some(b"test"));
613 pool.add_source(Box::new(MockSource::new("mock", (0..=255).collect())), 1.0);
614 for size in [1, 16, 32, 64, 100, 256] {
615 let bytes = pool.get_random_bytes(size);
616 assert_eq!(bytes.len(), size, "Expected {size} bytes");
617 }
618 }
619
620 #[test]
621 fn test_get_bytes_raw_mode() {
622 let mut pool = EntropyPool::new(Some(b"test"));
623 pool.add_source(Box::new(MockSource::new("mock", (0..=255).collect())), 1.0);
624 let bytes = pool.get_bytes(32, crate::conditioning::ConditioningMode::Raw);
625 assert_eq!(bytes.len(), 32);
626 }
627
628 #[test]
629 fn test_get_bytes_sha256_mode() {
630 let mut pool = EntropyPool::new(Some(b"test"));
631 pool.add_source(Box::new(MockSource::new("mock", (0..=255).collect())), 1.0);
632 let bytes = pool.get_bytes(32, crate::conditioning::ConditioningMode::Sha256);
633 assert_eq!(bytes.len(), 32);
634 }
635
636 #[test]
637 fn test_get_bytes_von_neumann_mode() {
638 let mut pool = EntropyPool::new(Some(b"test"));
639 pool.add_source(Box::new(MockSource::new("mock", (0..=255).collect())), 1.0);
640 let bytes = pool.get_bytes(16, crate::conditioning::ConditioningMode::VonNeumann);
641 assert!(bytes.len() <= 16);
643 }
644
645 #[test]
650 fn test_health_report_empty_pool() {
651 let pool = EntropyPool::new(Some(b"test"));
652 let report = pool.health_report();
653 assert_eq!(report.total, 0);
654 assert_eq!(report.healthy, 0);
655 assert_eq!(report.raw_bytes, 0);
656 assert_eq!(report.output_bytes, 0);
657 assert_eq!(report.buffer_size, 0);
658 assert!(report.sources.is_empty());
659 }
660
661 #[test]
662 fn test_health_report_after_collection() {
663 let mut pool = EntropyPool::new(Some(b"test"));
664 pool.add_source(
665 Box::new(MockSource::new("good_source", (0..=255).collect())),
666 1.0,
667 );
668 pool.collect_all();
669 let report = pool.health_report();
670 assert_eq!(report.total, 1);
671 assert!(report.raw_bytes > 0);
672 assert_eq!(report.sources.len(), 1);
673 assert_eq!(report.sources[0].name, "good_source");
674 assert!(report.sources[0].bytes > 0);
675 }
676
677 #[test]
678 fn test_health_report_failing_source() {
679 let mut pool = EntropyPool::new(Some(b"test"));
680 pool.add_source(Box::new(FailingSource::new("bad_source")), 1.0);
681 pool.collect_all();
682 let report = pool.health_report();
683 assert_eq!(report.total, 1);
684 assert_eq!(report.healthy, 0);
685 assert!(!report.sources[0].healthy);
686 assert_eq!(report.sources[0].failures, 1);
687 }
688
689 #[test]
690 fn test_health_report_mixed_sources() {
691 let mut pool = EntropyPool::new(Some(b"test"));
692 pool.add_source(Box::new(MockSource::new("good", (0..=255).collect())), 1.0);
693 pool.add_source(Box::new(FailingSource::new("bad")), 1.0);
694 pool.collect_all();
695 let report = pool.health_report();
696 assert_eq!(report.total, 2);
697 assert!(report.healthy >= 1);
699 assert_eq!(report.sources.len(), 2);
700 }
701
702 #[test]
703 fn test_health_report_tracks_output_bytes() {
704 let mut pool = EntropyPool::new(Some(b"test"));
705 pool.add_source(Box::new(MockSource::new("mock", (0..=255).collect())), 1.0);
706 let _ = pool.get_random_bytes(64);
707 let report = pool.health_report();
708 assert!(report.output_bytes >= 64);
709 }
710
711 #[test]
716 fn test_source_infos_empty() {
717 let pool = EntropyPool::new(Some(b"test"));
718 let infos = pool.source_infos();
719 assert!(infos.is_empty());
720 }
721
722 #[test]
723 fn test_source_infos_populated() {
724 let mut pool = EntropyPool::new(Some(b"test"));
725 pool.add_source(Box::new(MockSource::new("test_src", vec![1])), 1.0);
726 let infos = pool.source_infos();
727 assert_eq!(infos.len(), 1);
728 assert_eq!(infos[0].name, "test_src");
729 assert_eq!(infos[0].description, "mock source");
730 assert_eq!(infos[0].category, "system");
731 assert!((infos[0].entropy_rate_estimate - 1.0).abs() < f64::EPSILON);
732 }
733
734 #[test]
739 fn test_different_seeds_differ() {
740 let mut pool1 = EntropyPool::new(Some(b"seed_a"));
741 pool1.add_source(Box::new(MockSource::new("m", vec![42; 100])), 1.0);
742 let mut pool2 = EntropyPool::new(Some(b"seed_b"));
743 pool2.add_source(Box::new(MockSource::new("m", vec![42; 100])), 1.0);
744
745 let bytes1 = pool1.get_random_bytes(32);
746 let bytes2 = pool2.get_random_bytes(32);
747 assert_ne!(
748 bytes1, bytes2,
749 "Different seeds should produce different output"
750 );
751 }
752
753 #[test]
758 fn test_collect_from_empty_pool() {
759 let pool = EntropyPool::new(Some(b"test"));
760 let n = pool.collect_all();
761 assert_eq!(n, 0, "Empty pool should collect 0 bytes");
762 }
763
764 #[test]
765 fn test_collect_enabled_empty_list() {
766 let mut pool = EntropyPool::new(Some(b"test"));
767 pool.add_source(Box::new(MockSource::new("mock", vec![1])), 1.0);
768 let n = pool.collect_enabled(&[]);
769 assert_eq!(n, 0);
770 }
771}