use ipfrs_core::error::{Error, Result};
use std::collections::HashMap;
use std::sync::Arc;
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct ProtocolVersion {
pub major: u16,
pub minor: u16,
pub patch: u16,
}
impl ProtocolVersion {
pub fn new(major: u16, minor: u16, patch: u16) -> Self {
Self {
major,
minor,
patch,
}
}
pub fn parse(s: &str) -> Result<Self> {
let parts: Vec<&str> = s.split('.').collect();
if parts.len() != 3 {
return Err(Error::Network(format!("Invalid version string: {}", s)));
}
let major = parts[0]
.parse()
.map_err(|e| Error::Network(format!("Invalid major version: {}", e)))?;
let minor = parts[1]
.parse()
.map_err(|e| Error::Network(format!("Invalid minor version: {}", e)))?;
let patch = parts[2]
.parse()
.map_err(|e| Error::Network(format!("Invalid patch version: {}", e)))?;
Ok(Self::new(major, minor, patch))
}
pub fn is_compatible_with(&self, other: &ProtocolVersion) -> bool {
self.major == other.major && self.minor >= other.minor
}
}
impl std::fmt::Display for ProtocolVersion {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}.{}.{}", self.major, self.minor, self.patch)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ProtocolId {
pub name: String,
pub version: ProtocolVersion,
}
impl ProtocolId {
pub fn new(name: String, version: ProtocolVersion) -> Self {
Self { name, version }
}
pub fn to_protocol_string(&self) -> String {
format!("/ipfrs/{}/{}", self.name, self.version)
}
pub fn parse(s: &str) -> Result<Self> {
let parts: Vec<&str> = s.trim_matches('/').split('/').collect();
if parts.len() != 3 || parts[0] != "ipfrs" {
return Err(Error::Network(format!("Invalid protocol string: {}", s)));
}
let name = parts[1].to_string();
let version = ProtocolVersion::parse(parts[2])?;
Ok(Self::new(name, version))
}
}
impl std::fmt::Display for ProtocolId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.to_protocol_string())
}
}
#[derive(Debug, Clone)]
pub struct ProtocolCapabilities {
pub features: Vec<String>,
pub max_message_size: usize,
pub supports_streaming: bool,
}
impl Default for ProtocolCapabilities {
fn default() -> Self {
Self {
features: Vec::new(),
max_message_size: 1024 * 1024, supports_streaming: false,
}
}
}
type BoxedProtocolHandler = Arc<parking_lot::RwLock<Box<dyn ProtocolHandler>>>;
pub trait ProtocolHandler: Send + Sync {
fn protocol_id(&self) -> ProtocolId;
fn capabilities(&self) -> ProtocolCapabilities {
ProtocolCapabilities::default()
}
fn handle_request(&mut self, request: &[u8]) -> Result<Vec<u8>>;
fn initialize(&mut self) -> Result<()> {
Ok(())
}
fn shutdown(&mut self) -> Result<()> {
Ok(())
}
}
pub struct ProtocolRegistry {
handlers: parking_lot::RwLock<HashMap<ProtocolId, BoxedProtocolHandler>>,
aliases: parking_lot::RwLock<HashMap<String, Vec<ProtocolVersion>>>,
}
impl ProtocolRegistry {
pub fn new() -> Self {
Self {
handlers: parking_lot::RwLock::new(HashMap::new()),
aliases: parking_lot::RwLock::new(HashMap::new()),
}
}
pub fn register(&self, handler: Box<dyn ProtocolHandler>) -> Result<()> {
let protocol_id = handler.protocol_id();
let mut handlers = self.handlers.write();
if handlers.contains_key(&protocol_id) {
return Err(Error::Network(format!(
"Protocol already registered: {}",
protocol_id
)));
}
let mut aliases = self.aliases.write();
aliases
.entry(protocol_id.name.clone())
.or_default()
.push(protocol_id.version.clone());
handlers.insert(protocol_id, Arc::new(parking_lot::RwLock::new(handler)));
Ok(())
}
pub fn unregister(&self, protocol_id: &ProtocolId) -> Result<()> {
let mut handlers = self.handlers.write();
if let Some(handler) = handlers.remove(protocol_id) {
let mut handler = handler.write();
handler.shutdown()?;
let mut aliases = self.aliases.write();
if let Some(versions) = aliases.get_mut(&protocol_id.name) {
versions.retain(|v| v != &protocol_id.version);
if versions.is_empty() {
aliases.remove(&protocol_id.name);
}
}
Ok(())
} else {
Err(Error::Network(format!(
"Protocol not registered: {}",
protocol_id
)))
}
}
pub fn get(&self, protocol_id: &ProtocolId) -> Option<BoxedProtocolHandler> {
let handlers = self.handlers.read();
handlers.get(protocol_id).cloned()
}
pub fn find_compatible(&self, name: &str, min_version: &ProtocolVersion) -> Option<ProtocolId> {
let aliases = self.aliases.read();
if let Some(versions) = aliases.get(name) {
let mut compatible_versions: Vec<_> = versions
.iter()
.filter(|v| v.is_compatible_with(min_version))
.collect();
compatible_versions.sort_by(|a, b| b.cmp(a));
if let Some(version) = compatible_versions.first() {
return Some(ProtocolId::new(name.to_string(), (*version).clone()));
}
}
None
}
pub fn list_protocols(&self) -> Vec<ProtocolId> {
let handlers = self.handlers.read();
handlers.keys().cloned().collect()
}
pub fn handle_request(&self, protocol_id: &ProtocolId, request: &[u8]) -> Result<Vec<u8>> {
if let Some(handler) = self.get(protocol_id) {
let mut handler = handler.write();
handler.handle_request(request)
} else {
Err(Error::Network(format!(
"No handler registered for protocol: {}",
protocol_id
)))
}
}
pub fn get_capabilities(&self, protocol_id: &ProtocolId) -> Option<ProtocolCapabilities> {
if let Some(handler) = self.get(protocol_id) {
let handler = handler.read();
Some(handler.capabilities())
} else {
None
}
}
pub fn shutdown_all(&self) -> Result<()> {
let handlers = self.handlers.write();
for handler in handlers.values() {
let mut handler = handler.write();
handler.shutdown()?;
}
Ok(())
}
}
impl Default for ProtocolRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
struct MockProtocolHandler {
id: ProtocolId,
}
impl MockProtocolHandler {
fn new(name: &str, version: ProtocolVersion) -> Self {
Self {
id: ProtocolId::new(name.to_string(), version),
}
}
}
impl ProtocolHandler for MockProtocolHandler {
fn protocol_id(&self) -> ProtocolId {
self.id.clone()
}
fn handle_request(&mut self, request: &[u8]) -> Result<Vec<u8>> {
Ok(request.to_vec())
}
}
#[test]
fn test_protocol_version_creation() {
let version = ProtocolVersion::new(1, 2, 3);
assert_eq!(version.major, 1);
assert_eq!(version.minor, 2);
assert_eq!(version.patch, 3);
}
#[test]
fn test_protocol_version_parse() {
let version = ProtocolVersion::parse("1.2.3")
.expect("test: valid version string '1.2.3' should parse");
assert_eq!(version.major, 1);
assert_eq!(version.minor, 2);
assert_eq!(version.patch, 3);
assert!(ProtocolVersion::parse("invalid").is_err());
assert!(ProtocolVersion::parse("1.2").is_err());
}
#[test]
fn test_protocol_version_compatibility() {
let v1_0_0 = ProtocolVersion::new(1, 0, 0);
let v1_1_0 = ProtocolVersion::new(1, 1, 0);
let v1_2_0 = ProtocolVersion::new(1, 2, 0);
let v2_0_0 = ProtocolVersion::new(2, 0, 0);
assert!(v1_2_0.is_compatible_with(&v1_0_0));
assert!(v1_1_0.is_compatible_with(&v1_0_0));
assert!(v1_0_0.is_compatible_with(&v1_0_0));
assert!(!v1_0_0.is_compatible_with(&v1_1_0));
assert!(!v2_0_0.is_compatible_with(&v1_0_0));
assert!(!v1_0_0.is_compatible_with(&v2_0_0));
}
#[test]
fn test_protocol_version_display() {
let version = ProtocolVersion::new(1, 2, 3);
assert_eq!(format!("{}", version), "1.2.3");
}
#[test]
fn test_protocol_id_creation() {
let version = ProtocolVersion::new(1, 0, 0);
let id = ProtocolId::new("test".to_string(), version);
assert_eq!(id.name, "test");
assert_eq!(id.version.major, 1);
}
#[test]
fn test_protocol_id_to_string() {
let version = ProtocolVersion::new(1, 0, 0);
let id = ProtocolId::new("tensorswap".to_string(), version);
assert_eq!(id.to_protocol_string(), "/ipfrs/tensorswap/1.0.0");
}
#[test]
fn test_protocol_id_parse() {
let id = ProtocolId::parse("/ipfrs/tensorswap/1.0.0")
.expect("test: valid protocol string '/ipfrs/tensorswap/1.0.0' should parse");
assert_eq!(id.name, "tensorswap");
assert_eq!(id.version.major, 1);
assert!(ProtocolId::parse("/invalid/tensorswap/1.0.0").is_err());
assert!(ProtocolId::parse("/ipfrs/tensorswap").is_err());
}
#[test]
fn test_registry_creation() {
let registry = ProtocolRegistry::new();
assert_eq!(registry.list_protocols().len(), 0);
}
#[test]
fn test_register_handler() {
let registry = ProtocolRegistry::new();
let handler = Box::new(MockProtocolHandler::new(
"test",
ProtocolVersion::new(1, 0, 0),
));
registry
.register(handler)
.expect("test: registering a fresh handler should succeed");
assert_eq!(registry.list_protocols().len(), 1);
}
#[test]
fn test_register_duplicate() {
let registry = ProtocolRegistry::new();
let handler1 = Box::new(MockProtocolHandler::new(
"test",
ProtocolVersion::new(1, 0, 0),
));
let handler2 = Box::new(MockProtocolHandler::new(
"test",
ProtocolVersion::new(1, 0, 0),
));
registry
.register(handler1)
.expect("test: registering handler1 should succeed");
assert!(registry.register(handler2).is_err());
}
#[test]
fn test_get_handler() {
let registry = ProtocolRegistry::new();
let version = ProtocolVersion::new(1, 0, 0);
let protocol_id = ProtocolId::new("test".to_string(), version.clone());
let handler = Box::new(MockProtocolHandler::new("test", version));
registry
.register(handler)
.expect("test: registering handler should succeed");
let retrieved = registry.get(&protocol_id);
assert!(retrieved.is_some());
}
#[test]
fn test_find_compatible() {
let registry = ProtocolRegistry::new();
let handler1 = Box::new(MockProtocolHandler::new(
"test",
ProtocolVersion::new(1, 0, 0),
));
let handler2 = Box::new(MockProtocolHandler::new(
"test",
ProtocolVersion::new(1, 1, 0),
));
let handler3 = Box::new(MockProtocolHandler::new(
"test",
ProtocolVersion::new(1, 2, 0),
));
registry
.register(handler1)
.expect("test: registering handler1 should succeed");
registry
.register(handler2)
.expect("test: registering handler2 should succeed");
registry
.register(handler3)
.expect("test: registering handler3 should succeed");
let min_version = ProtocolVersion::new(1, 0, 0);
let compatible = registry.find_compatible("test", &min_version);
assert!(compatible.is_some());
let compatible = compatible
.expect("test: find_compatible should return a result for a registered version");
assert_eq!(compatible.version.major, 1);
assert_eq!(compatible.version.minor, 2);
}
#[test]
fn test_unregister_handler() {
let registry = ProtocolRegistry::new();
let version = ProtocolVersion::new(1, 0, 0);
let protocol_id = ProtocolId::new("test".to_string(), version.clone());
let handler = Box::new(MockProtocolHandler::new("test", version));
registry
.register(handler)
.expect("test: registering handler should succeed");
registry
.unregister(&protocol_id)
.expect("test: unregistering a registered protocol should succeed");
assert_eq!(registry.list_protocols().len(), 0);
}
#[test]
fn test_handle_request() {
let registry = ProtocolRegistry::new();
let version = ProtocolVersion::new(1, 0, 0);
let protocol_id = ProtocolId::new("test".to_string(), version.clone());
let handler = Box::new(MockProtocolHandler::new("test", version));
registry
.register(handler)
.expect("test: registering handler should succeed");
let request = b"test request";
let response = registry
.handle_request(&protocol_id, request)
.expect("test: handle_request should return response for registered protocol");
assert_eq!(response, request);
}
#[test]
fn test_protocol_capabilities_default() {
let caps = ProtocolCapabilities::default();
assert_eq!(caps.max_message_size, 1024 * 1024);
assert!(!caps.supports_streaming);
}
}