openigtlink_rust/protocol/
message.rs

1//! OpenIGTLink message trait and structures
2//!
3//! This module defines the common interface that all message types must implement,
4//! as well as the generic message wrapper structure.
5
6use crate::error::Result;
7use crate::protocol::header::Header;
8use std::collections::HashMap;
9
10/// Common interface for all OpenIGTLink message types
11///
12/// Each message type (TRANSFORM, IMAGE, STATUS, etc.) must implement this trait
13/// to provide encoding/decoding functionality.
14pub trait Message: Sized {
15    /// Returns the message type name (e.g., "TRANSFORM", "IMAGE")
16    ///
17    /// This must match the OpenIGTLink protocol specification.
18    fn message_type() -> &'static str;
19
20    /// Encode message content to bytes
21    ///
22    /// # Returns
23    /// Byte vector containing the encoded message content (without header)
24    fn encode_content(&self) -> Result<Vec<u8>>;
25
26    /// Decode message content from bytes
27    ///
28    /// # Arguments
29    /// * `data` - Byte slice containing the message content (without header)
30    ///
31    /// # Returns
32    /// Decoded message or error
33    fn decode_content(data: &[u8]) -> Result<Self>;
34}
35
36/// Complete OpenIGTLink message structure
37///
38/// Wraps a specific message type with header, optional extended header,
39/// and optional metadata.
40///
41/// # Type Parameters
42/// * `T` - Message type that implements the `Message` trait
43#[derive(Debug)]
44pub struct IgtlMessage<T: Message> {
45    /// Message header (58 bytes)
46    pub header: Header,
47    /// Extended header (Version 3 feature, optional)
48    pub extended_header: Option<Vec<u8>>,
49    /// Message content
50    pub content: T,
51    /// Metadata as key-value pairs (Version 3 feature, optional)
52    pub metadata: Option<HashMap<String, String>>,
53}
54
55impl<T: Message> IgtlMessage<T> {
56    /// Create a new message with the given content and device name
57    ///
58    /// # Arguments
59    /// * `content` - Message content
60    /// * `device_name` - Device name (max 20 characters)
61    ///
62    /// # Returns
63    /// New message with generated header
64    pub fn new(content: T, device_name: &str) -> Result<Self> {
65        use crate::protocol::header::{TypeName, DeviceName, Timestamp};
66
67        let timestamp = Timestamp::now();
68
69        let content_bytes = content.encode_content()?;
70        let body_size = content_bytes.len() as u64;
71
72        let header = Header {
73            version: 2, // Version 2 compatible
74            type_name: TypeName::new(T::message_type())?,
75            device_name: DeviceName::new(device_name)?,
76            timestamp,
77            body_size,
78            crc: 0, // Will be calculated during encode
79        };
80
81        Ok(IgtlMessage {
82            header,
83            extended_header: None,
84            content,
85            metadata: None,
86        })
87    }
88
89    /// Set extended header data (Version 3 feature)
90    ///
91    /// When extended header is set, the message version is automatically upgraded to 3.
92    ///
93    /// # Arguments
94    /// * `data` - Extended header data as byte vector
95    ///
96    /// # Examples
97    /// ```no_run
98    /// # use openigtlink_rust::protocol::{IgtlMessage, types::TransformMessage};
99    /// let transform = TransformMessage::identity();
100    /// let mut msg = IgtlMessage::new(transform, "Device").unwrap();
101    /// msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
102    /// ```
103    pub fn set_extended_header(&mut self, data: Vec<u8>) {
104        self.extended_header = Some(data);
105        // Upgrade to version 3 when extended header is used
106        if self.header.version < 3 {
107            self.header.version = 3;
108        }
109    }
110
111    /// Get extended header data reference (Version 3 feature)
112    ///
113    /// # Returns
114    /// Optional reference to extended header bytes
115    pub fn get_extended_header(&self) -> Option<&[u8]> {
116        self.extended_header.as_deref()
117    }
118
119    /// Remove extended header and optionally downgrade to Version 2
120    pub fn clear_extended_header(&mut self) {
121        self.extended_header = None;
122        // Downgrade to version 2 if no version 3 features are used
123        if self.metadata.is_none() && self.header.version == 3 {
124            self.header.version = 2;
125        }
126    }
127
128    /// Set metadata key-value pairs (Version 3 feature)
129    ///
130    /// When metadata is set, the message version is automatically upgraded to 3.
131    ///
132    /// # Arguments
133    /// * `metadata` - HashMap of key-value pairs
134    ///
135    /// # Examples
136    /// ```no_run
137    /// # use openigtlink_rust::protocol::{IgtlMessage, types::TransformMessage};
138    /// # use std::collections::HashMap;
139    /// let transform = TransformMessage::identity();
140    /// let mut msg = IgtlMessage::new(transform, "Device").unwrap();
141    /// let mut metadata = HashMap::new();
142    /// metadata.insert("priority".to_string(), "high".to_string());
143    /// msg.set_metadata(metadata);
144    /// ```
145    pub fn set_metadata(&mut self, metadata: HashMap<String, String>) {
146        self.metadata = Some(metadata);
147        // Upgrade to version 3 when metadata is used
148        if self.header.version < 3 {
149            self.header.version = 3;
150        }
151    }
152
153    /// Add a single metadata key-value pair (Version 3 feature)
154    ///
155    /// # Arguments
156    /// * `key` - Metadata key
157    /// * `value` - Metadata value
158    pub fn add_metadata(&mut self, key: String, value: String) {
159        if self.metadata.is_none() {
160            self.metadata = Some(HashMap::new());
161            if self.header.version < 3 {
162                self.header.version = 3;
163            }
164        }
165        self.metadata.as_mut().unwrap().insert(key, value);
166    }
167
168    /// Get metadata reference (Version 3 feature)
169    ///
170    /// # Returns
171    /// Optional reference to metadata HashMap
172    pub fn get_metadata(&self) -> Option<&HashMap<String, String>> {
173        self.metadata.as_ref()
174    }
175
176    /// Remove metadata and optionally downgrade to Version 2
177    pub fn clear_metadata(&mut self) {
178        self.metadata = None;
179        // Downgrade to version 2 if no version 3 features are used
180        if self.extended_header.is_none() && self.header.version == 3 {
181            self.header.version = 2;
182        }
183    }
184
185    /// Encode the complete message to bytes
186    ///
187    /// Version 2 format: Header (58) + Content
188    /// Version 3 format: Header (58) + ExtHdrSize (2) + ExtHdr (var) + Content + Metadata (var)
189    ///
190    /// Metadata format (Version 3):
191    /// - MetadataSize (2 bytes, big-endian)
192    /// - For each pair:
193    ///   - KeySize (2 bytes)
194    ///   - Key (KeySize bytes, UTF-8)
195    ///   - ValueSize (2 bytes)
196    ///   - Value (ValueSize bytes, UTF-8)
197    ///
198    /// # Returns
199    /// Complete message as byte vector
200    pub fn encode(&self) -> Result<Vec<u8>> {
201        use crate::protocol::crc::calculate_crc;
202
203        // 1. Encode content
204        let content_bytes = self.content.encode_content()?;
205
206        // 2. Encode metadata if present
207        let metadata_bytes = if self.header.version >= 3 && self.metadata.is_some() {
208            let metadata = self.metadata.as_ref().unwrap();
209            let mut meta_buf = Vec::new();
210
211            // Metadata header size (2 bytes)
212            let meta_header_size = (metadata.len() as u16).to_be_bytes();
213            meta_buf.extend_from_slice(&meta_header_size);
214
215            // Each key-value pair
216            for (key, value) in metadata.iter() {
217                // Key size and data
218                let key_bytes = key.as_bytes();
219                meta_buf.extend_from_slice(&(key_bytes.len() as u16).to_be_bytes());
220                meta_buf.extend_from_slice(key_bytes);
221
222                // Value size and data
223                let value_bytes = value.as_bytes();
224                meta_buf.extend_from_slice(&(value_bytes.len() as u16).to_be_bytes());
225                meta_buf.extend_from_slice(value_bytes);
226            }
227
228            meta_buf
229        } else {
230            Vec::new()
231        };
232
233        // 3. Build body based on version
234        let body_bytes = if self.header.version >= 3 && self.extended_header.is_some() {
235            // Version 3 with extended header
236            let ext_header = self.extended_header.as_ref().unwrap();
237            let ext_header_size = ext_header.len() as u16;
238
239            let mut body = Vec::with_capacity(
240                2 + ext_header.len() + content_bytes.len() + metadata_bytes.len()
241            );
242            // Extended header size (2 bytes, big-endian)
243            body.extend_from_slice(&ext_header_size.to_be_bytes());
244            // Extended header data
245            body.extend_from_slice(ext_header);
246            // Content
247            body.extend_from_slice(&content_bytes);
248            // Metadata
249            body.extend_from_slice(&metadata_bytes);
250
251            body
252        } else if self.header.version >= 3 && !metadata_bytes.is_empty() {
253            // Version 3 without extended header but with metadata
254            let mut body = Vec::with_capacity(2 + content_bytes.len() + metadata_bytes.len());
255            // Extended header size = 0
256            body.extend_from_slice(&0u16.to_be_bytes());
257            // Content
258            body.extend_from_slice(&content_bytes);
259            // Metadata
260            body.extend_from_slice(&metadata_bytes);
261
262            body
263        } else {
264            // Version 2 or Version 3 without extended header/metadata
265            content_bytes
266        };
267
268        // 4. Update header with correct body_size and CRC
269        let mut header = self.header.clone();
270        header.body_size = body_bytes.len() as u64;
271        header.crc = calculate_crc(&body_bytes);
272
273        // 5. Combine header + body
274        let mut buf = Vec::with_capacity(Header::SIZE + body_bytes.len());
275        buf.extend_from_slice(&header.encode());
276        buf.extend_from_slice(&body_bytes);
277
278        Ok(buf)
279    }
280
281    /// Decode a complete message from bytes with CRC verification
282    ///
283    /// Automatically detects Version 2 or Version 3 format based on header.
284    ///
285    /// # Arguments
286    /// * `data` - Byte slice containing the complete message
287    ///
288    /// # Returns
289    /// Decoded message or error
290    pub fn decode(data: &[u8]) -> Result<Self> {
291        Self::decode_with_options(data, true)
292    }
293
294    /// Decode a complete message from bytes with optional CRC verification
295    ///
296    /// Allows skipping CRC verification for performance in trusted environments.
297    ///
298    /// # Arguments
299    /// * `data` - Byte slice containing the complete message
300    /// * `verify_crc` - Whether to verify CRC (true = verify, false = skip)
301    ///
302    /// # Returns
303    /// Decoded message or error
304    ///
305    /// # Safety
306    /// Disabling CRC verification (`verify_crc = false`) should only be done in
307    /// trusted environments where data corruption is unlikely (e.g., loopback, local network).
308    /// Using this in production over unreliable networks may lead to silent data corruption.
309    ///
310    /// # Examples
311    /// ```no_run
312    /// # use openigtlink_rust::protocol::{IgtlMessage, types::TransformMessage};
313    /// # let data = vec![0u8; 106];
314    /// // Decode with CRC verification (recommended)
315    /// let msg = IgtlMessage::<TransformMessage>::decode_with_options(&data, true)?;
316    ///
317    /// // Decode without CRC verification (use with caution)
318    /// let msg_fast = IgtlMessage::<TransformMessage>::decode_with_options(&data, false)?;
319    /// # Ok::<(), openigtlink_rust::error::IgtlError>(())
320    /// ```
321    pub fn decode_with_options(data: &[u8], verify_crc: bool) -> Result<Self> {
322        use crate::protocol::crc::calculate_crc;
323        use crate::error::IgtlError;
324
325        if data.len() < Header::SIZE {
326            return Err(IgtlError::InvalidSize {
327                expected: Header::SIZE,
328                actual: data.len(),
329            });
330        }
331
332        // 1. Parse header
333        let header = Header::decode(&data[..Header::SIZE])?;
334
335        // 2. Extract body
336        let body_start = Header::SIZE;
337        let body_end = body_start + header.body_size as usize;
338
339        if data.len() < body_end {
340            return Err(IgtlError::InvalidSize {
341                expected: body_end,
342                actual: data.len(),
343            });
344        }
345
346        let body_bytes = &data[body_start..body_end];
347
348        // 3. Verify CRC (if requested)
349        if verify_crc {
350            let calculated_crc = calculate_crc(body_bytes);
351            if calculated_crc != header.crc {
352                return Err(IgtlError::CrcMismatch {
353                    expected: header.crc,
354                    actual: calculated_crc,
355                });
356            }
357        }
358
359        // 4. Parse body based on version
360        let (extended_header, remaining_bytes, has_ext_header_field) = if header.version >= 3 && body_bytes.len() >= 2 {
361            // Try to parse as Version 3 with extended header
362            let ext_header_size = u16::from_be_bytes([body_bytes[0], body_bytes[1]]) as usize;
363
364            if ext_header_size > 0 && body_bytes.len() >= 2 + ext_header_size {
365                // Version 3 with non-empty extended header
366                let ext_header_data = body_bytes[2..2 + ext_header_size].to_vec();
367                let content_start = 2 + ext_header_size;
368                (Some(ext_header_data), &body_bytes[content_start..], true)
369            } else if ext_header_size == 0 && body_bytes.len() >= 2 {
370                // Version 3 with empty extended header (size field = 0)
371                (Some(Vec::new()), &body_bytes[2..], true)
372            } else {
373                // Invalid extended header size
374                return Err(IgtlError::InvalidSize {
375                    expected: 2 + ext_header_size,
376                    actual: body_bytes.len(),
377                });
378            }
379        } else {
380            // Version 2 - entire body is content
381            (None, body_bytes, false)
382        };
383
384        // 5. Try to determine content size and parse metadata (Version 3 only)
385        let (content_bytes, metadata) = if header.version >= 3 && has_ext_header_field {
386            // For Version 3 with ext_header_size field (whether 0 or not)
387            // We need to figure out where content ends and metadata begins
388            // Strategy: Try to decode content. If successful, re-encode to find actual size.
389
390            // First, try decoding the entire remaining_bytes as content
391            match T::decode_content(remaining_bytes) {
392                Ok(content) => {
393                    // Content decoded successfully
394                    // Re-encode to find actual content size
395                    let actual_content_size = content.encode_content()?.len();
396
397                    if remaining_bytes.len() > actual_content_size {
398                        // There's metadata after the content
399                        let content_part = &remaining_bytes[..actual_content_size];
400                        let metadata_part = &remaining_bytes[actual_content_size..];
401
402                        // Parse metadata
403                        let parsed_metadata = Self::decode_metadata(metadata_part)?;
404                        (content_part, parsed_metadata)
405                    } else {
406                        // No metadata
407                        (remaining_bytes, None)
408                    }
409                }
410                Err(IgtlError::InvalidSize { expected, .. }) if remaining_bytes.len() > expected => {
411                    // Content decode failed due to size mismatch (fixed-size message with metadata)
412                    // This means we have: Content (expected bytes) + Metadata (rest)
413                    let content_part = &remaining_bytes[..expected];
414                    let metadata_part = &remaining_bytes[expected..];
415
416                    // Verify content decodes correctly
417                    if T::decode_content(content_part).is_ok() {
418                        // Parse metadata
419                        let parsed_metadata = Self::decode_metadata(metadata_part)?;
420                        (content_part, parsed_metadata)
421                    } else {
422                        // Content still doesn't decode - treat all as content
423                        (remaining_bytes, None)
424                    }
425                }
426                Err(_) => {
427                    // Some other error - treat all as content
428                    (remaining_bytes, None)
429                }
430            }
431        } else {
432            // Version 2 - no metadata
433            (remaining_bytes, None)
434        };
435
436        // 6. Decode content
437        let content = T::decode_content(content_bytes)?;
438
439        Ok(IgtlMessage {
440            header,
441            extended_header,
442            content,
443            metadata,
444        })
445    }
446
447    /// Decode metadata from bytes (helper function)
448    fn decode_metadata(data: &[u8]) -> Result<Option<HashMap<String, String>>> {
449        use crate::error::IgtlError;
450
451        if data.len() < 2 {
452            return Ok(None);
453        }
454
455        let metadata_count = u16::from_be_bytes([data[0], data[1]]) as usize;
456
457        if metadata_count == 0 {
458            return Ok(None);
459        }
460
461        let mut metadata = HashMap::new();
462        let mut offset = 2;
463
464        for _ in 0..metadata_count {
465            // Read key size
466            if offset + 2 > data.len() {
467                return Err(IgtlError::InvalidSize {
468                    expected: offset + 2,
469                    actual: data.len(),
470                });
471            }
472            let key_size = u16::from_be_bytes([data[offset], data[offset + 1]]) as usize;
473            offset += 2;
474
475            // Read key
476            if offset + key_size > data.len() {
477                return Err(IgtlError::InvalidSize {
478                    expected: offset + key_size,
479                    actual: data.len(),
480                });
481            }
482            let key = String::from_utf8(data[offset..offset + key_size].to_vec())?;
483            offset += key_size;
484
485            // Read value size
486            if offset + 2 > data.len() {
487                return Err(IgtlError::InvalidSize {
488                    expected: offset + 2,
489                    actual: data.len(),
490                });
491            }
492            let value_size = u16::from_be_bytes([data[offset], data[offset + 1]]) as usize;
493            offset += 2;
494
495            // Read value
496            if offset + value_size > data.len() {
497                return Err(IgtlError::InvalidSize {
498                    expected: offset + value_size,
499                    actual: data.len(),
500                });
501            }
502            let value = String::from_utf8(data[offset..offset + value_size].to_vec())?;
503            offset += value_size;
504
505            metadata.insert(key, value);
506        }
507
508        Ok(Some(metadata))
509    }
510}
511
512#[cfg(test)]
513mod tests {
514    use super::*;
515    use crate::protocol::types::{TransformMessage, StatusMessage, CapabilityMessage};
516
517    // Mock message type for testing
518    struct TestMessage {
519        data: Vec<u8>,
520    }
521
522    impl Message for TestMessage {
523        fn message_type() -> &'static str {
524            "TEST"
525        }
526
527        fn encode_content(&self) -> Result<Vec<u8>> {
528            Ok(self.data.clone())
529        }
530
531        fn decode_content(data: &[u8]) -> Result<Self> {
532            Ok(TestMessage {
533                data: data.to_vec(),
534            })
535        }
536    }
537
538    #[test]
539    fn test_message_trait() {
540        assert_eq!(TestMessage::message_type(), "TEST");
541    }
542
543    #[test]
544    fn test_message_encode_decode() {
545        let original = TestMessage {
546            data: vec![1, 2, 3, 4, 5],
547        };
548
549        let encoded = original.encode_content().unwrap();
550        let decoded = TestMessage::decode_content(&encoded).unwrap();
551
552        assert_eq!(original.data, decoded.data);
553    }
554
555    #[test]
556    fn test_full_message_roundtrip_transform() {
557        let transform = TransformMessage::identity();
558        let msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
559
560        let encoded = msg.encode().unwrap();
561        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
562
563        // Verify header fields
564        assert_eq!(decoded.header.version, 2);
565        assert_eq!(decoded.header.type_name.as_str().unwrap(), "TRANSFORM");
566        assert_eq!(decoded.header.device_name.as_str().unwrap(), "TestDevice");
567        assert_eq!(decoded.header.body_size, 48);
568
569        // Verify content
570        assert_eq!(decoded.content, transform);
571    }
572
573    #[test]
574    fn test_full_message_roundtrip_status() {
575        let status = StatusMessage::ok("Operation successful");
576        let msg = IgtlMessage::new(status.clone(), "StatusDevice").unwrap();
577
578        let encoded = msg.encode().unwrap();
579        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
580
581        assert_eq!(decoded.header.type_name.as_str().unwrap(), "STATUS");
582        assert_eq!(decoded.content, status);
583    }
584
585    #[test]
586    fn test_full_message_roundtrip_capability() {
587        let capability = CapabilityMessage::new(vec![
588            "TRANSFORM".to_string(),
589            "STATUS".to_string(),
590            "IMAGE".to_string(),
591        ]);
592        let msg = IgtlMessage::new(capability.clone(), "CapDevice").unwrap();
593
594        let encoded = msg.encode().unwrap();
595        let decoded = IgtlMessage::<CapabilityMessage>::decode(&encoded).unwrap();
596
597        assert_eq!(decoded.header.type_name.as_str().unwrap(), "CAPABILITY");
598        assert_eq!(decoded.content, capability);
599    }
600
601    #[test]
602    fn test_timestamp_reasonable() {
603        let transform = TransformMessage::identity();
604        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
605
606        // Timestamp should be recent (within last year and not in future)
607        let now = std::time::SystemTime::now()
608            .duration_since(std::time::UNIX_EPOCH)
609            .unwrap()
610            .as_secs() as u32;
611
612        let one_year_ago = now - (365 * 24 * 60 * 60);
613
614        assert!(msg.header.timestamp.seconds >= one_year_ago);
615        assert!(msg.header.timestamp.seconds <= now + 1); // +1 for clock skew
616    }
617
618    #[test]
619    fn test_crc_verification() {
620        let transform = TransformMessage::identity();
621        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
622
623        let mut encoded = msg.encode().unwrap();
624
625        // Corrupt the content
626        let content_start = Header::SIZE;
627        encoded[content_start] ^= 0xFF;
628
629        // Should fail CRC check
630        let result = IgtlMessage::<TransformMessage>::decode(&encoded);
631        assert!(matches!(result, Err(crate::error::IgtlError::CrcMismatch { .. })));
632    }
633
634    #[test]
635    fn test_message_size_calculation() {
636        let transform = TransformMessage::identity();
637        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
638
639        let encoded = msg.encode().unwrap();
640
641        // Total size should be: Header (58) + TRANSFORM content (48) = 106
642        assert_eq!(encoded.len(), 106);
643    }
644
645    #[test]
646    fn test_decode_short_buffer() {
647        let short_data = vec![0u8; 30];
648        let result = IgtlMessage::<TransformMessage>::decode(&short_data);
649        assert!(matches!(result, Err(crate::error::IgtlError::InvalidSize { .. })));
650    }
651
652    // Version 3 Extended Header Tests
653
654    #[test]
655    fn test_extended_header_set_get() {
656        let transform = TransformMessage::identity();
657        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
658
659        // Initially no extended header
660        assert_eq!(msg.get_extended_header(), None);
661        assert_eq!(msg.header.version, 2);
662
663        // Set extended header
664        let ext_header = vec![0x01, 0x02, 0x03, 0x04];
665        msg.set_extended_header(ext_header.clone());
666
667        // Should upgrade to version 3
668        assert_eq!(msg.header.version, 3);
669        assert_eq!(msg.get_extended_header(), Some(ext_header.as_slice()));
670    }
671
672    #[test]
673    fn test_extended_header_clear() {
674        let transform = TransformMessage::identity();
675        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
676
677        // Set extended header
678        msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
679        assert_eq!(msg.header.version, 3);
680
681        // Clear extended header
682        msg.clear_extended_header();
683        assert_eq!(msg.get_extended_header(), None);
684        // Should downgrade to version 2 (no metadata either)
685        assert_eq!(msg.header.version, 2);
686    }
687
688    #[test]
689    fn test_version3_encode_decode_with_extended_header() {
690        let transform = TransformMessage::identity();
691        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
692
693        // Add extended header
694        let ext_header = vec![0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF];
695        msg.set_extended_header(ext_header.clone());
696
697        // Encode
698        let encoded = msg.encode().unwrap();
699
700        // Decode
701        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
702
703        // Verify version
704        assert_eq!(decoded.header.version, 3);
705
706        // Verify extended header
707        assert_eq!(decoded.get_extended_header(), Some(ext_header.as_slice()));
708
709        // Verify content
710        assert_eq!(decoded.content, transform);
711    }
712
713    #[test]
714    fn test_version3_encode_decode_without_extended_header() {
715        let status = StatusMessage::ok("Test message");
716        let msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
717
718        // Encode as Version 2
719        let encoded = msg.encode().unwrap();
720
721        // Decode
722        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
723
724        // Should remain version 2
725        assert_eq!(decoded.header.version, 2);
726        assert_eq!(decoded.get_extended_header(), None);
727        assert_eq!(decoded.content, status);
728    }
729
730    #[test]
731    fn test_extended_header_body_size_calculation() {
732        let transform = TransformMessage::identity();
733        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
734
735        // Add extended header
736        let ext_header = vec![0x01, 0x02, 0x03, 0x04]; // 4 bytes
737        msg.set_extended_header(ext_header);
738
739        let encoded = msg.encode().unwrap();
740
741        // Total size should be:
742        // Header (58) + ExtHdrSize (2) + ExtHdr (4) + TRANSFORM content (48) = 112
743        assert_eq!(encoded.len(), 112);
744
745        // Verify body_size in header includes extended header overhead
746        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
747        assert_eq!(decoded.header.body_size, 2 + 4 + 48); // ExtHdrSize + ExtHdr + Content
748    }
749
750    #[test]
751    fn test_extended_header_empty() {
752        let transform = TransformMessage::identity();
753        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
754
755        // Set empty extended header
756        msg.set_extended_header(vec![]);
757
758        let encoded = msg.encode().unwrap();
759        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
760
761        // Should still be version 3
762        assert_eq!(decoded.header.version, 3);
763        // Extended header should be empty
764        assert_eq!(decoded.get_extended_header(), Some(&[] as &[u8]));
765        // Content should be intact
766        assert_eq!(decoded.content, transform);
767    }
768
769    #[test]
770    fn test_extended_header_large() {
771        let status = StatusMessage::ok("Test");
772        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
773
774        // Create large extended header (1 KB)
775        let ext_header = vec![0xAB; 1024];
776        msg.set_extended_header(ext_header.clone());
777
778        let encoded = msg.encode().unwrap();
779        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
780
781        assert_eq!(decoded.header.version, 3);
782        assert_eq!(decoded.get_extended_header(), Some(ext_header.as_slice()));
783        assert_eq!(decoded.content, status);
784    }
785
786    #[test]
787    fn test_version3_crc_includes_extended_header() {
788        let transform = TransformMessage::identity();
789        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
790
791        msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
792
793        let mut encoded = msg.encode().unwrap();
794
795        // Corrupt extended header
796        encoded[Header::SIZE + 2] ^= 0xFF; // First byte of extended header
797
798        // Should fail CRC check
799        let result = IgtlMessage::<TransformMessage>::decode(&encoded);
800        assert!(matches!(result, Err(crate::error::IgtlError::CrcMismatch { .. })));
801    }
802
803    #[test]
804    fn test_backward_compatibility_version2() {
805        // Create a Version 2 message
806        let capability = CapabilityMessage::new(vec!["TRANSFORM".to_string()]);
807        let msg = IgtlMessage::new(capability.clone(), "Device").unwrap();
808
809        assert_eq!(msg.header.version, 2);
810
811        let encoded = msg.encode().unwrap();
812        let decoded = IgtlMessage::<CapabilityMessage>::decode(&encoded).unwrap();
813
814        // Should decode correctly as Version 2
815        assert_eq!(decoded.header.version, 2);
816        assert_eq!(decoded.get_extended_header(), None);
817        assert_eq!(decoded.content, capability);
818    }
819
820    // Version 3 Metadata Tests
821
822    #[test]
823    fn test_metadata_set_get() {
824        let transform = TransformMessage::identity();
825        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
826
827        // Initially no metadata
828        assert_eq!(msg.get_metadata(), None);
829        assert_eq!(msg.header.version, 2);
830
831        // Set metadata
832        let mut metadata = HashMap::new();
833        metadata.insert("priority".to_string(), "high".to_string());
834        metadata.insert("sequence".to_string(), "42".to_string());
835        msg.set_metadata(metadata.clone());
836
837        // Should upgrade to version 3
838        assert_eq!(msg.header.version, 3);
839        assert_eq!(msg.get_metadata(), Some(&metadata));
840    }
841
842    #[test]
843    fn test_metadata_add() {
844        let status = StatusMessage::ok("Test");
845        let mut msg = IgtlMessage::new(status, "TestDevice").unwrap();
846
847        assert_eq!(msg.header.version, 2);
848
849        // Add metadata one by one
850        msg.add_metadata("key1".to_string(), "value1".to_string());
851        assert_eq!(msg.header.version, 3);
852
853        msg.add_metadata("key2".to_string(), "value2".to_string());
854
855        let metadata = msg.get_metadata().unwrap();
856        assert_eq!(metadata.get("key1"), Some(&"value1".to_string()));
857        assert_eq!(metadata.get("key2"), Some(&"value2".to_string()));
858    }
859
860    #[test]
861    fn test_metadata_clear() {
862        let transform = TransformMessage::identity();
863        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
864
865        // Set metadata
866        msg.add_metadata("test".to_string(), "value".to_string());
867        assert_eq!(msg.header.version, 3);
868
869        // Clear metadata
870        msg.clear_metadata();
871        assert_eq!(msg.get_metadata(), None);
872        // Should downgrade to version 2 (no extended header either)
873        assert_eq!(msg.header.version, 2);
874    }
875
876    #[test]
877    fn test_version3_encode_decode_with_metadata() {
878        let transform = TransformMessage::identity();
879        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
880
881        // Add metadata
882        let mut metadata = HashMap::new();
883        metadata.insert("priority".to_string(), "high".to_string());
884        metadata.insert("timestamp".to_string(), "123456".to_string());
885        msg.set_metadata(metadata.clone());
886
887        // Encode
888        let encoded = msg.encode().unwrap();
889
890        // Decode
891        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
892
893        // Verify version
894        assert_eq!(decoded.header.version, 3);
895
896        // Verify metadata
897        let decoded_metadata = decoded.get_metadata().unwrap();
898        assert_eq!(decoded_metadata.get("priority"), Some(&"high".to_string()));
899        assert_eq!(decoded_metadata.get("timestamp"), Some(&"123456".to_string()));
900
901        // Verify content
902        assert_eq!(decoded.content, transform);
903    }
904
905    #[test]
906    fn test_version3_with_extended_header_and_metadata() {
907        let status = StatusMessage::ok("Test message");
908        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
909
910        // Add both extended header and metadata
911        msg.set_extended_header(vec![0xAA, 0xBB, 0xCC, 0xDD]);
912        msg.add_metadata("key1".to_string(), "value1".to_string());
913        msg.add_metadata("key2".to_string(), "value2".to_string());
914
915        let encoded = msg.encode().unwrap();
916        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
917
918        // Verify version
919        assert_eq!(decoded.header.version, 3);
920
921        // Verify extended header
922        let expected_ext_header: &[u8] = &[0xAA, 0xBB, 0xCC, 0xDD];
923        assert_eq!(decoded.get_extended_header(), Some(expected_ext_header));
924
925        // Verify metadata
926        let metadata = decoded.get_metadata().unwrap();
927        assert_eq!(metadata.get("key1"), Some(&"value1".to_string()));
928        assert_eq!(metadata.get("key2"), Some(&"value2".to_string()));
929
930        // Verify content
931        assert_eq!(decoded.content, status);
932    }
933
934    #[test]
935    fn test_metadata_empty() {
936        let transform = TransformMessage::identity();
937        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
938
939        // Set empty metadata
940        msg.set_metadata(HashMap::new());
941
942        let encoded = msg.encode().unwrap();
943        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
944
945        // Should not have metadata (empty HashMap)
946        assert_eq!(decoded.get_metadata(), None);
947        assert_eq!(decoded.content, transform);
948    }
949
950    #[test]
951    fn test_metadata_utf8_values() {
952        let status = StatusMessage::ok("Test");
953        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
954
955        // Add UTF-8 metadata
956        msg.add_metadata("name".to_string(), "日本語".to_string());
957        msg.add_metadata("emoji".to_string(), "🎉✨".to_string());
958
959        let encoded = msg.encode().unwrap();
960        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
961
962        let metadata = decoded.get_metadata().unwrap();
963        assert_eq!(metadata.get("name"), Some(&"日本語".to_string()));
964        assert_eq!(metadata.get("emoji"), Some(&"🎉✨".to_string()));
965    }
966
967    // CRC Verification Tests
968
969    #[test]
970    fn test_decode_with_crc_verification_enabled() {
971        let transform = TransformMessage::identity();
972        let msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
973
974        let encoded = msg.encode().unwrap();
975
976        // Should decode successfully with CRC verification
977        let decoded = IgtlMessage::<TransformMessage>::decode_with_options(&encoded, true).unwrap();
978        assert_eq!(decoded.content, transform);
979    }
980
981    #[test]
982    fn test_decode_with_crc_verification_disabled() {
983        let transform = TransformMessage::identity();
984        let msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
985
986        let mut encoded = msg.encode().unwrap();
987
988        // Corrupt the data
989        encoded[Header::SIZE] ^= 0xFF;
990
991        // Should fail with CRC verification enabled
992        let result_with_crc = IgtlMessage::<TransformMessage>::decode_with_options(&encoded, true);
993        assert!(matches!(result_with_crc, Err(crate::error::IgtlError::CrcMismatch { .. })));
994
995        // Should succeed with CRC verification disabled (even with corrupted data)
996        let result_without_crc = IgtlMessage::<TransformMessage>::decode_with_options(&encoded, false);
997        assert!(result_without_crc.is_ok());
998    }
999
1000    #[test]
1001    fn test_decode_default_uses_crc_verification() {
1002        let transform = TransformMessage::identity();
1003        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
1004
1005        let mut encoded = msg.encode().unwrap();
1006
1007        // Corrupt the data
1008        encoded[Header::SIZE] ^= 0xFF;
1009
1010        // Default decode() should verify CRC and fail
1011        let result = IgtlMessage::<TransformMessage>::decode(&encoded);
1012        assert!(matches!(result, Err(crate::error::IgtlError::CrcMismatch { .. })));
1013    }
1014
1015    #[test]
1016    fn test_crc_skip_performance() {
1017        // This test demonstrates that skipping CRC works correctly
1018        let status = StatusMessage::ok("Performance test");
1019        let msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
1020
1021        let encoded = msg.encode().unwrap();
1022
1023        // Both should decode to the same content (when data is not corrupted)
1024        let decoded_with_crc = IgtlMessage::<StatusMessage>::decode_with_options(&encoded, true).unwrap();
1025        let decoded_without_crc = IgtlMessage::<StatusMessage>::decode_with_options(&encoded, false).unwrap();
1026
1027        assert_eq!(decoded_with_crc.content, decoded_without_crc.content);
1028        assert_eq!(decoded_with_crc.content, status);
1029    }
1030
1031    #[test]
1032    fn test_version3_crc_skip_with_extended_header() {
1033        let transform = TransformMessage::identity();
1034        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
1035
1036        msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
1037
1038        let encoded = msg.encode().unwrap();
1039
1040        // Should work with CRC disabled
1041        let decoded = IgtlMessage::<TransformMessage>::decode_with_options(&encoded, false).unwrap();
1042        assert_eq!(decoded.header.version, 3);
1043        let expected: &[u8] = &[0x01, 0x02, 0x03, 0x04];
1044        assert_eq!(decoded.get_extended_header(), Some(expected));
1045        assert_eq!(decoded.content, transform);
1046    }
1047
1048    #[test]
1049    fn test_version3_crc_skip_with_metadata() {
1050        let status = StatusMessage::ok("Test");
1051        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
1052
1053        msg.add_metadata("key1".to_string(), "value1".to_string());
1054        msg.add_metadata("key2".to_string(), "value2".to_string());
1055
1056        let encoded = msg.encode().unwrap();
1057
1058        // Should work with CRC disabled
1059        let decoded = IgtlMessage::<StatusMessage>::decode_with_options(&encoded, false).unwrap();
1060        assert_eq!(decoded.header.version, 3);
1061
1062        let metadata = decoded.get_metadata().unwrap();
1063        assert_eq!(metadata.get("key1"), Some(&"value1".to_string()));
1064        assert_eq!(metadata.get("key2"), Some(&"value2".to_string()));
1065        assert_eq!(decoded.content, status);
1066    }
1067}
openigtlink_rust/protocol/message.rs

openigtlink_rust/protocol/
message.rs
