1use std::collections::{HashMap, HashSet};
17
18pub type ProtocolVersion = String;
20
21#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct ProtocolMetadata {
24 pub name: String,
26 pub version: ProtocolVersion,
28 pub priority: i32,
30 pub description: String,
32}
33
34impl ProtocolMetadata {
35 pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
37 Self {
38 name: name.into(),
39 version: version.into(),
40 priority: 0,
41 description: String::new(),
42 }
43 }
44
45 pub fn with_priority(mut self, priority: i32) -> Self {
47 self.priority = priority;
48 self
49 }
50
51 pub fn with_description(mut self, desc: impl Into<String>) -> Self {
53 self.description = desc.into();
54 self
55 }
56
57 pub fn header_value(&self) -> String {
60 format!("{}.{}", self.name, self.version)
61 }
62}
63
64#[derive(Debug, Clone, PartialEq, Eq)]
66pub enum NegotiationOutcome {
67 Accepted {
69 header_value: String,
71 metadata: ProtocolMetadata,
73 },
74 NotRequested,
76 NoMatch {
78 requested: Vec<String>,
80 },
81}
82
83impl NegotiationOutcome {
84 pub fn is_accepted(&self) -> bool {
86 matches!(self, NegotiationOutcome::Accepted { .. })
87 }
88
89 pub fn header_value(&self) -> Option<&str> {
91 match self {
92 NegotiationOutcome::Accepted { header_value, .. } => Some(header_value),
93 _ => None,
94 }
95 }
96}
97
98#[derive(Debug, Clone, Default)]
103pub struct SubProtocolRegistry {
104 protocols: Vec<String>,
106}
107
108impl SubProtocolRegistry {
109 pub fn new() -> Self {
111 Self::default()
112 }
113
114 pub fn register(&mut self, name: &str) {
116 if !self.protocols.contains(&name.to_string()) {
117 self.protocols.push(name.to_string());
118 }
119 }
120
121 pub fn register_many(&mut self, names: &[&str]) {
123 for name in names {
124 self.register(name);
125 }
126 }
127
128 pub fn is_registered(&self, name: &str) -> bool {
130 self.protocols.contains(&name.to_string())
131 }
132
133 pub fn protocols(&self) -> &[String] {
135 &self.protocols
136 }
137
138 pub fn negotiate(&self, client_protocols: &[String]) -> Option<String> {
142 let registered: HashSet<&str> = self.protocols.iter().map(|s| s.as_str()).collect();
143 client_protocols
144 .iter()
145 .find(|p| registered.contains(p.as_str()))
146 .cloned()
147 }
148
149 pub fn clear(&mut self) {
151 self.protocols.clear();
152 }
153
154 pub fn len(&self) -> usize {
156 self.protocols.len()
157 }
158
159 pub fn is_empty(&self) -> bool {
161 self.protocols.is_empty()
162 }
163}
164
165#[derive(Debug, Clone, Default)]
167pub struct NegotiationStats {
168 pub total_negotiations: u64,
170 pub accepted: u64,
172 pub not_requested: u64,
174 pub no_match: u64,
176}
177
178impl NegotiationStats {
179 pub fn success_rate(&self) -> f64 {
181 if self.total_negotiations == 0 {
182 return 0.0;
183 }
184 self.accepted as f64 / self.total_negotiations as f64
185 }
186}
187
188pub struct VersionedNegotiator {
196 protocols: HashMap<String, ProtocolMetadata>,
198 stats: NegotiationStats,
200}
201
202impl Default for VersionedNegotiator {
203 fn default() -> Self {
204 Self::new()
205 }
206}
207
208impl VersionedNegotiator {
209 pub fn new() -> Self {
211 Self {
212 protocols: HashMap::new(),
213 stats: NegotiationStats::default(),
214 }
215 }
216
217 pub fn register(&mut self, metadata: ProtocolMetadata) {
219 let key = metadata.header_value();
220 self.protocols.insert(key, metadata);
221 }
222
223 pub fn register_simple(&mut self, name: &str, version: &str) {
225 self.register(ProtocolMetadata::new(name, version));
226 }
227
228 pub fn unregister(&mut self, header_value: &str) -> bool {
230 self.protocols.remove(header_value).is_some()
231 }
232
233 pub fn contains(&self, header_value: &str) -> bool {
235 self.protocols.contains_key(header_value)
236 }
237
238 pub fn len(&self) -> usize {
240 self.protocols.len()
241 }
242
243 pub fn is_empty(&self) -> bool {
245 self.protocols.is_empty()
246 }
247
248 pub fn stats(&self) -> NegotiationStats {
250 self.stats.clone()
251 }
252
253 pub fn registered_protocols(&self) -> Vec<String> {
255 let mut keys: Vec<String> = self.protocols.keys().cloned().collect();
256 keys.sort();
257 keys
258 }
259
260 pub fn protocols_by_priority(&self) -> Vec<&ProtocolMetadata> {
263 let mut list: Vec<&ProtocolMetadata> = self.protocols.values().collect();
264 list.sort_by(|a, b| {
265 b.priority
266 .cmp(&a.priority)
267 .then_with(|| a.header_value().cmp(&b.header_value()))
268 });
269 list
270 }
271
272 pub fn negotiate(&mut self, client_protocols: &[String]) -> NegotiationOutcome {
282 self.stats.total_negotiations += 1;
283
284 if client_protocols.is_empty() {
285 self.stats.not_requested += 1;
286 return NegotiationOutcome::NotRequested;
287 }
288
289 let candidates: Vec<&ProtocolMetadata> = client_protocols
291 .iter()
292 .filter_map(|c| self.protocols.get(c))
293 .collect();
294
295 if candidates.is_empty() {
296 self.stats.no_match += 1;
297 return NegotiationOutcome::NoMatch {
298 requested: client_protocols.to_vec(),
299 };
300 }
301
302 let best = candidates
304 .iter()
305 .max_by_key(|m| m.priority)
306 .copied()
307 .expect("candidates is non-empty");
308
309 self.stats.accepted += 1;
310 NegotiationOutcome::Accepted {
311 header_value: best.header_value(),
312 metadata: best.clone(),
313 }
314 }
315
316 pub fn clear(&mut self) {
318 self.protocols.clear();
319 self.stats = NegotiationStats::default();
320 }
321}
322
323#[cfg(test)]
324mod tests {
325 use super::*;
326
327 #[test]
330 fn test_subprotocol_registry_new() {
331 let reg = SubProtocolRegistry::new();
332 assert!(reg.is_empty());
333 assert_eq!(reg.len(), 0);
334 }
335
336 #[test]
337 fn test_subprotocol_register_and_check() {
338 let mut reg = SubProtocolRegistry::new();
339 reg.register("json");
340 assert!(reg.is_registered("json"));
341 assert!(!reg.is_registered("xml"));
342 assert_eq!(reg.len(), 1);
343 }
344
345 #[test]
346 fn test_subprotocol_no_duplicate() {
347 let mut reg = SubProtocolRegistry::new();
348 reg.register("json");
349 reg.register("json");
350 assert_eq!(reg.len(), 1);
351 }
352
353 #[test]
354 fn test_subprotocol_register_many() {
355 let mut reg = SubProtocolRegistry::new();
356 reg.register_many(&["json", "xml", "protobuf"]);
357 assert_eq!(reg.len(), 3);
358 assert!(reg.is_registered("protobuf"));
359 }
360
361 #[test]
362 fn test_subprotocol_negotiate_matches_first() {
363 let mut reg = SubProtocolRegistry::new();
364 reg.register_many(&["json", "protobuf"]);
365
366 let client = vec![
367 "xml".to_string(),
368 "json".to_string(),
369 "protobuf".to_string(),
370 ];
371 let result = reg.negotiate(&client);
372 assert_eq!(result, Some("json".to_string()));
373 }
374
375 #[test]
376 fn test_subprotocol_negotiate_no_match() {
377 let reg = SubProtocolRegistry::new();
378 let client = vec!["xml".to_string(), "msgpack".to_string()];
379 let result = reg.negotiate(&client);
380 assert!(result.is_none());
381 }
382
383 #[test]
384 fn test_subprotocol_negotiate_empty_client() {
385 let mut reg = SubProtocolRegistry::new();
386 reg.register("json");
387 let client: Vec<String> = vec![];
388 assert!(reg.negotiate(&client).is_none());
389 }
390
391 #[test]
392 fn test_subprotocol_negotiate_empty_registry() {
393 let reg = SubProtocolRegistry::new();
394 let client = vec!["json".to_string()];
395 assert!(reg.negotiate(&client).is_none());
396 }
397
398 #[test]
399 fn test_subprotocol_clear() {
400 let mut reg = SubProtocolRegistry::new();
401 reg.register_many(&["json", "xml"]);
402 assert_eq!(reg.len(), 2);
403 reg.clear();
404 assert!(reg.is_empty());
405 }
406
407 #[test]
408 fn test_subprotocol_protocols_list() {
409 let mut reg = SubProtocolRegistry::new();
410 reg.register_many(&["json", "xml"]);
411 let list = reg.protocols();
412 assert_eq!(list.len(), 2);
413 assert!(list.contains(&"json".to_string()));
414 }
415
416 #[test]
417 fn test_subprotocol_preserves_registration_order() {
418 let mut reg = SubProtocolRegistry::new();
419 reg.register("c");
420 reg.register("a");
421 reg.register("b");
422 assert_eq!(reg.protocols(), &["c", "a", "b"]);
424 }
425
426 #[test]
427 fn test_subprotocol_negotiate_returns_client_order_not_registry_order() {
428 let mut reg = SubProtocolRegistry::new();
430 reg.register("a");
431 reg.register("b");
432 let client = vec!["b".to_string(), "a".to_string()];
433 assert_eq!(reg.negotiate(&client), Some("b".to_string()));
434 }
435
436 #[test]
439 fn test_protocol_metadata_new() {
440 let meta = ProtocolMetadata::new("chat", "1.0");
441 assert_eq!(meta.name, "chat");
442 assert_eq!(meta.version, "1.0");
443 assert_eq!(meta.priority, 0);
444 assert!(meta.description.is_empty());
445 }
446
447 #[test]
448 fn test_protocol_metadata_with_priority() {
449 let meta = ProtocolMetadata::new("chat", "1.0").with_priority(10);
450 assert_eq!(meta.priority, 10);
451 }
452
453 #[test]
454 fn test_protocol_metadata_with_description() {
455 let meta = ProtocolMetadata::new("chat", "1.0").with_description("Chat protocol v1");
456 assert_eq!(meta.description, "Chat protocol v1");
457 }
458
459 #[test]
460 fn test_protocol_metadata_header_value() {
461 let meta = ProtocolMetadata::new("chat", "1.0");
462 assert_eq!(meta.header_value(), "chat.1.0");
463 }
464
465 #[test]
466 fn test_protocol_metadata_header_value_with_complex_version() {
467 let meta = ProtocolMetadata::new("rpc", "2.1.3");
468 assert_eq!(meta.header_value(), "rpc.2.1.3");
469 }
470
471 #[test]
472 fn test_protocol_metadata_builder_chain() {
473 let meta = ProtocolMetadata::new("jsonrpc", "2.0")
474 .with_priority(5)
475 .with_description("JSON-RPC 2.0");
476 assert_eq!(meta.priority, 5);
477 assert_eq!(meta.description, "JSON-RPC 2.0");
478 assert_eq!(meta.header_value(), "jsonrpc.2.0");
479 }
480
481 #[test]
484 fn test_negotiation_outcome_is_accepted() {
485 let accepted = NegotiationOutcome::Accepted {
486 header_value: "chat.1.0".to_string(),
487 metadata: ProtocolMetadata::new("chat", "1.0"),
488 };
489 assert!(accepted.is_accepted());
490
491 let not_requested = NegotiationOutcome::NotRequested;
492 assert!(!not_requested.is_accepted());
493
494 let no_match = NegotiationOutcome::NoMatch {
495 requested: vec!["xml".to_string()],
496 };
497 assert!(!no_match.is_accepted());
498 }
499
500 #[test]
501 fn test_negotiation_outcome_header_value() {
502 let accepted = NegotiationOutcome::Accepted {
503 header_value: "chat.1.0".to_string(),
504 metadata: ProtocolMetadata::new("chat", "1.0"),
505 };
506 assert_eq!(accepted.header_value(), Some("chat.1.0"));
507
508 let not_requested = NegotiationOutcome::NotRequested;
509 assert_eq!(not_requested.header_value(), None);
510
511 let no_match = NegotiationOutcome::NoMatch { requested: vec![] };
512 assert_eq!(no_match.header_value(), None);
513 }
514
515 #[test]
518 fn test_negotiation_stats_default() {
519 let stats = NegotiationStats::default();
520 assert_eq!(stats.total_negotiations, 0);
521 assert_eq!(stats.accepted, 0);
522 assert_eq!(stats.not_requested, 0);
523 assert_eq!(stats.no_match, 0);
524 assert_eq!(stats.success_rate(), 0.0);
525 }
526
527 #[test]
528 fn test_negotiation_stats_success_rate_all_success() {
529 let stats = NegotiationStats {
530 total_negotiations: 10,
531 accepted: 10,
532 not_requested: 0,
533 no_match: 0,
534 };
535 assert!((stats.success_rate() - 1.0).abs() < 1e-9);
536 }
537
538 #[test]
539 fn test_negotiation_stats_success_rate_half() {
540 let stats = NegotiationStats {
541 total_negotiations: 10,
542 accepted: 5,
543 not_requested: 3,
544 no_match: 2,
545 };
546 assert!((stats.success_rate() - 0.5).abs() < 1e-9);
547 }
548
549 #[test]
550 fn test_negotiation_stats_success_rate_zero_total() {
551 let stats = NegotiationStats::default();
552 assert_eq!(stats.success_rate(), 0.0);
553 }
554
555 #[test]
558 fn test_versioned_negotiator_new_empty() {
559 let neg = VersionedNegotiator::new();
560 assert!(neg.is_empty());
561 assert_eq!(neg.len(), 0);
562 let stats = neg.stats();
563 assert_eq!(stats.total_negotiations, 0);
564 }
565
566 #[test]
567 fn test_versioned_negotiator_register() {
568 let mut neg = VersionedNegotiator::new();
569 neg.register(ProtocolMetadata::new("chat", "1.0"));
570 assert_eq!(neg.len(), 1);
571 assert!(neg.contains("chat.1.0"));
572 }
573
574 #[test]
575 fn test_versioned_negotiator_register_simple() {
576 let mut neg = VersionedNegotiator::new();
577 neg.register_simple("jsonrpc", "2.0");
578 assert!(neg.contains("jsonrpc.2.0"));
579 assert_eq!(neg.len(), 1);
580 }
581
582 #[test]
583 fn test_versioned_negotiator_unregister() {
584 let mut neg = VersionedNegotiator::new();
585 neg.register_simple("chat", "1.0");
586 assert!(neg.unregister("chat.1.0"));
587 assert!(!neg.contains("chat.1.0"));
588 assert_eq!(neg.len(), 0);
589 }
590
591 #[test]
592 fn test_versioned_negotiator_unregister_missing() {
593 let mut neg = VersionedNegotiator::new();
594 assert!(!neg.unregister("nonexistent"));
595 }
596
597 #[test]
598 fn test_versioned_negotiator_negotiate_success() {
599 let mut neg = VersionedNegotiator::new();
600 neg.register(ProtocolMetadata::new("chat", "1.0").with_priority(5));
601
602 let client = vec!["chat.1.0".to_string()];
603 let outcome = neg.negotiate(&client);
604 assert!(outcome.is_accepted());
605 assert_eq!(outcome.header_value(), Some("chat.1.0"));
606 }
607
608 #[test]
609 fn test_versioned_negotiator_not_requested() {
610 let mut neg = VersionedNegotiator::new();
611 neg.register_simple("chat", "1.0");
612
613 let client: Vec<String> = vec![];
614 let outcome = neg.negotiate(&client);
615 assert_eq!(outcome, NegotiationOutcome::NotRequested);
616
617 let stats = neg.stats();
618 assert_eq!(stats.not_requested, 1);
619 assert_eq!(stats.total_negotiations, 1);
620 }
621
622 #[test]
623 fn test_versioned_negotiator_no_match() {
624 let mut neg = VersionedNegotiator::new();
625 neg.register_simple("chat", "1.0");
626
627 let client = vec!["xml.1.0".to_string(), "msgpack.1.0".to_string()];
628 let outcome = neg.negotiate(&client);
629 match outcome {
630 NegotiationOutcome::NoMatch { requested } => {
631 assert_eq!(requested, client);
632 }
633 _ => panic!("expected NoMatch"),
634 }
635
636 let stats = neg.stats();
637 assert_eq!(stats.no_match, 1);
638 }
639
640 #[test]
641 fn test_versioned_negotiator_selects_highest_priority() {
642 let mut neg = VersionedNegotiator::new();
643 neg.register(ProtocolMetadata::new("chat", "1.0").with_priority(1));
644 neg.register(ProtocolMetadata::new("chat", "2.0").with_priority(10));
645 neg.register(ProtocolMetadata::new("chat", "1.5").with_priority(5));
646
647 let client = vec![
649 "chat.1.0".to_string(),
650 "chat.2.0".to_string(),
651 "chat.1.5".to_string(),
652 ];
653 let outcome = neg.negotiate(&client);
654 assert_eq!(outcome.header_value(), Some("chat.2.0"));
655 }
656
657 #[test]
658 fn test_versioned_negotiator_priority_tiebreak_client_order() {
659 let mut neg = VersionedNegotiator::new();
661 neg.register(ProtocolMetadata::new("a", "1.0").with_priority(5));
662 neg.register(ProtocolMetadata::new("b", "1.0").with_priority(5));
663
664 let client = vec!["b.1.0".to_string(), "a.1.0".to_string()];
665 let outcome = neg.negotiate(&client);
666 let header = outcome.header_value().expect("should accept");
670 assert!(header == "a.1.0" || header == "b.1.0");
671 }
672
673 #[test]
674 fn test_versioned_negotiator_stats_tracked_across_calls() {
675 let mut neg = VersionedNegotiator::new();
676 neg.register_simple("chat", "1.0");
677
678 neg.negotiate(&["chat.1.0".to_string()]);
680 neg.negotiate(&[]);
682 neg.negotiate(&["xml.1.0".to_string()]);
684 neg.negotiate(&["chat.1.0".to_string()]);
686
687 let stats = neg.stats();
688 assert_eq!(stats.total_negotiations, 4);
689 assert_eq!(stats.accepted, 2);
690 assert_eq!(stats.not_requested, 1);
691 assert_eq!(stats.no_match, 1);
692 assert!((stats.success_rate() - 0.5).abs() < 1e-9);
693 }
694
695 #[test]
696 fn test_versioned_negotiator_registered_protocols_sorted() {
697 let mut neg = VersionedNegotiator::new();
698 neg.register_simple("zebra", "1.0");
699 neg.register_simple("alpha", "1.0");
700 neg.register_simple("mango", "1.0");
701
702 let list = neg.registered_protocols();
703 assert_eq!(list, vec!["alpha.1.0", "mango.1.0", "zebra.1.0"]);
704 }
705
706 #[test]
707 fn test_versioned_negotiator_protocols_by_priority_descending() {
708 let mut neg = VersionedNegotiator::new();
709 neg.register(ProtocolMetadata::new("low", "1.0").with_priority(1));
710 neg.register(ProtocolMetadata::new("high", "1.0").with_priority(10));
711 neg.register(ProtocolMetadata::new("mid", "1.0").with_priority(5));
712
713 let sorted = neg.protocols_by_priority();
714 assert_eq!(sorted[0].name, "high");
715 assert_eq!(sorted[1].name, "mid");
716 assert_eq!(sorted[2].name, "low");
717 }
718
719 #[test]
720 fn test_versioned_negotiator_protocols_by_priority_tiebreak_alpha() {
721 let mut neg = VersionedNegotiator::new();
723 neg.register(ProtocolMetadata::new("zeta", "1.0").with_priority(5));
724 neg.register(ProtocolMetadata::new("alpha", "1.0").with_priority(5));
725
726 let sorted = neg.protocols_by_priority();
727 assert_eq!(sorted[0].name, "alpha");
728 assert_eq!(sorted[1].name, "zeta");
729 }
730
731 #[test]
732 fn test_versioned_negotiator_clear() {
733 let mut neg = VersionedNegotiator::new();
734 neg.register_simple("chat", "1.0");
735 neg.negotiate(&["chat.1.0".to_string()]);
736
737 neg.clear();
738 assert!(neg.is_empty());
739 let stats = neg.stats();
740 assert_eq!(stats.total_negotiations, 0);
741 }
742
743 #[test]
744 fn test_versioned_negotiator_overwrite_registration() {
745 let mut neg = VersionedNegotiator::new();
747 neg.register(ProtocolMetadata::new("chat", "1.0").with_priority(1));
748 neg.register(ProtocolMetadata::new("chat", "1.0").with_priority(10));
749
750 assert_eq!(neg.len(), 1);
751 let client = vec!["chat.1.0".to_string()];
752 let outcome = neg.negotiate(&client);
753 if let NegotiationOutcome::Accepted { metadata, .. } = outcome {
754 assert_eq!(metadata.priority, 10);
755 } else {
756 panic!("expected Accepted");
757 }
758 }
759
760 #[test]
761 fn test_versioned_negotiator_partial_client_match() {
762 let mut neg = VersionedNegotiator::new();
763 neg.register(ProtocolMetadata::new("chat", "1.0").with_priority(5));
764 neg.register(ProtocolMetadata::new("rpc", "2.0").with_priority(3));
765
766 let client = vec![
768 "xml.1.0".to_string(),
769 "rpc.2.0".to_string(),
770 "chat.1.0".to_string(),
771 ];
772 let outcome = neg.negotiate(&client);
773 assert_eq!(outcome.header_value(), Some("chat.1.0"));
775 }
776
777 #[test]
778 fn test_versioned_negotiator_default() {
779 let neg = VersionedNegotiator::default();
780 assert!(neg.is_empty());
781 }
782}