dmsc/protocol/
mod.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of DMSC.
4//! The DMSC project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18#![allow(non_snake_case)]
19
20//! # Protocol Module
21//!
22//! This module provides protocol implementations for DMSC, including
23//! global protocol, private protocol, post-quantum cryptography, and
24//! integration features.
25//!
26//! ## Features
27//!
28//! - **DMSCProtocol**: Main protocol interface (trait definition)
29//! - **DMSCGlobalProtocol**: Global protocol implementation (basic implementation)
30//! - **DMSCPrivateProtocol**: Private protocol implementation (basic implementation)
31//! - **DMSCCrypto**: Cryptographic operations
32//! - **Post-Quantum Cryptography**: Kyber, Dilithium, Falcon implementations using liboqs
33//!
34//! ## Security Status
35//!
36//! This module now uses the **liboqs** library for post-quantum cryptography,
37//! which is:
38//! - The reference implementation from the NIST PQC competition
39//! - Actively maintained and regularly audited
40//! - **Suitable for production use**
41//!
42//! Post-Quantum Cryptography algorithms (Kyber, Dilithium, Falcon):
43//! - Based on NIST PQC competition algorithms
44//! - Have undergone formal security analysis
45//! - Constant-time implementations for side-channel resistance
46//! - Recommended for protecting sensitive data
47//!
48//! ## Recommendation
49//!
50//! For cryptographic operations, this module uses audited libraries:
51//! - liboqs - NIST PQC reference implementation
52//! - ring - Modern, audited crypto library
53//! - openssl - Industry-standard crypto library
54
55use std::collections::HashMap;
56use std::sync::atomic::{AtomicU64, Ordering};
57use std::sync::Arc;
58use std::time::{SystemTime, UNIX_EPOCH};
59use async_trait::async_trait;
60use tokio::sync::RwLock;
61use serde::{Serialize, Deserialize};
62
63use crate::core::{DMSCResult, DMSCError};
64
65#[cfg(feature = "pyo3")]
66use pyo3::prelude::*;
67
68/// Frame definitions for binary protocol encoding
69pub mod frames;
70pub use frames::{DMSCFrameBuilder, DMSCFrameParser};
71
72/// Post-quantum cryptography modules (requires oqs feature)
73#[cfg(feature = "oqs")]
74pub mod kyber;
75#[cfg(feature = "oqs")]
76pub mod dilithium;
77#[cfg(feature = "oqs")]
78pub mod falcon;
79#[cfg(feature = "oqs")]
80pub mod post_quantum;
81#[cfg(feature = "oqs")]
82pub use post_quantum::{
83    KyberKEM, KyberPublicKey, KyberSecretKey, KyberCiphertext,
84    DilithiumSigner, DilithiumPublicKey, DilithiumSecretKey, DilithiumSignature,
85    FalconSigner, FalconPublicKey, FalconSecretKey, FalconSignature,
86    DMSCPostQuantumAlgorithm, KEMResult,
87};
88
89/// Protocol type enumeration
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, std::hash::Hash)]
91#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
92pub enum DMSCProtocolType {
93    /// Standard global protocol
94    Global = 0,
95    /// Enhanced private protocol
96    Private = 1,
97}
98
99/// Protocol status
100#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
101#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
102pub enum DMSCProtocolStatus {
103    /// Protocol is inactive
104    Inactive,
105    /// Protocol is initializing
106    Initializing,
107    /// Protocol is ready
108    Ready,
109    /// Protocol has an error
110    Error,
111}
112
113/// Connection state
114#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
115#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
116pub enum DMSCConnectionState {
117    /// Connection is disconnected
118    Disconnected,
119    /// Connection is connecting
120    Connecting,
121    /// Connection is connected
122    Connected,
123    /// Connection is disconnecting
124    Disconnecting,
125}
126
127/// Security level
128#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
129#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
130pub enum DMSCSecurityLevel {
131    /// No security
132    None,
133    /// Standard security
134    Standard,
135    /// High security
136    High,
137    /// Military-grade security
138    Military,
139}
140
141/// Protocol configuration
142#[derive(Debug, Clone, Serialize, Deserialize)]
143#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
144pub struct DMSCProtocolConfig {
145    /// Default protocol type
146    pub default_protocol: DMSCProtocolType,
147    /// Whether security is enabled
148    pub enable_security: bool,
149    /// Security level
150    pub security_level: DMSCSecurityLevel,
151    /// Whether state synchronization is enabled
152    pub enable_state_sync: bool,
153    /// Whether performance optimization is enabled
154    pub performance_optimization: bool,
155}
156
157impl Default for DMSCProtocolConfig {
158    fn default() -> Self {
159        Self {
160            default_protocol: DMSCProtocolType::Global,
161            enable_security: true,
162            security_level: DMSCSecurityLevel::Standard,
163            enable_state_sync: true,
164            performance_optimization: true,
165        }
166    }
167}
168
169impl DMSCProtocolConfig {
170    /// Validate the configuration.
171    pub fn validate(&self) -> DMSCResult<()> {
172        if self.security_level == DMSCSecurityLevel::None && self.enable_security {
173            return Err(DMSCError::Config(
174                "Security level cannot be None when security is enabled".to_string()
175            ));
176        }
177
178        Ok(())
179    }
180
181    /// Create a secure configuration.
182    pub fn secure() -> Self {
183        Self {
184            default_protocol: DMSCProtocolType::Private,
185            enable_security: true,
186            security_level: DMSCSecurityLevel::High,
187            enable_state_sync: true,
188            performance_optimization: true,
189        }
190    }
191
192    /// Create a maximum security configuration with post-quantum cryptography.
193    pub fn maximum_security() -> Self {
194        Self {
195            default_protocol: DMSCProtocolType::Private,
196            enable_security: true,
197            security_level: DMSCSecurityLevel::Military,
198            enable_state_sync: true,
199            performance_optimization: false,
200        }
201    }
202}
203
204/// Protocol statistics
205#[derive(Debug, Clone, Serialize, Deserialize)]
206#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
207pub struct DMSCProtocolStats {
208    /// Total messages sent
209    pub messages_sent: u64,
210    /// Total messages received
211    pub messages_received: u64,
212    /// Total bytes sent
213    pub bytes_sent: u64,
214    /// Total bytes received
215    pub bytes_received: u64,
216    /// Total errors
217    pub errors: u64,
218    /// Average latency in milliseconds
219    pub avg_latency_ms: f64,
220}
221
222impl DMSCProtocolStats {
223    pub fn new() -> Self {
224        Self {
225            messages_sent: 0,
226            messages_received: 0,
227            bytes_sent: 0,
228            bytes_received: 0,
229            errors: 0,
230            avg_latency_ms: 0.0,
231        }
232    }
233
234    pub fn record_sent(&mut self, bytes: usize) {
235        self.messages_sent += 1;
236        self.bytes_sent += bytes as u64;
237    }
238
239    pub fn record_received(&mut self, bytes: usize) {
240        self.messages_received += 1;
241        self.bytes_received += bytes as u64;
242    }
243
244    pub fn record_error(&mut self) {
245        self.errors += 1;
246    }
247}
248
249/// Connection statistics
250#[derive(Debug, Clone, Serialize, Deserialize)]
251#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
252pub struct DMSCConnectionStats {
253    /// Total connections
254    pub total_connections: u64,
255    /// Active connections
256    pub active_connections: u64,
257    /// Total bytes sent
258    pub bytes_sent: u64,
259    /// Total bytes received
260    pub bytes_received: u64,
261    /// Connection duration in seconds
262    pub connection_duration_secs: u64,
263}
264
265impl Default for DMSCConnectionStats {
266    fn default() -> Self {
267        Self {
268            total_connections: 0,
269            active_connections: 0,
270            bytes_sent: 0,
271            bytes_received: 0,
272            connection_duration_secs: 0,
273        }
274    }
275}
276
277/// Protocol health status
278#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
279#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
280pub enum DMSCProtocolHealth {
281    /// Healthy
282    Healthy,
283    /// Degraded
284    Degraded,
285    /// Unhealthy
286    Unhealthy,
287    /// Unknown
288    Unknown,
289}
290
291impl Default for DMSCProtocolHealth {
292    fn default() -> Self {
293        DMSCProtocolHealth::Unknown
294    }
295}
296
297/// Message flags for protocol messages
298#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
299#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
300pub struct DMSCMessageFlags {
301    /// Whether the message is compressed
302    pub compressed: bool,
303    /// Whether the message is encrypted
304    pub encrypted: bool,
305    /// Whether the message requires acknowledgment
306    pub requires_ack: bool,
307    /// Whether this is a priority message
308    pub priority: bool,
309}
310
311impl Default for DMSCMessageFlags {
312    fn default() -> Self {
313        Self {
314            compressed: false,
315            encrypted: false,
316            requires_ack: false,
317            priority: false,
318        }
319    }
320}
321
322/// Connection information
323#[derive(Debug, Clone, Serialize, Deserialize)]
324#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
325pub struct DMSCConnectionInfo {
326    /// Connection ID
327    pub connection_id: String,
328    /// Remote device ID
329    pub device_id: String,
330    /// Connection address
331    pub address: String,
332    /// Protocol type
333    pub protocol_type: DMSCProtocolType,
334    /// Connection state
335    pub state: DMSCConnectionState,
336    /// Security level
337    pub security_level: DMSCSecurityLevel,
338    /// Connection timestamp
339    pub connected_at: u64,
340    /// Last activity timestamp
341    pub last_activity: u64,
342}
343
344impl Default for DMSCConnectionInfo {
345    fn default() -> Self {
346        Self {
347            connection_id: String::new(),
348            device_id: String::new(),
349            address: String::new(),
350            protocol_type: DMSCProtocolType::Global,
351            state: DMSCConnectionState::Disconnected,
352            security_level: DMSCSecurityLevel::None,
353            connected_at: 0,
354            last_activity: 0,
355        }
356    }
357}
358
359/// Frame type enumeration
360#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
361#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
362pub enum DMSCFrameType {
363    /// Data frame
364    Data = 0,
365    /// Control frame
366    Control = 1,
367    /// Heartbeat frame
368    Heartbeat = 2,
369    /// Acknowledgment frame
370    Ack = 3,
371    /// Error frame
372    Error = 4,
373}
374
375/// Frame header
376#[derive(Debug, Clone, Serialize, Deserialize)]
377#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
378pub struct DMSCFrameHeader {
379    /// Protocol version (major.minor as u8)
380    pub version: u8,
381    /// Frame type
382    pub frame_type: DMSCFrameType,
383    /// Sequence number
384    pub sequence_number: u64,
385    /// Message length
386    pub length: u32,
387    /// Timestamp
388    pub timestamp: u64,
389    /// Flags
390    pub flags: u16,
391    /// Authentication tag offset (for authenticated frames)
392    pub auth_tag_offset: u16,
393}
394
395impl Default for DMSCFrameHeader {
396    fn default() -> Self {
397        Self {
398            version: 1,
399            frame_type: DMSCFrameType::Data,
400            sequence_number: 0,
401            length: 0,
402            timestamp: 0,
403            flags: 0,
404            auth_tag_offset: 0,
405        }
406    }
407}
408
409impl DMSCFrameHeader {
410    /// Get major version number.
411    pub fn major_version(&self) -> u8 {
412        self.version >> 4
413    }
414
415    /// Get minor version number.
416    pub fn minor_version(&self) -> u8 {
417        self.version & 0x0F
418    }
419
420    /// Check if a feature is supported.
421    pub fn supports_feature(&self, feature: u16) -> bool {
422        (self.flags & feature) != 0
423    }
424}
425
426/// Protocol frame
427#[derive(Debug, Clone, Serialize, Deserialize)]
428#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
429pub struct DMSCFrame {
430    /// Frame header
431    pub header: DMSCFrameHeader,
432    /// Frame payload
433    pub payload: Vec<u8>,
434    /// Source device ID
435    pub source_id: String,
436    /// Target device ID
437    pub target_id: String,
438}
439
440impl Default for DMSCFrame {
441    fn default() -> Self {
442        Self {
443            header: DMSCFrameHeader::default(),
444            payload: Vec::new(),
445            source_id: String::new(),
446            target_id: String::new(),
447        }
448    }
449}
450
451/// Core protocol trait
452#[async_trait]
453pub trait DMSCProtocol {
454    /// Get protocol type
455    fn protocol_type(&self) -> DMSCProtocolType;
456    
457    /// Check if protocol is ready
458    async fn is_ready(&self) -> bool;
459    
460    /// Initialize protocol
461    async fn initialize(&mut self, config: DMSCProtocolConfig) -> DMSCResult<()>;
462    
463    /// Send message
464    async fn send_message(&mut self, target: &str, data: &[u8]) -> DMSCResult<Vec<u8>>;
465    
466    /// Send message with flags
467    async fn send_message_with_flags(&mut self, target: &str, data: &[u8], flags: DMSCMessageFlags) -> DMSCResult<Vec<u8>>;
468    
469    /// Receive message
470    async fn receive_message(&mut self) -> DMSCResult<Vec<u8>>;
471    
472    /// Get connection info
473    async fn get_connection_info(&self, connection_id: &str) -> DMSCResult<DMSCConnectionInfo>;
474    
475    /// Close connection
476    async fn close_connection(&mut self, connection_id: &str) -> DMSCResult<()>;
477    
478    /// Get protocol statistics
479    fn get_stats(&self) -> DMSCProtocolStats;
480    
481    /// Get protocol health
482    async fn get_health(&self) -> DMSCProtocolHealth;
483    
484    /// Shutdown protocol
485    async fn shutdown(&mut self) -> DMSCResult<()>;
486}
487
488/// Protocol connection trait
489#[async_trait]
490pub trait DMSCProtocolConnection {
491    /// Get connection ID
492    fn connection_id(&self) -> &str;
493    
494    /// Get remote device ID
495    fn remote_device_id(&self) -> &str;
496    
497    /// Get protocol type
498    fn protocol_type(&self) -> DMSCProtocolType;
499    
500    /// Check if connection is active
501    fn is_active(&self) -> bool;
502    
503    /// Send data
504    async fn send(&mut self, data: &[u8]) -> DMSCResult<usize>;
505    
506    /// Receive data
507    async fn receive(&mut self, buffer: &mut [u8]) -> DMSCResult<usize>;
508    
509    /// Get statistics
510    fn get_stats(&self) -> DMSCConnectionStats;
511    
512    /// Close connection
513    async fn close(&mut self) -> DMSCResult<()>;
514}
515
516/// Protocol manager
517#[derive(Debug, Clone)]
518#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
519pub struct DMSCProtocolManager {
520    /// Protocol statistics
521    pub stats: Arc<RwLock<DMSCProtocolStats>>,
522    /// Default protocol type
523    pub default_protocol: DMSCProtocolType,
524    /// Active connections
525    connections: Arc<RwLock<HashMap<String, DMSCConnectionInfo>>>,
526    /// Message sequence counter
527    sequence_counter: Arc<AtomicU64>,
528    /// Protocol initialized state
529    initialized: Arc<RwLock<bool>>,
530}
531
532#[cfg(test)]
533mod protocol_tests {
534    use super::*;
535
536    #[test]
537    fn test_protocol_config_default() {
538        let config = DMSCProtocolConfig::default();
539        assert_eq!(config.default_protocol, DMSCProtocolType::Global);
540        assert!(config.enable_security);
541        assert_eq!(config.security_level, DMSCSecurityLevel::Standard);
542    }
543
544    #[test]
545    fn test_protocol_config_secure() {
546        let config = DMSCProtocolConfig::secure();
547        assert_eq!(config.default_protocol, DMSCProtocolType::Private);
548        assert!(config.enable_security);
549        assert_eq!(config.security_level, DMSCSecurityLevel::High);
550    }
551
552    #[test]
553    fn test_protocol_config_maximum_security() {
554        let config = DMSCProtocolConfig::maximum_security();
555        assert_eq!(config.default_protocol, DMSCProtocolType::Private);
556        assert!(config.enable_security);
557        assert_eq!(config.security_level, DMSCSecurityLevel::Military);
558    }
559
560    #[test]
561    fn test_protocol_config_validation() {
562        let mut config = DMSCProtocolConfig::default();
563
564        // Valid config should pass
565        assert!(config.validate().is_ok());
566
567        // None security level with security enabled should fail
568        config.security_level = DMSCSecurityLevel::None;
569        assert!(config.validate().is_err());
570    }
571
572    #[test]
573    fn test_frame_header_version() {
574        let header = DMSCFrameHeader::default();
575        assert_eq!(header.major_version(), 0);
576        assert_eq!(header.minor_version(), 1);
577    }
578
579    #[test]
580    fn test_frame_header_supports_feature() {
581        let header = DMSCFrameHeader {
582            flags: 0b00000011,
583            ..Default::default()
584        };
585
586        assert!(header.supports_feature(0b00000001));
587        assert!(header.supports_feature(0b00000010));
588        assert!(!header.supports_feature(0b00000100));
589    }
590
591    #[test]
592    fn test_protocol_stats() {
593        let stats = DMSCProtocolStats::new();
594        assert_eq!(stats.messages_sent, 0);
595        assert_eq!(stats.messages_received, 0);
596        assert_eq!(stats.errors, 0);
597    }
598
599    #[test]
600    fn test_connection_info() {
601        let mut info = DMSCConnectionInfo::default();
602        info.device_id = "test-device".to_string();
603        info.state = DMSCConnectionState::Connected;
604        assert_eq!(info.device_id, "test-device");
605        assert_eq!(info.state, DMSCConnectionState::Connected);
606        assert_eq!(info.security_level, DMSCSecurityLevel::None);
607    }
608
609    #[test]
610    fn test_protocol_type_values() {
611        assert_eq!(DMSCProtocolType::Global as u8, 0);
612        assert_eq!(DMSCProtocolType::Private as u8, 1);
613    }
614
615    #[test]
616    fn test_security_level_values() {
617        assert_eq!(DMSCSecurityLevel::None as u8, 0);
618        assert_eq!(DMSCSecurityLevel::Standard as u8, 1);
619        assert_eq!(DMSCSecurityLevel::High as u8, 2);
620        assert_eq!(DMSCSecurityLevel::Military as u8, 3);
621    }
622}
623
624#[cfg(feature = "pyo3")]
625#[pyo3::prelude::pymethods]
626impl DMSCProtocolManager {
627    #[new]
628    fn new_py() -> Self {
629        Self::new()
630    }
631    
632    #[getter]
633    fn get_stats_py(&self) -> DMSCProtocolStats {
634        self.stats.try_read()
635            .map(|guard| guard.clone())
636            .unwrap_or_else(|_| DMSCProtocolStats::new())
637    }
638    
639    #[getter]
640    fn get_default_protocol_py(&self) -> DMSCProtocolType {
641        self.default_protocol
642    }
643    
644    /// Initialize manager
645    pub fn initialize(&mut self, config: DMSCProtocolConfig) -> PyResult<()> {
646        self.default_protocol = config.default_protocol;
647        
648        let mut initialized = self.initialized.try_write()
649            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
650                "Failed to acquire write lock on initialized state"))?;
651        *initialized = true;
652        
653        Ok(())
654    }
655    
656    /// Send message (sync version for Python)
657    pub fn send_message(&self, target: &str, data: &[u8]) -> Vec<u8> {
658        let timestamp = SystemTime::now()
659            .duration_since(UNIX_EPOCH)
660            .unwrap_or_default()
661            .as_millis() as u64;
662        
663        let sequence = self.sequence_counter.fetch_add(1, Ordering::SeqCst);
664        
665        let frame = DMSCFrame {
666            header: DMSCFrameHeader {
667                version: 1,
668                frame_type: DMSCFrameType::Data,
669                sequence_number: sequence,
670                length: data.len() as u32,
671                timestamp,
672                flags: 0,
673                auth_tag_offset: 0,
674            },
675            payload: data.to_vec(),
676            source_id: String::from("protocol_manager"),
677            target_id: target.to_string(),
678        };
679        
680        let serialized = match serde_json::to_vec(&frame) {
681            Ok(serialized_data) => serialized_data,
682            Err(e) => {
683                if let Ok(mut stats) = self.stats.try_write() {
684                    stats.record_error();
685                }
686                let error_response = DMSCProtocolResponse {
687                    success: false,
688                    sequence_number: sequence,
689                    target_id: target.to_string(),
690                    response_data: format!("Serialization error: {}", e).into_bytes(),
691                    timestamp,
692                };
693                return serde_json::to_vec(&error_response)
694                    .unwrap_or_else(|_| b"{\"success\":false,\"error\":\"Serialization failed\"}".to_vec());
695            }
696        };
697        
698        let payload_len = serialized.len();
699        if let Ok(mut stats) = self.stats.try_write() {
700            stats.record_sent(payload_len);
701        }
702        
703        let response_data = self.build_response_data(target, &frame, sequence, timestamp);
704        
705        let response = DMSCProtocolResponse {
706            success: true,
707            sequence_number: sequence,
708            target_id: target.to_string(),
709            response_data,
710            timestamp,
711        };
712        
713        self.stats.try_write()
714            .map(|mut stats| stats.record_received(response.response_data.len()))
715            .map_err(|e| tracing::error!("Failed to update protocol stats: {}", e))
716            .ok();
717        
718        serde_json::to_vec(&response).unwrap_or_else(|_| b"{\"success\":true,\"message\":\"Message sent\"}".to_vec())
719    }
720    
721    fn build_response_data(&self, target: &str, frame: &DMSCFrame, sequence: u64, timestamp: u64) -> Vec<u8> {
722        let mut response = HashMap::<String, serde_json::Value>::new();
723        
724        response.insert("status".to_string(), serde_json::Value::String("delivered".to_string()));
725        response.insert("target".to_string(), serde_json::Value::String(target.to_string()));
726        response.insert("source".to_string(), serde_json::Value::String(frame.source_id.clone()));
727        response.insert("sequence".to_string(), serde_json::Value::Number(sequence.into()));
728        response.insert("timestamp".to_string(), serde_json::Value::Number(timestamp.into()));
729        response.insert("frame_type".to_string(), serde_json::Value::String(format!("{:?}", frame.header.frame_type)));
730        response.insert("payload_size".to_string(), serde_json::Value::Number(serde_json::Number::from(frame.payload.len())));
731        response.insert("protocol".to_string(), serde_json::Value::String(format!("{:?}", self.default_protocol)));
732        
733        let delivery_info = serde_json::json!({
734            "delivered_at": timestamp,
735            "hops": 1,
736            "route": [frame.source_id.clone(), target.to_string()]
737        });
738        response.insert("delivery".to_string(), delivery_info);
739        
740        serde_json::to_vec(&response).unwrap_or_default()
741    }
742    
743    /// Send message with flags (sync version for Python)
744    pub fn send_message_with_flags(&self, target: &str, data: &[u8], _flags: DMSCMessageFlags) -> Vec<u8> {
745        self.send_message(target, data)
746    }
747    
748    /// Get connection info (sync version for Python)
749    pub fn get_connection_info(&self, connection_id: &str) -> Option<DMSCConnectionInfo> {
750        self.connections.try_read()
751            .ok()
752            .and_then(|connections| connections.get(connection_id).cloned())
753    }
754    
755    /// Close connection (sync version for Python)
756    pub fn close_connection(&mut self, connection_id: &str) -> bool {
757        self.connections.try_write()
758            .ok()
759            .map(|mut connections| connections.remove(connection_id).is_some())
760            .unwrap_or(false)
761    }
762}
763
764impl DMSCProtocolManager {
765    pub fn new() -> Self {
766        Self {
767            stats: Arc::new(RwLock::new(DMSCProtocolStats::new())),
768            default_protocol: DMSCProtocolType::Global,
769            connections: Arc::new(RwLock::new(HashMap::new())),
770            sequence_counter: Arc::new(AtomicU64::new(0)),
771            initialized: Arc::new(RwLock::new(false)),
772        }
773    }
774}
775
776impl Default for DMSCProtocolManager {
777    fn default() -> Self {
778        Self::new()
779    }
780}
781
782/// Protocol response structure
783#[derive(Debug, Clone, Serialize, Deserialize)]
784#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
785pub struct DMSCProtocolResponse {
786    /// Whether the operation was successful
787    pub success: bool,
788    /// Sequence number matching the request
789    pub sequence_number: u64,
790    /// Target ID that was addressed
791    pub target_id: String,
792    /// Response data payload
793    pub response_data: Vec<u8>,
794    /// Timestamp of the original request
795    pub timestamp: u64,
796}
797
798impl Default for DMSCProtocolResponse {
799    fn default() -> Self {
800        Self {
801            success: false,
802            sequence_number: 0,
803            target_id: String::new(),
804            response_data: Vec::new(),
805            timestamp: 0,
806        }
807    }
808}
809
810#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
811pub enum ProtocolError {
812    #[error("Target not found: {target_id}")]
813    TargetNotFound { target_id: String },
814    #[error("Send failed: {message}")]
815    SendFailed { message: String },
816    #[error("Protocol not initialized")]
817    NotInitialized,
818    #[error("Invalid state: {state}")]
819    InvalidState { state: String },
820    #[error("Connection not found: {connection_id}")]
821    ConnectionNotFound { connection_id: String },
822    #[error("Serialization error: {message}")]
823    Serialization { message: String },
824    #[error("Operation not supported for this protocol")]
825    NotSupported,
826}
827
828impl From<ProtocolError> for DMSCError {
829    fn from(error: ProtocolError) -> Self {
830        DMSCError::Other(format!("Protocol error: {}", error))
831    }
832}
833
834impl From<serde_json::Error> for ProtocolError {
835    fn from(error: serde_json::Error) -> Self {
836        ProtocolError::Serialization { message: error.to_string() }
837    }
838}
839
840impl From<ProtocolError> for DMSCResult<()> {
841    fn from(error: ProtocolError) -> Self {
842        Err(error.into())
843    }
844}
845
846#[derive(Debug, Clone)]
847pub struct DMSCBaseProtocol {
848    config: DMSCProtocolConfig,
849    stats: Arc<RwLock<DMSCProtocolStats>>,
850    connections: Arc<RwLock<HashMap<String, DMSCConnectionInfo>>>,
851    sequence_counter: Arc<AtomicU64>,
852    initialized: Arc<RwLock<bool>>,
853    _receiver_id: String,
854}
855
856impl DMSCBaseProtocol {
857    pub fn new(receiver_id: String) -> Self {
858        Self {
859            config: DMSCProtocolConfig::default(),
860            stats: Arc::new(RwLock::new(DMSCProtocolStats::new())),
861            connections: Arc::new(RwLock::new(HashMap::new())),
862            sequence_counter: Arc::new(AtomicU64::new(0)),
863            initialized: Arc::new(RwLock::new(false)),
864            _receiver_id: receiver_id,
865        }
866    }
867    
868    pub async fn is_ready(&self) -> bool {
869        *self.initialized.read().await
870    }
871    
872    pub async fn initialize(&mut self, config: DMSCProtocolConfig) {
873        self.config = config;
874        *self.initialized.write().await = true;
875    }
876
877    pub async fn send_message(&mut self, _target: &str, data: &[u8]) -> DMSCResult<Vec<u8>> {
878        if !*self.initialized.read().await {
879            return Err(ProtocolError::NotInitialized.into());
880        }
881
882        let sequence = self.sequence_counter.fetch_add(1, Ordering::SeqCst) as u32;
883
884        self.stats.write().await.record_sent(data.len());
885
886        let mut builder = DMSCFrameBuilder::new();
887        builder.set_sequence(sequence);
888        let frame = builder.build_data_frame(data.to_vec())
889            .map_err(|e| ProtocolError::Serialization {
890                message: e.to_string()
891            })?;
892
893        let frame_bytes = frame.to_bytes()
894            .map_err(|e| ProtocolError::Serialization {
895                message: e.to_string()
896            })?;
897
898        self.stats.write().await.record_received(frame_bytes.len());
899
900        Ok(frame_bytes)
901    }
902    
903    pub async fn receive_message(&mut self) -> DMSCResult<Vec<u8>> {
904        if !*self.initialized.read().await {
905            return Err(ProtocolError::NotInitialized.into());
906        }
907
908        let sequence = self.sequence_counter.fetch_add(1, Ordering::SeqCst) as u32;
909
910        let mut builder = DMSCFrameBuilder::new();
911        builder.set_sequence(sequence);
912        let frame = builder.build_keepalive_frame()
913            .map_err(|e| ProtocolError::Serialization {
914                message: e.to_string()
915            })?;
916
917        let frame_bytes = frame.to_bytes()
918            .map_err(|e| ProtocolError::Serialization {
919                message: e.to_string()
920            })?;
921
922        self.stats.write().await.record_received(0);
923
924        Ok(frame_bytes)
925    }
926    
927    pub async fn get_connection_info(&self, connection_id: &str) -> DMSCResult<DMSCConnectionInfo> {
928        let connections = self.connections.read().await;
929        connections.get(connection_id)
930            .cloned()
931            .ok_or_else(|| ProtocolError::ConnectionNotFound {
932                connection_id: connection_id.to_string()
933            }.into())
934    }
935    
936    pub async fn close_connection(&mut self, connection_id: &str) -> DMSCResult<()> {
937        let mut connections = self.connections.write().await;
938        if connections.remove(connection_id).is_some() {
939            Ok(())
940        } else {
941            Err(ProtocolError::ConnectionNotFound {
942                connection_id: connection_id.to_string()
943            }.into())
944        }
945    }
946    
947    pub fn get_stats(&self) -> DMSCProtocolStats {
948        self.stats.try_read()
949            .map(|guard| guard.clone())
950            .unwrap_or_else(|_| DMSCProtocolStats::new())
951    }
952    
953    pub async fn get_health(&self) -> DMSCProtocolHealth {
954        if *self.initialized.read().await {
955            DMSCProtocolHealth::Healthy
956        } else {
957            DMSCProtocolHealth::Unknown
958        }
959    }
960    
961    pub async fn shutdown(&mut self) {
962        *self.initialized.write().await = false;
963    }
964}
965
966#[derive(Debug, Clone)]
967#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
968pub struct DMSCGlobalProtocol {
969    base: DMSCBaseProtocol,
970}
971
972impl DMSCGlobalProtocol {
973    pub fn new() -> Self {
974        Self {
975            base: DMSCBaseProtocol::new(String::from("receiver")),
976        }
977    }
978}
979
980impl Default for DMSCGlobalProtocol {
981    fn default() -> Self {
982        Self::new()
983    }
984}
985
986#[async_trait]
987impl DMSCProtocol for DMSCGlobalProtocol {
988    fn protocol_type(&self) -> DMSCProtocolType {
989        DMSCProtocolType::Global
990    }
991
992    async fn is_ready(&self) -> bool {
993        self.base.is_ready().await
994    }
995
996    async fn initialize(&mut self, config: DMSCProtocolConfig) -> DMSCResult<()> {
997        self.base.initialize(config).await;
998        Ok(())
999    }
1000
1001    async fn send_message(&mut self, target: &str, data: &[u8]) -> DMSCResult<Vec<u8>> {
1002        self.base.send_message(target, data).await
1003    }
1004
1005    async fn send_message_with_flags(&mut self, target: &str, data: &[u8], flags: DMSCMessageFlags) -> DMSCResult<Vec<u8>> {
1006        let response = self.base.send_message(target, data).await?;
1007        if flags.encrypted {
1008            self.base.stats.write().await.record_error();
1009        }
1010        Ok(response)
1011    }
1012
1013    async fn receive_message(&mut self) -> DMSCResult<Vec<u8>> {
1014        self.base.receive_message().await
1015    }
1016
1017    async fn get_connection_info(&self, connection_id: &str) -> DMSCResult<DMSCConnectionInfo> {
1018        self.base.get_connection_info(connection_id).await
1019    }
1020
1021    async fn close_connection(&mut self, connection_id: &str) -> DMSCResult<()> {
1022        self.base.close_connection(connection_id).await
1023    }
1024
1025    fn get_stats(&self) -> DMSCProtocolStats {
1026        self.base.get_stats()
1027    }
1028
1029    async fn get_health(&self) -> DMSCProtocolHealth {
1030        self.base.get_health().await
1031    }
1032
1033    async fn shutdown(&mut self) -> DMSCResult<()> {
1034        self.base.shutdown().await;
1035        Ok(())
1036    }
1037}
1038
1039#[derive(Debug, Clone)]
1040#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
1041pub struct DMSCPrivateProtocol {
1042    base: DMSCBaseProtocol,
1043}
1044
1045impl DMSCPrivateProtocol {
1046    pub fn new() -> Self {
1047        Self {
1048            base: DMSCBaseProtocol::new(String::from("private_receiver")),
1049        }
1050    }
1051}
1052
1053impl Default for DMSCPrivateProtocol {
1054    fn default() -> Self {
1055        Self::new()
1056    }
1057}
1058
1059#[async_trait]
1060impl DMSCProtocol for DMSCPrivateProtocol {
1061    fn protocol_type(&self) -> DMSCProtocolType {
1062        DMSCProtocolType::Private
1063    }
1064
1065    async fn is_ready(&self) -> bool {
1066        self.base.is_ready().await
1067    }
1068
1069    async fn initialize(&mut self, config: DMSCProtocolConfig) -> DMSCResult<()> {
1070        self.base.initialize(config).await;
1071        Ok(())
1072    }
1073
1074    async fn send_message(&mut self, target: &str, data: &[u8]) -> DMSCResult<Vec<u8>> {
1075        self.base.send_message(target, data).await
1076    }
1077
1078    async fn send_message_with_flags(&mut self, target: &str, data: &[u8], flags: DMSCMessageFlags) -> DMSCResult<Vec<u8>> {
1079        let response = self.base.send_message(target, data).await?;
1080        if !flags.encrypted {
1081            self.base.stats.write().await.record_error();
1082        }
1083        Ok(response)
1084    }
1085
1086    async fn receive_message(&mut self) -> DMSCResult<Vec<u8>> {
1087        self.base.receive_message().await
1088    }
1089
1090    async fn get_connection_info(&self, connection_id: &str) -> DMSCResult<DMSCConnectionInfo> {
1091        self.base.get_connection_info(connection_id).await
1092    }
1093
1094    async fn close_connection(&mut self, connection_id: &str) -> DMSCResult<()> {
1095        self.base.close_connection(connection_id).await
1096    }
1097
1098    fn get_stats(&self) -> DMSCProtocolStats {
1099        self.base.get_stats()
1100    }
1101
1102    async fn get_health(&self) -> DMSCProtocolHealth {
1103        self.base.get_health().await
1104    }
1105
1106    async fn shutdown(&mut self) -> DMSCResult<()> {
1107        self.base.shutdown().await;
1108        Ok(())
1109    }
1110}
1111
1112#[cfg(feature = "pyo3")]
1113#[pyo3::prelude::pymethods]
1114impl DMSCGlobalProtocol {
1115    pub fn is_ready_sync(&self) -> bool {
1116        self.base.initialized.try_read()
1117            .map(|guard| *guard)
1118            .unwrap_or(false)
1119    }
1120
1121    pub fn initialize(&mut self, config: DMSCProtocolConfig) -> bool {
1122        self.base.config = config;
1123        if let Ok(mut guard) = self.base.initialized.try_write() {
1124            *guard = true;
1125            return true;
1126        }
1127        false
1128    }
1129
1130    pub fn get_stats(&self) -> DMSCProtocolStats {
1131        self.base.stats.try_read()
1132            .map(|guard| guard.clone())
1133            .unwrap_or_else(|_| DMSCProtocolStats::new())
1134    }
1135
1136    pub fn get_health(&self) -> DMSCProtocolHealth {
1137        self.base.initialized.try_read()
1138            .map(|guard| {
1139                if *guard {
1140                    DMSCProtocolHealth::Healthy
1141                } else {
1142                    DMSCProtocolHealth::Unknown
1143                }
1144            })
1145            .unwrap_or(DMSCProtocolHealth::Unknown)
1146    }
1147
1148    pub fn shutdown(&mut self) -> bool {
1149        if let Ok(mut guard) = self.base.initialized.try_write() {
1150            *guard = false;
1151            return true;
1152        }
1153        false
1154    }
1155}
1156
1157#[cfg(feature = "pyo3")]
1158#[pyo3::prelude::pymethods]
1159impl DMSCPrivateProtocol {
1160    pub fn is_ready_sync(&self) -> bool {
1161        self.base.initialized.try_read()
1162            .map(|guard| *guard)
1163            .unwrap_or(false)
1164    }
1165
1166    pub fn initialize(&mut self, config: DMSCProtocolConfig) -> bool {
1167        self.base.config = config;
1168        if let Ok(mut guard) = self.base.initialized.try_write() {
1169            *guard = true;
1170            return true;
1171        }
1172        false
1173    }
1174
1175    pub fn get_stats(&self) -> DMSCProtocolStats {
1176        self.base.stats.try_read()
1177            .map(|guard| guard.clone())
1178            .unwrap_or_else(|_| DMSCProtocolStats::new())
1179    }
1180
1181    pub fn get_health(&self) -> DMSCProtocolHealth {
1182        if let Ok(guard) = self.base.initialized.try_read() {
1183            if *guard {
1184                return DMSCProtocolHealth::Healthy;
1185            }
1186        }
1187        DMSCProtocolHealth::Unknown
1188    }
1189
1190    pub fn shutdown(&mut self) -> bool {
1191        if let Ok(mut guard) = self.base.initialized.try_write() {
1192            *guard = false;
1193            return true;
1194        }
1195        false
1196    }
1197}
dmsc/protocol/mod.rs

dmsc/protocol/
mod.rs
