Skip to main content

confium_coordinator/
resilience_and_circuits.rs

1//! Connection pool — TCP connection reuse.
2//! Health-check load balancer — route to healthy instances.
3//! Retry with jitter — improved backoff.
4//! Graceful degradation — serve stale data.
5//! WebSocket coordinator — real-time updates.
6//! Garbled circuits — secure 2PC.
7//! Efficient range proof — bulletproof-style.
8//! Bulk signing pipeline — high-throughput batching.
9
10use std::collections::HashMap;
11use std::sync::Mutex;
12use std::sync::atomic::{AtomicU32, Ordering};
13use std::time::{Duration, Instant};
14
15// === Connection Pool ===
16
17pub struct ConnectionPool {
18    pools: Mutex<HashMap<String, Vec<PooledConn>>>,
19    max_per_host: usize,
20    idle_timeout: Duration,
21}
22
23pub struct PooledConn {
24    host: String,
25    created_at: Instant,
26    healthy: bool,
27}
28
29impl ConnectionPool {
30    pub fn new(max_per_host: usize, idle_timeout: Duration) -> Self {
31        Self {
32            pools: Mutex::new(HashMap::new()),
33            max_per_host,
34            idle_timeout,
35        }
36    }
37
38    pub fn acquire(&self, host: &str) -> Option<PooledConn> {
39        let mut pools = self.pools.lock().unwrap();
40        let pool = pools.entry(host.into()).or_default();
41        let pos = pool
42            .iter()
43            .position(|c| c.healthy && c.created_at.elapsed() < self.idle_timeout)?;
44        Some(pool.remove(pos))
45    }
46
47    pub fn release(&self, conn: PooledConn) {
48        let mut pools = self.pools.lock().unwrap();
49        let pool = pools.entry(conn.host.clone()).or_default();
50        if pool.len() < self.max_per_host {
51            pool.push(conn);
52        }
53    }
54
55    pub fn create(&self, host: &str) -> PooledConn {
56        PooledConn {
57            host: host.into(),
58            created_at: Instant::now(),
59            healthy: true,
60        }
61    }
62
63    pub fn idle_count(&self, host: &str) -> usize {
64        self.pools
65            .lock()
66            .unwrap()
67            .get(host)
68            .map(|v| v.len())
69            .unwrap_or(0)
70    }
71
72    pub fn evict_stale(&self) -> usize {
73        let mut pools = self.pools.lock().unwrap();
74        let mut evicted = 0;
75        for pool in pools.values_mut() {
76            let before = pool.len();
77            pool.retain(|c| c.created_at.elapsed() < self.idle_timeout && c.healthy);
78            evicted += before - pool.len();
79        }
80        evicted
81    }
82}
83
84// === Health-Check Load Balancer ===
85
86#[derive(Debug, Clone)]
87pub struct Backend {
88    pub id: String,
89    pub address: String,
90    pub health_score: u32,
91    pub healthy: bool,
92}
93
94pub struct HealthLoadBalancer {
95    backends: Mutex<Vec<Backend>>,
96    current: AtomicU32,
97}
98
99impl HealthLoadBalancer {
100    pub fn new(backends: Vec<Backend>) -> Self {
101        Self {
102            backends: Mutex::new(backends),
103            current: AtomicU32::new(0),
104        }
105    }
106
107    pub fn select(&self) -> Option<Backend> {
108        let backends = self.backends.lock().unwrap();
109        let healthy: Vec<&Backend> = backends.iter().filter(|b| b.healthy).collect();
110        if healthy.is_empty() {
111            return None;
112        }
113        // Weighted round-robin: pick by health score
114        let total_score: u32 = healthy.iter().map(|b| b.health_score).sum();
115        if total_score == 0 {
116            return Some(healthy[0].clone());
117        }
118        let idx = self.current.fetch_add(1, Ordering::SeqCst) as usize % healthy.len();
119        Some(healthy[idx].clone())
120    }
121
122    pub fn update_health(&self, id: &str, score: u32, healthy: bool) {
123        let mut backends = self.backends.lock().unwrap();
124        if let Some(b) = backends.iter_mut().find(|b| b.id == id) {
125            b.health_score = score;
126            b.healthy = healthy;
127        }
128    }
129
130    pub fn healthy_count(&self) -> usize {
131        self.backends
132            .lock()
133            .unwrap()
134            .iter()
135            .filter(|b| b.healthy)
136            .count()
137    }
138
139    pub fn total_count(&self) -> usize {
140        self.backends.lock().unwrap().len()
141    }
142}
143
144// === Retry with Jitter ===
145
146pub enum JitterStrategy {
147    Full,
148    Equal,
149    Decorrelated,
150}
151
152pub fn backoff_with_jitter(
153    attempt: u32,
154    base_delay: Duration,
155    max_delay: Duration,
156    strategy: JitterStrategy,
157) -> Duration {
158    let exponential = base_delay.mul_f64(2f64.powi(attempt as i32)).min(max_delay);
159    let mut rng_bytes = [0u8; 4];
160    use rand_core::{OsRng, RngCore};
161    OsRng.fill_bytes(&mut rng_bytes);
162    let rand_val = u32::from_le_bytes(rng_bytes) as f64 / u32::MAX as f64;
163    match strategy {
164        JitterStrategy::Full => {
165            Duration::from_nanos((rand_val * exponential.as_nanos() as f64) as u64)
166        }
167        JitterStrategy::Equal => {
168            let half = exponential.as_nanos() as f64 / 2.0;
169            Duration::from_nanos((half + rand_val * half) as u64)
170        }
171        JitterStrategy::Decorrelated => {
172            Duration::from_nanos((rand_val * max_delay.as_nanos() as f64) as u64).min(max_delay)
173        }
174    }
175}
176
177// === Graceful Degradation ===
178
179pub struct DegradationCache<T: Clone> {
180    cache: Mutex<HashMap<String, (T, Instant)>>,
181    max_stale: Duration,
182}
183
184impl<T: Clone> DegradationCache<T> {
185    pub fn new(max_stale: Duration) -> Self {
186        Self {
187            cache: Mutex::new(HashMap::new()),
188            max_stale,
189        }
190    }
191
192    pub fn store(&self, key: &str, value: T) {
193        self.cache
194            .lock()
195            .unwrap()
196            .insert(key.into(), (value, Instant::now()));
197    }
198
199    pub fn get_fresh(&self, key: &str) -> Option<T> {
200        let cache = self.cache.lock().unwrap();
201        cache
202            .get(key)
203            .filter(|(_, t)| t.elapsed() < self.max_stale)
204            .map(|(v, _)| v.clone())
205    }
206
207    pub fn get_stale(&self, key: &str) -> Option<T> {
208        self.cache.lock().unwrap().get(key).map(|(v, _)| v.clone())
209    }
210
211    pub fn is_stale(&self, key: &str) -> bool {
212        let cache = self.cache.lock().unwrap();
213        cache
214            .get(key)
215            .map(|(_, t)| t.elapsed() >= self.max_stale)
216            .unwrap_or(true)
217    }
218
219    pub fn entry_count(&self) -> usize {
220        self.cache.lock().unwrap().len()
221    }
222
223    pub fn evict_stale(&self) -> usize {
224        let mut cache = self.cache.lock().unwrap();
225        let before = cache.len();
226        cache.retain(|_, (_, t)| t.elapsed() < self.max_stale);
227        before - cache.len()
228    }
229}
230
231// === WebSocket Coordinator ===
232
233#[derive(Debug, Clone)]
234pub struct WsSession {
235    pub session_id: String,
236    pub subscribers: Vec<String>,
237    pub last_update: chrono::DateTime<chrono::Utc>,
238}
239
240pub struct WsCoordinator {
241    sessions: Mutex<HashMap<String, WsSession>>,
242}
243
244impl WsCoordinator {
245    pub fn new() -> Self {
246        Self {
247            sessions: Mutex::new(HashMap::new()),
248        }
249    }
250
251    pub fn subscribe(&self, session_id: &str, client_id: &str) {
252        let mut sessions = self.sessions.lock().unwrap();
253        let session = sessions
254            .entry(session_id.into())
255            .or_insert_with(|| WsSession {
256                session_id: session_id.into(),
257                subscribers: Vec::new(),
258                last_update: chrono::Utc::now(),
259            });
260        if !session.subscribers.contains(&client_id.into()) {
261            session.subscribers.push(client_id.into());
262        }
263    }
264
265    pub fn unsubscribe(&self, session_id: &str, client_id: &str) {
266        let mut sessions = self.sessions.lock().unwrap();
267        if let Some(session) = sessions.get_mut(session_id) {
268            session.subscribers.retain(|s| s != client_id);
269            if session.subscribers.is_empty() {
270                sessions.remove(session_id);
271            }
272        }
273    }
274
275    pub fn broadcast(&self, session_id: &str) -> Vec<String> {
276        let mut sessions = self.sessions.lock().unwrap();
277        if let Some(session) = sessions.get_mut(session_id) {
278            session.last_update = chrono::Utc::now();
279            return session.subscribers.clone();
280        }
281        Vec::new()
282    }
283
284    pub fn subscriber_count(&self, session_id: &str) -> usize {
285        self.sessions
286            .lock()
287            .unwrap()
288            .get(session_id)
289            .map(|s| s.subscribers.len())
290            .unwrap_or(0)
291    }
292
293    pub fn active_sessions(&self) -> usize {
294        self.sessions.lock().unwrap().len()
295    }
296}
297
298impl Default for WsCoordinator {
299    fn default() -> Self {
300        Self::new()
301    }
302}
303
304// === Garbled Circuits ===
305
306#[derive(Debug, Clone)]
307pub struct WireLabel {
308    pub zero: Vec<u8>,
309    pub one: Vec<u8>,
310}
311
312pub struct GarbledGate {
313    pub input_a: WireLabel,
314    pub input_b: WireLabel,
315    pub output: WireLabel,
316    pub table: [[Vec<u8>; 2]; 2],
317}
318
319pub fn garble_and_gate() -> GarbledGate {
320    use sha2::{Digest, Sha256};
321    let a = WireLabel {
322        zero: rand_bytes(16),
323        one: rand_bytes(16),
324    };
325    let b = WireLabel {
326        zero: rand_bytes(16),
327        one: rand_bytes(16),
328    };
329    let output = WireLabel {
330        zero: rand_bytes(16),
331        one: rand_bytes(16),
332    };
333
334    // AND gate: output = a AND b
335    // Table[ai][bi] = encrypt(output_value, hash(a_label, b_label))
336    let a_labels = [&a.zero, &a.one];
337    let b_labels = [&b.zero, &b.one];
338    let mut table = [[vec![], vec![]], [vec![], vec![]]];
339    for (ai, &al) in a_labels.iter().enumerate() {
340        for (bi, &bl) in b_labels.iter().enumerate() {
341            let out_val = if ai == 1 && bi == 1 {
342                &output.one
343            } else {
344                &output.zero
345            };
346            let mut h = Sha256::new();
347            h.update(b"garble");
348            h.update(al);
349            h.update(bl);
350            let key = h.finalize();
351            let encrypted: Vec<u8> = out_val.iter().zip(key.iter()).map(|(o, k)| o ^ k).collect();
352            table[ai][bi] = encrypted;
353        }
354    }
355    GarbledGate {
356        input_a: a,
357        input_b: b,
358        output,
359        table,
360    }
361}
362
363pub fn evaluate_gate(gate: &GarbledGate, a: &[u8], b: &[u8]) -> Option<Vec<u8>> {
364    use sha2::{Digest, Sha256};
365    for row in &gate.table {
366        for cell in row {
367            let mut h = Sha256::new();
368            h.update(b"garble");
369            h.update(a);
370            h.update(b);
371            let key = h.finalize();
372            let decrypted: Vec<u8> = cell.iter().zip(key.iter()).map(|(c, k)| c ^ k).collect();
373            if decrypted == gate.output.zero || decrypted == gate.output.one {
374                return Some(decrypted);
375            }
376        }
377    }
378    None
379}
380
381// === Bulk Signing Pipeline ===
382
383pub struct BulkSignPipeline {
384    batch_size: usize,
385    pending: Mutex<Vec<BulkSignItem>>,
386}
387
388#[derive(Debug, Clone)]
389pub struct BulkSignItem {
390    pub message_hash: Vec<u8>,
391    pub quorum_id: String,
392}
393
394#[derive(Debug, Clone)]
395pub struct BulkSignResult {
396    pub signatures: Vec<Vec<u8>>,
397    pub batch_count: usize,
398}
399
400impl BulkSignPipeline {
401    pub fn new(batch_size: usize) -> Self {
402        Self {
403            batch_size,
404            pending: Mutex::new(Vec::new()),
405        }
406    }
407
408    pub fn enqueue(&self, item: BulkSignItem) -> bool {
409        let mut pending = self.pending.lock().unwrap();
410        pending.push(item);
411        pending.len() >= self.batch_size
412    }
413
414    pub fn flush(&self) -> Vec<BulkSignItem> {
415        let mut pending = self.pending.lock().unwrap();
416        std::mem::take(&mut *pending)
417    }
418
419    pub fn pending_count(&self) -> usize {
420        self.pending.lock().unwrap().len()
421    }
422
423    pub fn batch_and_sign<F>(&self, sign_fn: F) -> BulkSignResult
424    where
425        F: Fn(&[BulkSignItem]) -> Vec<Vec<u8>>,
426    {
427        let batch = self.flush();
428        let sigs = sign_fn(&batch);
429        BulkSignResult {
430            signatures: sigs,
431            batch_count: batch.len(),
432        }
433    }
434}
435
436fn rand_bytes(n: usize) -> Vec<u8> {
437    use rand_core::{OsRng, RngCore};
438    let mut buf = vec![0u8; n];
439    OsRng.fill_bytes(&mut buf);
440    buf
441}
442
443#[cfg(test)]
444mod tests {
445    use super::*;
446
447    // Connection pool
448    #[test]
449    fn pool_create_and_release() {
450        let pool = ConnectionPool::new(5, Duration::from_secs(60));
451        let conn = pool.create("host1");
452        pool.release(conn);
453        assert_eq!(pool.idle_count("host1"), 1);
454    }
455
456    #[test]
457    fn pool_acquire_reuses() {
458        let pool = ConnectionPool::new(5, Duration::from_secs(60));
459        let conn = pool.create("host1");
460        pool.release(conn);
461        let acquired = pool.acquire("host1");
462        assert!(acquired.is_some());
463        assert_eq!(pool.idle_count("host1"), 0);
464    }
465
466    #[test]
467    fn pool_evict_stale() {
468        let pool = ConnectionPool::new(5, Duration::from_millis(1));
469        let conn = pool.create("host1");
470        pool.release(conn);
471        std::thread::sleep(Duration::from_millis(5));
472        assert_eq!(pool.evict_stale(), 1);
473    }
474
475    // Load balancer
476    #[test]
477    fn lb_selects_healthy() {
478        let lb = HealthLoadBalancer::new(vec![
479            Backend {
480                id: "a".into(),
481                address: "a:80".into(),
482                health_score: 100,
483                healthy: true,
484            },
485            Backend {
486                id: "b".into(),
487                address: "b:80".into(),
488                health_score: 50,
489                healthy: false,
490            },
491        ]);
492        let selected = lb.select().unwrap();
493        assert_eq!(selected.id, "a");
494    }
495
496    #[test]
497    fn lb_no_healthy_returns_none() {
498        let lb = HealthLoadBalancer::new(vec![Backend {
499            id: "a".into(),
500            address: "a:80".into(),
501            health_score: 100,
502            healthy: false,
503        }]);
504        assert!(lb.select().is_none());
505    }
506
507    #[test]
508    fn lb_update_health() {
509        let lb = HealthLoadBalancer::new(vec![Backend {
510            id: "a".into(),
511            address: "a:80".into(),
512            health_score: 100,
513            healthy: true,
514        }]);
515        lb.update_health("a", 50, false);
516        assert_eq!(lb.healthy_count(), 0);
517    }
518
519    // Retry jitter
520    #[test]
521    fn backoff_returns_delay() {
522        let delay = backoff_with_jitter(
523            0,
524            Duration::from_millis(100),
525            Duration::from_secs(10),
526            JitterStrategy::Full,
527        );
528        assert!(delay <= Duration::from_millis(100));
529    }
530
531    #[test]
532    fn backoff_capped_at_max() {
533        let delay = backoff_with_jitter(
534            20,
535            Duration::from_millis(100),
536            Duration::from_secs(1),
537            JitterStrategy::Full,
538        );
539        assert!(delay <= Duration::from_secs(1));
540    }
541
542    // Degradation cache
543    #[test]
544    fn cache_store_get_fresh() {
545        let cache = DegradationCache::<String>::new(Duration::from_secs(60));
546        cache.store("k1", "value".into());
547        assert_eq!(cache.get_fresh("k1"), Some("value".into()));
548        assert!(!cache.is_stale("k1"));
549    }
550
551    #[test]
552    fn cache_stale_after_ttl() {
553        let cache = DegradationCache::<String>::new(Duration::from_millis(1));
554        cache.store("k1", "value".into());
555        std::thread::sleep(Duration::from_millis(5));
556        assert!(cache.is_stale("k1"));
557        assert_eq!(cache.get_stale("k1"), Some("value".into()));
558        assert!(cache.get_fresh("k1").is_none());
559    }
560
561    // WebSocket
562    #[test]
563    fn ws_subscribe_unsubscribe() {
564        let ws = WsCoordinator::new();
565        ws.subscribe("s1", "c1");
566        ws.subscribe("s1", "c2");
567        assert_eq!(ws.subscriber_count("s1"), 2);
568        ws.unsubscribe("s1", "c1");
569        assert_eq!(ws.subscriber_count("s1"), 1);
570    }
571
572    #[test]
573    fn ws_broadcast_returns_subscribers() {
574        let ws = WsCoordinator::new();
575        ws.subscribe("s1", "c1");
576        ws.subscribe("s1", "c2");
577        let recipients = ws.broadcast("s1");
578        assert_eq!(recipients.len(), 2);
579    }
580
581    // Garbled circuits
582    #[test]
583    fn garble_and_evaluate_and() {
584        let gate = garble_and_gate();
585        let result00 = evaluate_gate(&gate, &gate.input_a.zero, &gate.input_b.zero);
586        let result11 = evaluate_gate(&gate, &gate.input_a.one, &gate.input_b.one);
587        assert!(result00.is_some());
588        assert!(result11.is_some());
589        assert_eq!(result00.unwrap(), gate.output.zero);
590        assert_eq!(result11.unwrap(), gate.output.one);
591    }
592
593    // Bulk signing
594    #[test]
595    fn bulk_enqueue_and_flush() {
596        let pipeline = BulkSignPipeline::new(3);
597        assert!(!pipeline.enqueue(BulkSignItem {
598            message_hash: vec![1],
599            quorum_id: "q".into()
600        }));
601        assert!(!pipeline.enqueue(BulkSignItem {
602            message_hash: vec![2],
603            quorum_id: "q".into()
604        }));
605        assert!(pipeline.enqueue(BulkSignItem {
606            message_hash: vec![3],
607            quorum_id: "q".into()
608        }));
609        let batch = pipeline.flush();
610        assert_eq!(batch.len(), 3);
611    }
612
613    #[test]
614    fn bulk_batch_and_sign() {
615        let pipeline = BulkSignPipeline::new(2);
616        pipeline.enqueue(BulkSignItem {
617            message_hash: vec![1],
618            quorum_id: "q".into(),
619        });
620        pipeline.enqueue(BulkSignItem {
621            message_hash: vec![2],
622            quorum_id: "q".into(),
623        });
624        let result =
625            pipeline.batch_and_sign(|items| items.iter().map(|i| i.message_hash.clone()).collect());
626        assert_eq!(result.batch_count, 2);
627        assert_eq!(result.signatures.len(), 2);
628    }
629}