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::{DeviceName, Timestamp, TypeName};
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::error::IgtlError;
323        use crate::protocol::crc::calculate_crc;
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) =
361            if header.version >= 3 && body_bytes.len() >= 2 {
362                // Try to parse as Version 3 with extended header
363                let ext_header_size = u16::from_be_bytes([body_bytes[0], body_bytes[1]]) as usize;
364
365                if ext_header_size > 0 && body_bytes.len() >= 2 + ext_header_size {
366                    // Version 3 with non-empty extended header
367                    let ext_header_data = body_bytes[2..2 + ext_header_size].to_vec();
368                    let content_start = 2 + ext_header_size;
369                    (Some(ext_header_data), &body_bytes[content_start..], true)
370                } else if ext_header_size == 0 && body_bytes.len() >= 2 {
371                    // Version 3 with empty extended header (size field = 0)
372                    (Some(Vec::new()), &body_bytes[2..], true)
373                } else {
374                    // Invalid extended header size
375                    return Err(IgtlError::InvalidSize {
376                        expected: 2 + ext_header_size,
377                        actual: body_bytes.len(),
378                    });
379                }
380            } else {
381                // Version 2 - entire body is content
382                (None, body_bytes, false)
383            };
384
385        // 5. Try to determine content size and parse metadata (Version 3 only)
386        let (content_bytes, metadata) = if header.version >= 3 && has_ext_header_field {
387            // For Version 3 with ext_header_size field (whether 0 or not)
388            // We need to figure out where content ends and metadata begins
389            // Strategy: Try to decode content. If successful, re-encode to find actual size.
390
391            // First, try decoding the entire remaining_bytes as content
392            match T::decode_content(remaining_bytes) {
393                Ok(content) => {
394                    // Content decoded successfully
395                    // Re-encode to find actual content size
396                    let actual_content_size = content.encode_content()?.len();
397
398                    if remaining_bytes.len() > actual_content_size {
399                        // There's metadata after the content
400                        let content_part = &remaining_bytes[..actual_content_size];
401                        let metadata_part = &remaining_bytes[actual_content_size..];
402
403                        // Parse metadata
404                        let parsed_metadata = Self::decode_metadata(metadata_part)?;
405                        (content_part, parsed_metadata)
406                    } else {
407                        // No metadata
408                        (remaining_bytes, None)
409                    }
410                }
411                Err(IgtlError::InvalidSize { expected, .. })
412                    if remaining_bytes.len() > expected =>
413                {
414                    // Content decode failed due to size mismatch (fixed-size message with metadata)
415                    // This means we have: Content (expected bytes) + Metadata (rest)
416                    let content_part = &remaining_bytes[..expected];
417                    let metadata_part = &remaining_bytes[expected..];
418
419                    // Verify content decodes correctly
420                    if T::decode_content(content_part).is_ok() {
421                        // Parse metadata
422                        let parsed_metadata = Self::decode_metadata(metadata_part)?;
423                        (content_part, parsed_metadata)
424                    } else {
425                        // Content still doesn't decode - treat all as content
426                        (remaining_bytes, None)
427                    }
428                }
429                Err(_) => {
430                    // Some other error - treat all as content
431                    (remaining_bytes, None)
432                }
433            }
434        } else {
435            // Version 2 - no metadata
436            (remaining_bytes, None)
437        };
438
439        // 6. Decode content
440        let content = T::decode_content(content_bytes)?;
441
442        Ok(IgtlMessage {
443            header,
444            extended_header,
445            content,
446            metadata,
447        })
448    }
449
450    /// Decode metadata from bytes (helper function)
451    fn decode_metadata(data: &[u8]) -> Result<Option<HashMap<String, String>>> {
452        use crate::error::IgtlError;
453
454        if data.len() < 2 {
455            return Ok(None);
456        }
457
458        let metadata_count = u16::from_be_bytes([data[0], data[1]]) as usize;
459
460        if metadata_count == 0 {
461            return Ok(None);
462        }
463
464        let mut metadata = HashMap::new();
465        let mut offset = 2;
466
467        for _ in 0..metadata_count {
468            // Read key size
469            if offset + 2 > data.len() {
470                return Err(IgtlError::InvalidSize {
471                    expected: offset + 2,
472                    actual: data.len(),
473                });
474            }
475            let key_size = u16::from_be_bytes([data[offset], data[offset + 1]]) as usize;
476            offset += 2;
477
478            // Read key
479            if offset + key_size > data.len() {
480                return Err(IgtlError::InvalidSize {
481                    expected: offset + key_size,
482                    actual: data.len(),
483                });
484            }
485            let key = String::from_utf8(data[offset..offset + key_size].to_vec())?;
486            offset += key_size;
487
488            // Read value size
489            if offset + 2 > data.len() {
490                return Err(IgtlError::InvalidSize {
491                    expected: offset + 2,
492                    actual: data.len(),
493                });
494            }
495            let value_size = u16::from_be_bytes([data[offset], data[offset + 1]]) as usize;
496            offset += 2;
497
498            // Read value
499            if offset + value_size > data.len() {
500                return Err(IgtlError::InvalidSize {
501                    expected: offset + value_size,
502                    actual: data.len(),
503                });
504            }
505            let value = String::from_utf8(data[offset..offset + value_size].to_vec())?;
506            offset += value_size;
507
508            metadata.insert(key, value);
509        }
510
511        Ok(Some(metadata))
512    }
513}
514
515#[cfg(test)]
516mod tests {
517    use super::*;
518    use crate::protocol::types::{CapabilityMessage, StatusMessage, TransformMessage};
519
520    // Mock message type for testing
521    struct TestMessage {
522        data: Vec<u8>,
523    }
524
525    impl Message for TestMessage {
526        fn message_type() -> &'static str {
527            "TEST"
528        }
529
530        fn encode_content(&self) -> Result<Vec<u8>> {
531            Ok(self.data.clone())
532        }
533
534        fn decode_content(data: &[u8]) -> Result<Self> {
535            Ok(TestMessage {
536                data: data.to_vec(),
537            })
538        }
539    }
540
541    #[test]
542    fn test_message_trait() {
543        assert_eq!(TestMessage::message_type(), "TEST");
544    }
545
546    #[test]
547    fn test_message_encode_decode() {
548        let original = TestMessage {
549            data: vec![1, 2, 3, 4, 5],
550        };
551
552        let encoded = original.encode_content().unwrap();
553        let decoded = TestMessage::decode_content(&encoded).unwrap();
554
555        assert_eq!(original.data, decoded.data);
556    }
557
558    #[test]
559    fn test_full_message_roundtrip_transform() {
560        let transform = TransformMessage::identity();
561        let msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
562
563        let encoded = msg.encode().unwrap();
564        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
565
566        // Verify header fields
567        assert_eq!(decoded.header.version, 2);
568        assert_eq!(decoded.header.type_name.as_str().unwrap(), "TRANSFORM");
569        assert_eq!(decoded.header.device_name.as_str().unwrap(), "TestDevice");
570        assert_eq!(decoded.header.body_size, 48);
571
572        // Verify content
573        assert_eq!(decoded.content, transform);
574    }
575
576    #[test]
577    fn test_full_message_roundtrip_status() {
578        let status = StatusMessage::ok("Operation successful");
579        let msg = IgtlMessage::new(status.clone(), "StatusDevice").unwrap();
580
581        let encoded = msg.encode().unwrap();
582        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
583
584        assert_eq!(decoded.header.type_name.as_str().unwrap(), "STATUS");
585        assert_eq!(decoded.content, status);
586    }
587
588    #[test]
589    fn test_full_message_roundtrip_capability() {
590        let capability = CapabilityMessage::new(vec![
591            "TRANSFORM".to_string(),
592            "STATUS".to_string(),
593            "IMAGE".to_string(),
594        ]);
595        let msg = IgtlMessage::new(capability.clone(), "CapDevice").unwrap();
596
597        let encoded = msg.encode().unwrap();
598        let decoded = IgtlMessage::<CapabilityMessage>::decode(&encoded).unwrap();
599
600        assert_eq!(decoded.header.type_name.as_str().unwrap(), "CAPABILITY");
601        assert_eq!(decoded.content, capability);
602    }
603
604    #[test]
605    fn test_timestamp_reasonable() {
606        let transform = TransformMessage::identity();
607        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
608
609        // Timestamp should be recent (within last year and not in future)
610        let now = std::time::SystemTime::now()
611            .duration_since(std::time::UNIX_EPOCH)
612            .unwrap()
613            .as_secs() as u32;
614
615        let one_year_ago = now - (365 * 24 * 60 * 60);
616
617        assert!(msg.header.timestamp.seconds >= one_year_ago);
618        assert!(msg.header.timestamp.seconds <= now + 1); // +1 for clock skew
619    }
620
621    #[test]
622    fn test_crc_verification() {
623        let transform = TransformMessage::identity();
624        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
625
626        let mut encoded = msg.encode().unwrap();
627
628        // Corrupt the content
629        let content_start = Header::SIZE;
630        encoded[content_start] ^= 0xFF;
631
632        // Should fail CRC check
633        let result = IgtlMessage::<TransformMessage>::decode(&encoded);
634        assert!(matches!(
635            result,
636            Err(crate::error::IgtlError::CrcMismatch { .. })
637        ));
638    }
639
640    #[test]
641    fn test_message_size_calculation() {
642        let transform = TransformMessage::identity();
643        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
644
645        let encoded = msg.encode().unwrap();
646
647        // Total size should be: Header (58) + TRANSFORM content (48) = 106
648        assert_eq!(encoded.len(), 106);
649    }
650
651    #[test]
652    fn test_decode_short_buffer() {
653        let short_data = vec![0u8; 30];
654        let result = IgtlMessage::<TransformMessage>::decode(&short_data);
655        assert!(matches!(
656            result,
657            Err(crate::error::IgtlError::InvalidSize { .. })
658        ));
659    }
660
661    // Version 3 Extended Header Tests
662
663    #[test]
664    fn test_extended_header_set_get() {
665        let transform = TransformMessage::identity();
666        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
667
668        // Initially no extended header
669        assert_eq!(msg.get_extended_header(), None);
670        assert_eq!(msg.header.version, 2);
671
672        // Set extended header
673        let ext_header = vec![0x01, 0x02, 0x03, 0x04];
674        msg.set_extended_header(ext_header.clone());
675
676        // Should upgrade to version 3
677        assert_eq!(msg.header.version, 3);
678        assert_eq!(msg.get_extended_header(), Some(ext_header.as_slice()));
679    }
680
681    #[test]
682    fn test_extended_header_clear() {
683        let transform = TransformMessage::identity();
684        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
685
686        // Set extended header
687        msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
688        assert_eq!(msg.header.version, 3);
689
690        // Clear extended header
691        msg.clear_extended_header();
692        assert_eq!(msg.get_extended_header(), None);
693        // Should downgrade to version 2 (no metadata either)
694        assert_eq!(msg.header.version, 2);
695    }
696
697    #[test]
698    fn test_version3_encode_decode_with_extended_header() {
699        let transform = TransformMessage::identity();
700        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
701
702        // Add extended header
703        let ext_header = vec![0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF];
704        msg.set_extended_header(ext_header.clone());
705
706        // Encode
707        let encoded = msg.encode().unwrap();
708
709        // Decode
710        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
711
712        // Verify version
713        assert_eq!(decoded.header.version, 3);
714
715        // Verify extended header
716        assert_eq!(decoded.get_extended_header(), Some(ext_header.as_slice()));
717
718        // Verify content
719        assert_eq!(decoded.content, transform);
720    }
721
722    #[test]
723    fn test_version3_encode_decode_without_extended_header() {
724        let status = StatusMessage::ok("Test message");
725        let msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
726
727        // Encode as Version 2
728        let encoded = msg.encode().unwrap();
729
730        // Decode
731        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
732
733        // Should remain version 2
734        assert_eq!(decoded.header.version, 2);
735        assert_eq!(decoded.get_extended_header(), None);
736        assert_eq!(decoded.content, status);
737    }
738
739    #[test]
740    fn test_extended_header_body_size_calculation() {
741        let transform = TransformMessage::identity();
742        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
743
744        // Add extended header
745        let ext_header = vec![0x01, 0x02, 0x03, 0x04]; // 4 bytes
746        msg.set_extended_header(ext_header);
747
748        let encoded = msg.encode().unwrap();
749
750        // Total size should be:
751        // Header (58) + ExtHdrSize (2) + ExtHdr (4) + TRANSFORM content (48) = 112
752        assert_eq!(encoded.len(), 112);
753
754        // Verify body_size in header includes extended header overhead
755        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
756        assert_eq!(decoded.header.body_size, 2 + 4 + 48); // ExtHdrSize + ExtHdr + Content
757    }
758
759    #[test]
760    fn test_extended_header_empty() {
761        let transform = TransformMessage::identity();
762        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
763
764        // Set empty extended header
765        msg.set_extended_header(vec![]);
766
767        let encoded = msg.encode().unwrap();
768        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
769
770        // Should still be version 3
771        assert_eq!(decoded.header.version, 3);
772        // Extended header should be empty
773        assert_eq!(decoded.get_extended_header(), Some(&[] as &[u8]));
774        // Content should be intact
775        assert_eq!(decoded.content, transform);
776    }
777
778    #[test]
779    fn test_extended_header_large() {
780        let status = StatusMessage::ok("Test");
781        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
782
783        // Create large extended header (1 KB)
784        let ext_header = vec![0xAB; 1024];
785        msg.set_extended_header(ext_header.clone());
786
787        let encoded = msg.encode().unwrap();
788        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
789
790        assert_eq!(decoded.header.version, 3);
791        assert_eq!(decoded.get_extended_header(), Some(ext_header.as_slice()));
792        assert_eq!(decoded.content, status);
793    }
794
795    #[test]
796    fn test_version3_crc_includes_extended_header() {
797        let transform = TransformMessage::identity();
798        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
799
800        msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
801
802        let mut encoded = msg.encode().unwrap();
803
804        // Corrupt extended header
805        encoded[Header::SIZE + 2] ^= 0xFF; // First byte of extended header
806
807        // Should fail CRC check
808        let result = IgtlMessage::<TransformMessage>::decode(&encoded);
809        assert!(matches!(
810            result,
811            Err(crate::error::IgtlError::CrcMismatch { .. })
812        ));
813    }
814
815    #[test]
816    fn test_backward_compatibility_version2() {
817        // Create a Version 2 message
818        let capability = CapabilityMessage::new(vec!["TRANSFORM".to_string()]);
819        let msg = IgtlMessage::new(capability.clone(), "Device").unwrap();
820
821        assert_eq!(msg.header.version, 2);
822
823        let encoded = msg.encode().unwrap();
824        let decoded = IgtlMessage::<CapabilityMessage>::decode(&encoded).unwrap();
825
826        // Should decode correctly as Version 2
827        assert_eq!(decoded.header.version, 2);
828        assert_eq!(decoded.get_extended_header(), None);
829        assert_eq!(decoded.content, capability);
830    }
831
832    // Version 3 Metadata Tests
833
834    #[test]
835    fn test_metadata_set_get() {
836        let transform = TransformMessage::identity();
837        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
838
839        // Initially no metadata
840        assert_eq!(msg.get_metadata(), None);
841        assert_eq!(msg.header.version, 2);
842
843        // Set metadata
844        let mut metadata = HashMap::new();
845        metadata.insert("priority".to_string(), "high".to_string());
846        metadata.insert("sequence".to_string(), "42".to_string());
847        msg.set_metadata(metadata.clone());
848
849        // Should upgrade to version 3
850        assert_eq!(msg.header.version, 3);
851        assert_eq!(msg.get_metadata(), Some(&metadata));
852    }
853
854    #[test]
855    fn test_metadata_add() {
856        let status = StatusMessage::ok("Test");
857        let mut msg = IgtlMessage::new(status, "TestDevice").unwrap();
858
859        assert_eq!(msg.header.version, 2);
860
861        // Add metadata one by one
862        msg.add_metadata("key1".to_string(), "value1".to_string());
863        assert_eq!(msg.header.version, 3);
864
865        msg.add_metadata("key2".to_string(), "value2".to_string());
866
867        let metadata = msg.get_metadata().unwrap();
868        assert_eq!(metadata.get("key1"), Some(&"value1".to_string()));
869        assert_eq!(metadata.get("key2"), Some(&"value2".to_string()));
870    }
871
872    #[test]
873    fn test_metadata_clear() {
874        let transform = TransformMessage::identity();
875        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
876
877        // Set metadata
878        msg.add_metadata("test".to_string(), "value".to_string());
879        assert_eq!(msg.header.version, 3);
880
881        // Clear metadata
882        msg.clear_metadata();
883        assert_eq!(msg.get_metadata(), None);
884        // Should downgrade to version 2 (no extended header either)
885        assert_eq!(msg.header.version, 2);
886    }
887
888    #[test]
889    fn test_version3_encode_decode_with_metadata() {
890        let transform = TransformMessage::identity();
891        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
892
893        // Add metadata
894        let mut metadata = HashMap::new();
895        metadata.insert("priority".to_string(), "high".to_string());
896        metadata.insert("timestamp".to_string(), "123456".to_string());
897        msg.set_metadata(metadata.clone());
898
899        // Encode
900        let encoded = msg.encode().unwrap();
901
902        // Decode
903        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
904
905        // Verify version
906        assert_eq!(decoded.header.version, 3);
907
908        // Verify metadata
909        let decoded_metadata = decoded.get_metadata().unwrap();
910        assert_eq!(decoded_metadata.get("priority"), Some(&"high".to_string()));
911        assert_eq!(
912            decoded_metadata.get("timestamp"),
913            Some(&"123456".to_string())
914        );
915
916        // Verify content
917        assert_eq!(decoded.content, transform);
918    }
919
920    #[test]
921    fn test_version3_with_extended_header_and_metadata() {
922        let status = StatusMessage::ok("Test message");
923        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
924
925        // Add both extended header and metadata
926        msg.set_extended_header(vec![0xAA, 0xBB, 0xCC, 0xDD]);
927        msg.add_metadata("key1".to_string(), "value1".to_string());
928        msg.add_metadata("key2".to_string(), "value2".to_string());
929
930        let encoded = msg.encode().unwrap();
931        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
932
933        // Verify version
934        assert_eq!(decoded.header.version, 3);
935
936        // Verify extended header
937        let expected_ext_header: &[u8] = &[0xAA, 0xBB, 0xCC, 0xDD];
938        assert_eq!(decoded.get_extended_header(), Some(expected_ext_header));
939
940        // Verify metadata
941        let metadata = decoded.get_metadata().unwrap();
942        assert_eq!(metadata.get("key1"), Some(&"value1".to_string()));
943        assert_eq!(metadata.get("key2"), Some(&"value2".to_string()));
944
945        // Verify content
946        assert_eq!(decoded.content, status);
947    }
948
949    #[test]
950    fn test_metadata_empty() {
951        let transform = TransformMessage::identity();
952        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
953
954        // Set empty metadata
955        msg.set_metadata(HashMap::new());
956
957        let encoded = msg.encode().unwrap();
958        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
959
960        // Should not have metadata (empty HashMap)
961        assert_eq!(decoded.get_metadata(), None);
962        assert_eq!(decoded.content, transform);
963    }
964
965    #[test]
966    fn test_metadata_utf8_values() {
967        let status = StatusMessage::ok("Test");
968        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
969
970        // Add UTF-8 metadata
971        msg.add_metadata("name".to_string(), "日本語".to_string());
972        msg.add_metadata("emoji".to_string(), "🎉✨".to_string());
973
974        let encoded = msg.encode().unwrap();
975        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
976
977        let metadata = decoded.get_metadata().unwrap();
978        assert_eq!(metadata.get("name"), Some(&"日本語".to_string()));
979        assert_eq!(metadata.get("emoji"), Some(&"🎉✨".to_string()));
980    }
981
982    // CRC Verification Tests
983
984    #[test]
985    fn test_decode_with_crc_verification_enabled() {
986        let transform = TransformMessage::identity();
987        let msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
988
989        let encoded = msg.encode().unwrap();
990
991        // Should decode successfully with CRC verification
992        let decoded = IgtlMessage::<TransformMessage>::decode_with_options(&encoded, true).unwrap();
993        assert_eq!(decoded.content, transform);
994    }
995
996    #[test]
997    fn test_decode_with_crc_verification_disabled() {
998        let transform = TransformMessage::identity();
999        let msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
1000
1001        let mut encoded = msg.encode().unwrap();
1002
1003        // Corrupt the data
1004        encoded[Header::SIZE] ^= 0xFF;
1005
1006        // Should fail with CRC verification enabled
1007        let result_with_crc = IgtlMessage::<TransformMessage>::decode_with_options(&encoded, true);
1008        assert!(matches!(
1009            result_with_crc,
1010            Err(crate::error::IgtlError::CrcMismatch { .. })
1011        ));
1012
1013        // Should succeed with CRC verification disabled (even with corrupted data)
1014        let result_without_crc =
1015            IgtlMessage::<TransformMessage>::decode_with_options(&encoded, false);
1016        assert!(result_without_crc.is_ok());
1017    }
1018
1019    #[test]
1020    fn test_decode_default_uses_crc_verification() {
1021        let transform = TransformMessage::identity();
1022        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
1023
1024        let mut encoded = msg.encode().unwrap();
1025
1026        // Corrupt the data
1027        encoded[Header::SIZE] ^= 0xFF;
1028
1029        // Default decode() should verify CRC and fail
1030        let result = IgtlMessage::<TransformMessage>::decode(&encoded);
1031        assert!(matches!(
1032            result,
1033            Err(crate::error::IgtlError::CrcMismatch { .. })
1034        ));
1035    }
1036
1037    #[test]
1038    fn test_crc_skip_performance() {
1039        // This test demonstrates that skipping CRC works correctly
1040        let status = StatusMessage::ok("Performance test");
1041        let msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
1042
1043        let encoded = msg.encode().unwrap();
1044
1045        // Both should decode to the same content (when data is not corrupted)
1046        let decoded_with_crc =
1047            IgtlMessage::<StatusMessage>::decode_with_options(&encoded, true).unwrap();
1048        let decoded_without_crc =
1049            IgtlMessage::<StatusMessage>::decode_with_options(&encoded, false).unwrap();
1050
1051        assert_eq!(decoded_with_crc.content, decoded_without_crc.content);
1052        assert_eq!(decoded_with_crc.content, status);
1053    }
1054
1055    #[test]
1056    fn test_version3_crc_skip_with_extended_header() {
1057        let transform = TransformMessage::identity();
1058        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
1059
1060        msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
1061
1062        let encoded = msg.encode().unwrap();
1063
1064        // Should work with CRC disabled
1065        let decoded =
1066            IgtlMessage::<TransformMessage>::decode_with_options(&encoded, false).unwrap();
1067        assert_eq!(decoded.header.version, 3);
1068        let expected: &[u8] = &[0x01, 0x02, 0x03, 0x04];
1069        assert_eq!(decoded.get_extended_header(), Some(expected));
1070        assert_eq!(decoded.content, transform);
1071    }
1072
1073    #[test]
1074    fn test_version3_crc_skip_with_metadata() {
1075        let status = StatusMessage::ok("Test");
1076        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
1077
1078        msg.add_metadata("key1".to_string(), "value1".to_string());
1079        msg.add_metadata("key2".to_string(), "value2".to_string());
1080
1081        let encoded = msg.encode().unwrap();
1082
1083        // Should work with CRC disabled
1084        let decoded = IgtlMessage::<StatusMessage>::decode_with_options(&encoded, false).unwrap();
1085        assert_eq!(decoded.header.version, 3);
1086
1087        let metadata = decoded.get_metadata().unwrap();
1088        assert_eq!(metadata.get("key1"), Some(&"value1".to_string()));
1089        assert_eq!(metadata.get("key2"), Some(&"value2".to_string()));
1090        assert_eq!(decoded.content, status);
1091    }
1092}
openigtlink_rust/protocol/message.rs

openigtlink_rust/protocol/
message.rs
