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::extended_header::ExtendedHeader;
8use crate::protocol::header::Header;
9use std::collections::HashMap;
10
11/// Common interface for all OpenIGTLink message types
12///
13/// Each message type (TRANSFORM, IMAGE, STATUS, etc.) must implement this trait
14/// to provide encoding/decoding functionality.
15pub trait Message: Sized {
16    /// Returns the message type name (e.g., "TRANSFORM", "IMAGE")
17    ///
18    /// This must match the OpenIGTLink protocol specification.
19    fn message_type() -> &'static str;
20
21    /// Encode message content to bytes
22    ///
23    /// # Returns
24    /// Byte vector containing the encoded message content (without header)
25    fn encode_content(&self) -> Result<Vec<u8>>;
26
27    /// Decode message content from bytes
28    ///
29    /// # Arguments
30    /// * `data` - Byte slice containing the message content (without header)
31    ///
32    /// # Returns
33    /// Decoded message or error
34    fn decode_content(data: &[u8]) -> Result<Self>;
35}
36
37/// Complete OpenIGTLink message structure
38///
39/// Wraps a specific message type with header, optional extended header,
40/// and optional metadata.
41///
42/// # Type Parameters
43/// * `T` - Message type that implements the `Message` trait
44#[derive(Debug)]
45pub struct IgtlMessage<T: Message> {
46    /// Message header (58 bytes)
47    pub header: Header,
48    /// Extended header (Version 3 feature, optional)
49    /// Contains metadata information and message ID when present
50    pub extended_header: Option<ExtendedHeader>,
51    /// Message content
52    pub content: T,
53    /// Metadata as key-value pairs (Version 3 feature, optional)
54    pub metadata: Option<HashMap<String, String>>,
55}
56
57impl<T: Message> IgtlMessage<T> {
58    /// Create a new message with the given content and device name
59    ///
60    /// # Arguments
61    /// * `content` - Message content
62    /// * `device_name` - Device name (max 20 characters)
63    ///
64    /// # Returns
65    /// New message with generated header
66    pub fn new(content: T, device_name: &str) -> Result<Self> {
67        use crate::protocol::header::{DeviceName, Timestamp, TypeName};
68
69        let timestamp = Timestamp::now();
70
71        let content_bytes = content.encode_content()?;
72        let body_size = content_bytes.len() as u64;
73
74        let header = Header {
75            version: 2, // Version 2 compatible
76            type_name: TypeName::new(T::message_type())?,
77            device_name: DeviceName::new(device_name)?,
78            timestamp,
79            body_size,
80            crc: 0, // Will be calculated during encode
81        };
82
83        Ok(IgtlMessage {
84            header,
85            extended_header: None,
86            content,
87            metadata: None,
88        })
89    }
90
91    /// Set extended header data (Version 3 feature)
92    ///
93    /// When extended header is set, the message version is automatically upgraded to 3.
94    ///
95    /// # Arguments
96    /// * `data` - Extended header data as byte vector
97    ///
98    /// # Examples
99    /// ```no_run
100    /// # use openigtlink_rust::protocol::{IgtlMessage, types::TransformMessage};
101    /// let transform = TransformMessage::identity();
102    /// let mut msg = IgtlMessage::new(transform, "Device").unwrap();
103    /// msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
104    /// ```
105    pub fn set_extended_header(&mut self, data: Vec<u8>) {
106        // Try to parse as structured ExtendedHeader, fallback to custom data
107        if let Ok(ext_header) = ExtendedHeader::decode(&data) {
108            self.extended_header = Some(ext_header);
109        } else {
110            // Custom extended header data - wrap in ExtendedHeader with additional_fields
111            let mut ext_header = ExtendedHeader::new();
112            ext_header.set_additional_fields(data);
113            self.extended_header = Some(ext_header);
114        }
115        // Upgrade to version 3 when extended header is used
116        if self.header.version < 3 {
117            self.header.version = 3;
118        }
119    }
120
121    /// Set structured Extended Header (Version 3 feature)
122    ///
123    /// # Arguments
124    /// * `ext_header` - Structured Extended Header
125    pub fn set_extended_header_struct(&mut self, ext_header: ExtendedHeader) {
126        self.extended_header = Some(ext_header);
127        // Upgrade to version 3 when extended header is used
128        if self.header.version < 3 {
129            self.header.version = 3;
130        }
131    }
132
133    /// Get extended header data reference (Version 3 feature)
134    ///
135    /// # Returns
136    /// Optional reference to extended header bytes (additional_fields only, not the full structure)
137    pub fn get_extended_header(&self) -> Option<Vec<u8>> {
138        self.extended_header
139            .as_ref()
140            .map(|h| h.additional_fields.clone())
141    }
142
143    /// Get structured Extended Header reference (Version 3 feature)
144    ///
145    /// # Returns
146    /// Optional reference to structured Extended Header
147    pub fn get_extended_header_struct(&self) -> Option<&ExtendedHeader> {
148        self.extended_header.as_ref()
149    }
150
151    /// Get message ID from Extended Header
152    ///
153    /// # Returns
154    /// Message ID if Extended Header exists, None otherwise
155    pub fn get_message_id(&self) -> Option<u32> {
156        self.extended_header.as_ref().map(|h| h.get_message_id())
157    }
158
159    /// Set message ID in Extended Header
160    ///
161    /// Creates Extended Header if it doesn't exist.
162    ///
163    /// # Arguments
164    /// * `message_id` - Unique message identifier
165    pub fn set_message_id(&mut self, message_id: u32) {
166        if let Some(ext_header) = &mut self.extended_header {
167            ext_header.message_id = message_id;
168        } else {
169            self.extended_header = Some(ExtendedHeader::with_message_id(message_id));
170        }
171        // Upgrade to version 3 when extended header is used
172        if self.header.version < 3 {
173            self.header.version = 3;
174        }
175    }
176
177    /// Remove extended header and optionally downgrade to Version 2
178    pub fn clear_extended_header(&mut self) {
179        self.extended_header = None;
180        // Downgrade to version 2 if no version 3 features are used
181        if self.metadata.is_none() && self.header.version == 3 {
182            self.header.version = 2;
183        }
184    }
185
186    /// Set metadata key-value pairs (Version 3 feature)
187    ///
188    /// When metadata is set, the message version is automatically upgraded to 3.
189    ///
190    /// # Arguments
191    /// * `metadata` - HashMap of key-value pairs
192    ///
193    /// # Examples
194    /// ```no_run
195    /// # use openigtlink_rust::protocol::{IgtlMessage, types::TransformMessage};
196    /// # use std::collections::HashMap;
197    /// let transform = TransformMessage::identity();
198    /// let mut msg = IgtlMessage::new(transform, "Device").unwrap();
199    /// let mut metadata = HashMap::new();
200    /// metadata.insert("priority".to_string(), "high".to_string());
201    /// msg.set_metadata(metadata);
202    /// ```
203    pub fn set_metadata(&mut self, metadata: HashMap<String, String>) {
204        self.metadata = Some(metadata);
205        // Upgrade to version 3 when metadata is used
206        if self.header.version < 3 {
207            self.header.version = 3;
208        }
209    }
210
211    /// Add a single metadata key-value pair (Version 3 feature)
212    ///
213    /// # Arguments
214    /// * `key` - Metadata key
215    /// * `value` - Metadata value
216    pub fn add_metadata(&mut self, key: String, value: String) {
217        if self.metadata.is_none() {
218            self.metadata = Some(HashMap::new());
219            if self.header.version < 3 {
220                self.header.version = 3;
221            }
222        }
223        self.metadata.as_mut().unwrap().insert(key, value);
224    }
225
226    /// Get metadata reference (Version 3 feature)
227    ///
228    /// # Returns
229    /// Optional reference to metadata HashMap
230    pub fn get_metadata(&self) -> Option<&HashMap<String, String>> {
231        self.metadata.as_ref()
232    }
233
234    /// Remove metadata and optionally downgrade to Version 2
235    pub fn clear_metadata(&mut self) {
236        self.metadata = None;
237        // Downgrade to version 2 if no version 3 features are used
238        if self.extended_header.is_none() && self.header.version == 3 {
239            self.header.version = 2;
240        }
241    }
242
243    /// Encode the complete message to bytes
244    ///
245    /// Message format is determined by the presence of extended_header and metadata fields,
246    /// NOT by the version field, as version information may be unreliable.
247    ///
248    /// Format without Extended Header: Header (58) + Content
249    /// Format with Extended Header: Header (58) + ExtHdrSize (2) + ExtHdr (var) + Content + Metadata (var)
250    ///
251    /// Extended Header Size field:
252    /// - 0: No extended header present, content follows immediately after the size field
253    /// - >0: Extended header present, value indicates the size (including this field)
254    ///
255    /// Metadata format (when present):
256    /// - MetadataSize (2 bytes, big-endian)
257    /// - For each pair:
258    ///   - KeySize (2 bytes)
259    ///   - Key (KeySize bytes, UTF-8)
260    ///   - ValueSize (2 bytes)
261    ///   - Value (ValueSize bytes, UTF-8)
262    ///
263    /// # Returns
264    /// Complete message as byte vector
265    pub fn encode(&self) -> Result<Vec<u8>> {
266        use crate::protocol::crc::calculate_crc;
267
268        // 1. Encode content
269        let content_bytes = self.content.encode_content()?;
270
271        // 2. Encode metadata if present
272        // Metadata is encoded based on its presence, not version number
273        // V3 Format: Separate header and body
274        // Header: INDEX_COUNT (2) + [KEY_SIZE (2) + VALUE_ENCODING (2) + VALUE_SIZE (4)]...
275        // Body: [KEY + VALUE]...
276        let (metadata_header, metadata_body) = if self.metadata.is_some() {
277            let metadata = self.metadata.as_ref().unwrap();
278            let mut meta_header = Vec::new();
279            let mut meta_body = Vec::new();
280
281            // INDEX_COUNT (2 bytes)
282            meta_header.extend_from_slice(&(metadata.len() as u16).to_be_bytes());
283
284            // Each key-value pair
285            for (key, value) in metadata.iter() {
286                let key_bytes = key.as_bytes();
287                let value_bytes = value.as_bytes();
288
289                // Header entry: KEY_SIZE (2) + VALUE_ENCODING (2) + VALUE_SIZE (4)
290                meta_header.extend_from_slice(&(key_bytes.len() as u16).to_be_bytes());
291                meta_header.extend_from_slice(&0u16.to_be_bytes()); // VALUE_ENCODING = 0 (UTF-8)
292                meta_header.extend_from_slice(&(value_bytes.len() as u32).to_be_bytes());
293
294                // Body: KEY + VALUE
295                meta_body.extend_from_slice(key_bytes);
296                meta_body.extend_from_slice(value_bytes);
297            }
298
299            (meta_header, meta_body)
300        } else {
301            (Vec::new(), Vec::new())
302        };
303
304        // 3. Build body based on extended header and metadata presence (NOT version)
305        // Extended Header format is determined by the presence of extended_header or metadata fields,
306        // not by the version number, as version information may be unreliable.
307        let body_bytes = if self.extended_header.is_some() {
308            // With extended header
309            let ext_header = self.extended_header.as_ref().unwrap();
310
311            // Update Extended Header with current metadata information
312            let mut ext_header_to_encode = ext_header.clone();
313            if !metadata_header.is_empty() {
314                ext_header_to_encode.metadata_header_size = metadata_header.len() as u16;
315                ext_header_to_encode.metadata_size = metadata_body.len() as u32;
316            }
317
318            let ext_header_encoded = ext_header_to_encode.encode();
319            let ext_header_size = ext_header_encoded.len();
320
321            let mut body = Vec::with_capacity(
322                ext_header_size + content_bytes.len() + metadata_header.len() + metadata_body.len(),
323            );
324            // Extended header (includes size field within it)
325            body.extend_from_slice(&ext_header_encoded);
326            // Content
327            body.extend_from_slice(&content_bytes);
328            // Metadata Header
329            body.extend_from_slice(&metadata_header);
330            // Metadata Body
331            body.extend_from_slice(&metadata_body);
332
333            body
334        } else if !metadata_header.is_empty() {
335            // Without extended header but with metadata - create minimal Extended Header
336            let ext_header = ExtendedHeader::with_metadata(
337                metadata_header.len() as u16,
338                metadata_body.len() as u32,
339            );
340            let ext_header_encoded = ext_header.encode();
341
342            let mut body = Vec::with_capacity(
343                ext_header_encoded.len()
344                    + content_bytes.len()
345                    + metadata_header.len()
346                    + metadata_body.len(),
347            );
348            // Extended header
349            body.extend_from_slice(&ext_header_encoded);
350            // Content
351            body.extend_from_slice(&content_bytes);
352            // Metadata Header
353            body.extend_from_slice(&metadata_header);
354            // Metadata Body
355            body.extend_from_slice(&metadata_body);
356
357            body
358        } else {
359            // No extended header and no metadata - just content
360            content_bytes
361        };
362
363        // 4. Update header with correct body_size and CRC
364        let mut header = self.header.clone();
365        header.body_size = body_bytes.len() as u64;
366        header.crc = calculate_crc(&body_bytes);
367
368        // 5. Combine header + body
369        let mut buf = Vec::with_capacity(Header::SIZE + body_bytes.len());
370        buf.extend_from_slice(&header.encode());
371        buf.extend_from_slice(&body_bytes);
372
373        Ok(buf)
374    }
375
376    /// Decode a complete message from bytes with CRC verification
377    ///
378    /// Automatically detects Extended Header presence based on the extended_header_size field,
379    /// NOT the version field, as version information may be unreliable.
380    ///
381    /// # Arguments
382    /// * `data` - Byte slice containing the complete message
383    ///
384    /// # Returns
385    /// Decoded message or error
386    pub fn decode(data: &[u8]) -> Result<Self> {
387        Self::decode_with_options(data, true)
388    }
389
390    /// Decode a complete message from bytes with optional CRC verification
391    ///
392    /// Allows skipping CRC verification for performance in trusted environments.
393    ///
394    /// # Arguments
395    /// * `data` - Byte slice containing the complete message
396    /// * `verify_crc` - Whether to verify CRC (true = verify, false = skip)
397    ///
398    /// # Returns
399    /// Decoded message or error
400    ///
401    /// # Safety
402    /// Disabling CRC verification (`verify_crc = false`) should only be done in
403    /// trusted environments where data corruption is unlikely (e.g., loopback, local network).
404    /// Using this in production over unreliable networks may lead to silent data corruption.
405    ///
406    /// # Examples
407    /// ```no_run
408    /// # use openigtlink_rust::protocol::{IgtlMessage, types::TransformMessage};
409    /// # let data = vec![0u8; 106];
410    /// // Decode with CRC verification (recommended)
411    /// let msg = IgtlMessage::<TransformMessage>::decode_with_options(&data, true)?;
412    ///
413    /// // Decode without CRC verification (use with caution)
414    /// let msg_fast = IgtlMessage::<TransformMessage>::decode_with_options(&data, false)?;
415    /// # Ok::<(), openigtlink_rust::error::IgtlError>(())
416    /// ```
417    pub fn decode_with_options(data: &[u8], verify_crc: bool) -> Result<Self> {
418        use crate::error::IgtlError;
419        use crate::protocol::crc::calculate_crc;
420
421        if data.len() < Header::SIZE {
422            return Err(IgtlError::InvalidSize {
423                expected: Header::SIZE,
424                actual: data.len(),
425            });
426        }
427
428        // 1. Parse header
429        let header = Header::decode(&data[..Header::SIZE])?;
430
431        // 2. Extract body
432        let body_start = Header::SIZE;
433        let body_end = body_start + header.body_size as usize;
434
435        if data.len() < body_end {
436            return Err(IgtlError::InvalidSize {
437                expected: body_end,
438                actual: data.len(),
439            });
440        }
441
442        let body_bytes = &data[body_start..body_end];
443
444        // 3. Verify CRC (if requested)
445        if verify_crc {
446            let calculated_crc = calculate_crc(body_bytes);
447            if calculated_crc != header.crc {
448                return Err(IgtlError::CrcMismatch {
449                    expected: header.crc,
450                    actual: calculated_crc,
451                });
452            }
453        }
454
455        // 4. Parse body based on Extended Header size field
456        // Extended Header is a Version 3 feature, but version field may be unreliable.
457        //
458        // Detection strategy:
459        // 1. If version < 3: No Extended Header field, treat entire body as content
460        // 2. If version >= 3 and body >= 2 bytes:
461        //    - Check first 2 bytes (extended_header_size):
462        //      - 0: Version 3 without Extended Header → skip 2-byte field
463        //      - >= 12: Extended Header present → parse it
464        //      - 1-11: Invalid → fallback to version check
465        // 3. If ext_header_size >= 12 but version < 3:
466        //    - This indicates unreliable version field
467        //    - Trust ext_header_size and parse Extended Header
468        //
469        // Extended Header structure (OpenIGTLink Version 3):
470        // - extended_header_size (2 bytes) - total size including this field
471        // - metadata_header_size (2 bytes) - size of metadata header section
472        // - metadata_size (4 bytes) - size of metadata data section
473        // - message_id (4 bytes)
474        // - additional fields (variable, if extended_header_size > 12)
475        //
476        // Body structure (C++ implementation):
477        // [Extended Header][Content][Metadata Header][Metadata Data]
478        //                            ^---- at end ----^
479        let (extended_header, content_bytes, metadata) = if body_bytes.len() >= 2 {
480            // Read potential extended_header_size field (first 2 bytes)
481            let ext_header_size = u16::from_be_bytes([body_bytes[0], body_bytes[1]]) as usize;
482
483            if ext_header_size >= ExtendedHeader::MIN_SIZE && body_bytes.len() >= ext_header_size {
484                // ext_header_size looks valid (>= 12), try to parse Extended Header
485                // This works even if version < 3 (unreliable version field)
486                // Extended header is present - try to parse it
487                match ExtendedHeader::decode(&body_bytes[..ext_header_size]) {
488                    Ok(ext_header) => {
489                        let metadata_header_size = ext_header.get_metadata_header_size();
490                        let metadata_size = ext_header.get_metadata_size();
491                        let body_size = body_bytes.len();
492
493                        // Content: from Extended Header to start of metadata
494                        // Metadata: at the end of body
495                        let content_size =
496                            body_size - ext_header_size - metadata_header_size - metadata_size;
497                        let content_start = ext_header_size;
498                        let content_end = content_start + content_size;
499
500                        let content_part = &body_bytes[content_start..content_end];
501
502                        // Parse metadata if present
503                        let parsed_metadata = if metadata_header_size > 0 && metadata_size > 0 {
504                            let meta_header_start =
505                                body_size - metadata_header_size - metadata_size;
506                            let meta_data_start = body_size - metadata_size;
507
508                            let meta_header_part = &body_bytes[meta_header_start..meta_data_start];
509                            let meta_data_part = &body_bytes[meta_data_start..];
510
511                            Self::decode_metadata_v3(meta_header_part, meta_data_part)?
512                        } else {
513                            None
514                        };
515
516                        (Some(ext_header), content_part, parsed_metadata)
517                    }
518                    Err(_) => {
519                        // Failed to parse as standard Extended Header
520                        // Treat entire body as content (Version 1/2 format)
521                        (None, body_bytes, None)
522                    }
523                }
524            } else if ext_header_size == 0 && header.version >= 3 {
525                // ext_header_size = 0 with version >= 3:
526                // Version 3 format explicitly without Extended Header
527                // Skip the 2-byte size field
528                (None, &body_bytes[2..], None)
529            } else {
530                // ext_header_size is 1-11 (invalid) or too small, or version < 3
531                // This is a Version 1/2 message without Extended Header field
532                // Treat entire body as content (don't skip any bytes)
533                (None, body_bytes, None)
534            }
535        } else {
536            // Body too small to contain extended_header_size field - treat as content
537            // This is likely a version 1 message
538            (None, body_bytes, None)
539        };
540
541        // 5. Decode content
542        let content = T::decode_content(content_bytes)?;
543
544        Ok(IgtlMessage {
545            header,
546            extended_header,
547            content,
548            metadata,
549        })
550    }
551
552    /// Decode metadata from V3 format (separate header and data)
553    ///
554    /// # Arguments
555    /// * `header_data` - Metadata header bytes (INDEX_COUNT + entries)
556    /// * `body_data` - Metadata body bytes (keys + values)
557    ///
558    /// # Format
559    /// Header: INDEX_COUNT (2) + [KEY_SIZE (2) + VALUE_ENCODING (2) + VALUE_SIZE (4)]...
560    /// Body: [KEY + VALUE]...
561    fn decode_metadata_v3(
562        header_data: &[u8],
563        body_data: &[u8],
564    ) -> Result<Option<HashMap<String, String>>> {
565        use crate::error::IgtlError;
566
567        if header_data.len() < 2 {
568            return Ok(None);
569        }
570
571        // Parse INDEX_COUNT
572        let index_count = u16::from_be_bytes([header_data[0], header_data[1]]) as usize;
573
574        if index_count == 0 {
575            return Ok(None);
576        }
577
578        // Parse metadata header entries
579        let mut entries = Vec::with_capacity(index_count);
580        let mut header_offset = 2;
581
582        for _ in 0..index_count {
583            if header_offset + 8 > header_data.len() {
584                return Err(IgtlError::InvalidSize {
585                    expected: header_offset + 8,
586                    actual: header_data.len(),
587                });
588            }
589
590            let key_size =
591                u16::from_be_bytes([header_data[header_offset], header_data[header_offset + 1]]);
592            let _value_encoding = u16::from_be_bytes([
593                header_data[header_offset + 2],
594                header_data[header_offset + 3],
595            ]);
596            let value_size = u32::from_be_bytes([
597                header_data[header_offset + 4],
598                header_data[header_offset + 5],
599                header_data[header_offset + 6],
600                header_data[header_offset + 7],
601            ]);
602
603            entries.push((key_size, value_size));
604            header_offset += 8;
605        }
606
607        // Parse metadata body using header entries
608        let mut metadata = HashMap::new();
609        let mut body_offset = 0;
610
611        for (key_size, value_size) in entries {
612            let key_size = key_size as usize;
613            let value_size = value_size as usize;
614
615            // Read key
616            if body_offset + key_size > body_data.len() {
617                return Err(IgtlError::InvalidSize {
618                    expected: body_offset + key_size,
619                    actual: body_data.len(),
620                });
621            }
622            let key = String::from_utf8(body_data[body_offset..body_offset + key_size].to_vec())?;
623            body_offset += key_size;
624
625            // Read value
626            if body_offset + value_size > body_data.len() {
627                return Err(IgtlError::InvalidSize {
628                    expected: body_offset + value_size,
629                    actual: body_data.len(),
630                });
631            }
632            let value =
633                String::from_utf8(body_data[body_offset..body_offset + value_size].to_vec())?;
634            body_offset += value_size;
635
636            metadata.insert(key, value);
637        }
638
639        Ok(Some(metadata))
640    }
641}
642
643#[cfg(test)]
644mod tests {
645    use super::*;
646    use crate::protocol::types::{CapabilityMessage, StatusMessage, TransformMessage};
647
648    // Mock message type for testing
649    struct TestMessage {
650        data: Vec<u8>,
651    }
652
653    impl Message for TestMessage {
654        fn message_type() -> &'static str {
655            "TEST"
656        }
657
658        fn encode_content(&self) -> Result<Vec<u8>> {
659            Ok(self.data.clone())
660        }
661
662        fn decode_content(data: &[u8]) -> Result<Self> {
663            Ok(TestMessage {
664                data: data.to_vec(),
665            })
666        }
667    }
668
669    #[test]
670    fn test_message_trait() {
671        assert_eq!(TestMessage::message_type(), "TEST");
672    }
673
674    #[test]
675    fn test_message_encode_decode() {
676        let original = TestMessage {
677            data: vec![1, 2, 3, 4, 5],
678        };
679
680        let encoded = original.encode_content().unwrap();
681        let decoded = TestMessage::decode_content(&encoded).unwrap();
682
683        assert_eq!(original.data, decoded.data);
684    }
685
686    #[test]
687    fn test_full_message_roundtrip_transform() {
688        let transform = TransformMessage::identity();
689        let msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
690
691        let encoded = msg.encode().unwrap();
692        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
693
694        // Verify header fields
695        assert_eq!(decoded.header.version, 2);
696        assert_eq!(decoded.header.type_name.as_str().unwrap(), "TRANSFORM");
697        assert_eq!(decoded.header.device_name.as_str().unwrap(), "TestDevice");
698        assert_eq!(decoded.header.body_size, 48);
699
700        // Verify content
701        assert_eq!(decoded.content, transform);
702    }
703
704    #[test]
705    fn test_full_message_roundtrip_status() {
706        let status = StatusMessage::ok("Operation successful");
707        let msg = IgtlMessage::new(status.clone(), "StatusDevice").unwrap();
708
709        let encoded = msg.encode().unwrap();
710        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
711
712        assert_eq!(decoded.header.type_name.as_str().unwrap(), "STATUS");
713        assert_eq!(decoded.content, status);
714    }
715
716    #[test]
717    fn test_full_message_roundtrip_capability() {
718        let capability = CapabilityMessage::new(vec![
719            "TRANSFORM".to_string(),
720            "STATUS".to_string(),
721            "IMAGE".to_string(),
722        ]);
723        let msg = IgtlMessage::new(capability.clone(), "CapDevice").unwrap();
724
725        let encoded = msg.encode().unwrap();
726        let decoded = IgtlMessage::<CapabilityMessage>::decode(&encoded).unwrap();
727
728        assert_eq!(decoded.header.type_name.as_str().unwrap(), "CAPABILITY");
729        assert_eq!(decoded.content, capability);
730    }
731
732    #[test]
733    fn test_timestamp_reasonable() {
734        let transform = TransformMessage::identity();
735        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
736
737        // Timestamp should be recent (within last year and not in future)
738        let now = std::time::SystemTime::now()
739            .duration_since(std::time::UNIX_EPOCH)
740            .unwrap()
741            .as_secs() as u32;
742
743        let one_year_ago = now - (365 * 24 * 60 * 60);
744
745        assert!(msg.header.timestamp.seconds >= one_year_ago);
746        assert!(msg.header.timestamp.seconds <= now + 1); // +1 for clock skew
747    }
748
749    #[test]
750    fn test_crc_verification() {
751        let transform = TransformMessage::identity();
752        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
753
754        let mut encoded = msg.encode().unwrap();
755
756        // Corrupt the content
757        let content_start = Header::SIZE;
758        encoded[content_start] ^= 0xFF;
759
760        // Should fail CRC check
761        let result = IgtlMessage::<TransformMessage>::decode(&encoded);
762        assert!(matches!(
763            result,
764            Err(crate::error::IgtlError::CrcMismatch { .. })
765        ));
766    }
767
768    #[test]
769    fn test_message_size_calculation() {
770        let transform = TransformMessage::identity();
771        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
772
773        let encoded = msg.encode().unwrap();
774
775        // Total size should be: Header (58) + TRANSFORM content (48) = 106
776        assert_eq!(encoded.len(), 106);
777    }
778
779    #[test]
780    fn test_decode_short_buffer() {
781        let short_data = vec![0u8; 30];
782        let result = IgtlMessage::<TransformMessage>::decode(&short_data);
783        assert!(matches!(
784            result,
785            Err(crate::error::IgtlError::InvalidSize { .. })
786        ));
787    }
788
789    // Version 3 Extended Header Tests
790
791    #[test]
792    fn test_extended_header_set_get() {
793        let transform = TransformMessage::identity();
794        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
795
796        // Initially no extended header
797        assert_eq!(msg.get_extended_header(), None);
798        assert_eq!(msg.header.version, 2);
799
800        // Set extended header
801        let ext_header = vec![0x01, 0x02, 0x03, 0x04];
802        msg.set_extended_header(ext_header.clone());
803
804        // Should upgrade to version 3
805        assert_eq!(msg.header.version, 3);
806        assert_eq!(msg.get_extended_header(), Some(ext_header.clone()));
807    }
808
809    #[test]
810    fn test_extended_header_clear() {
811        let transform = TransformMessage::identity();
812        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
813
814        // Set extended header
815        msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
816        assert_eq!(msg.header.version, 3);
817
818        // Clear extended header
819        msg.clear_extended_header();
820        assert_eq!(msg.get_extended_header(), None);
821        // Should downgrade to version 2 (no metadata either)
822        assert_eq!(msg.header.version, 2);
823    }
824
825    #[test]
826    fn test_version3_encode_decode_with_extended_header() {
827        let transform = TransformMessage::identity();
828        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
829
830        // Add extended header
831        let ext_header = vec![0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF];
832        msg.set_extended_header(ext_header.clone());
833
834        // Encode
835        let encoded = msg.encode().unwrap();
836
837        // Decode
838        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
839
840        // Verify version
841        assert_eq!(decoded.header.version, 3);
842
843        // Verify extended header
844        assert_eq!(decoded.get_extended_header(), Some(ext_header.clone()));
845
846        // Verify content
847        assert_eq!(decoded.content, transform);
848    }
849
850    #[test]
851    fn test_version3_encode_decode_without_extended_header() {
852        let status = StatusMessage::ok("Test message");
853        let msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
854
855        // Encode as Version 2
856        let encoded = msg.encode().unwrap();
857
858        // Decode
859        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
860
861        // Should remain version 2
862        assert_eq!(decoded.header.version, 2);
863        assert_eq!(decoded.get_extended_header(), None);
864        assert_eq!(decoded.content, status);
865    }
866
867    #[test]
868    fn test_extended_header_body_size_calculation() {
869        let transform = TransformMessage::identity();
870        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
871
872        // Add extended header
873        let ext_header = vec![0x01, 0x02, 0x03, 0x04]; // 4 bytes
874        msg.set_extended_header(ext_header);
875
876        let encoded = msg.encode().unwrap();
877
878        // Total size should be:
879        // Header (58) + ExtendedHeader (12 fixed + 4 additional_fields) + TRANSFORM content (48) = 122
880        assert_eq!(encoded.len(), 122);
881
882        // Verify body_size in header includes extended header overhead
883        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
884        assert_eq!(decoded.header.body_size, 16 + 48); // ExtendedHeader (12+4) + Content
885    }
886
887    #[test]
888    fn test_extended_header_empty() {
889        let transform = TransformMessage::identity();
890        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
891
892        // Set empty extended header
893        msg.set_extended_header(vec![]);
894
895        let encoded = msg.encode().unwrap();
896        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
897
898        // Should still be version 3
899        assert_eq!(decoded.header.version, 3);
900        // Extended header should be empty
901        assert_eq!(decoded.get_extended_header(), Some(vec![]));
902        // Content should be intact
903        assert_eq!(decoded.content, transform);
904    }
905
906    #[test]
907    fn test_extended_header_large() {
908        let status = StatusMessage::ok("Test");
909        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
910
911        // Create large extended header (1 KB)
912        let ext_header = vec![0xAB; 1024];
913        msg.set_extended_header(ext_header.clone());
914
915        let encoded = msg.encode().unwrap();
916        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
917
918        assert_eq!(decoded.header.version, 3);
919        assert_eq!(decoded.get_extended_header(), Some(ext_header.clone()));
920        assert_eq!(decoded.content, status);
921    }
922
923    #[test]
924    fn test_version3_crc_includes_extended_header() {
925        let transform = TransformMessage::identity();
926        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
927
928        msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
929
930        let mut encoded = msg.encode().unwrap();
931
932        // Corrupt extended header
933        encoded[Header::SIZE + 2] ^= 0xFF; // First byte of extended header
934
935        // Should fail CRC check
936        let result = IgtlMessage::<TransformMessage>::decode(&encoded);
937        assert!(matches!(
938            result,
939            Err(crate::error::IgtlError::CrcMismatch { .. })
940        ));
941    }
942
943    #[test]
944    fn test_backward_compatibility_version2() {
945        // Create a Version 2 message
946        let capability = CapabilityMessage::new(vec!["TRANSFORM".to_string()]);
947        let msg = IgtlMessage::new(capability.clone(), "Device").unwrap();
948
949        assert_eq!(msg.header.version, 2);
950
951        let encoded = msg.encode().unwrap();
952        let decoded = IgtlMessage::<CapabilityMessage>::decode(&encoded).unwrap();
953
954        // Should decode correctly as Version 2
955        assert_eq!(decoded.header.version, 2);
956        assert_eq!(decoded.get_extended_header(), None);
957        assert_eq!(decoded.content, capability);
958    }
959
960    // Version 3 Metadata Tests
961
962    #[test]
963    fn test_metadata_set_get() {
964        let transform = TransformMessage::identity();
965        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
966
967        // Initially no metadata
968        assert_eq!(msg.get_metadata(), None);
969        assert_eq!(msg.header.version, 2);
970
971        // Set metadata
972        let mut metadata = HashMap::new();
973        metadata.insert("priority".to_string(), "high".to_string());
974        metadata.insert("sequence".to_string(), "42".to_string());
975        msg.set_metadata(metadata.clone());
976
977        // Should upgrade to version 3
978        assert_eq!(msg.header.version, 3);
979        assert_eq!(msg.get_metadata(), Some(&metadata));
980    }
981
982    #[test]
983    fn test_metadata_add() {
984        let status = StatusMessage::ok("Test");
985        let mut msg = IgtlMessage::new(status, "TestDevice").unwrap();
986
987        assert_eq!(msg.header.version, 2);
988
989        // Add metadata one by one
990        msg.add_metadata("key1".to_string(), "value1".to_string());
991        assert_eq!(msg.header.version, 3);
992
993        msg.add_metadata("key2".to_string(), "value2".to_string());
994
995        let metadata = msg.get_metadata().unwrap();
996        assert_eq!(metadata.get("key1"), Some(&"value1".to_string()));
997        assert_eq!(metadata.get("key2"), Some(&"value2".to_string()));
998    }
999
1000    #[test]
1001    fn test_metadata_clear() {
1002        let transform = TransformMessage::identity();
1003        let mut msg = IgtlMessage::new(transform, "TestDevice").unwrap();
1004
1005        // Set metadata
1006        msg.add_metadata("test".to_string(), "value".to_string());
1007        assert_eq!(msg.header.version, 3);
1008
1009        // Clear metadata
1010        msg.clear_metadata();
1011        assert_eq!(msg.get_metadata(), None);
1012        // Should downgrade to version 2 (no extended header either)
1013        assert_eq!(msg.header.version, 2);
1014    }
1015
1016    #[test]
1017    fn test_version3_encode_decode_with_metadata() {
1018        let transform = TransformMessage::identity();
1019        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
1020
1021        // Add metadata
1022        let mut metadata = HashMap::new();
1023        metadata.insert("priority".to_string(), "high".to_string());
1024        metadata.insert("timestamp".to_string(), "123456".to_string());
1025        msg.set_metadata(metadata.clone());
1026
1027        // Encode
1028        let encoded = msg.encode().unwrap();
1029
1030        // Decode
1031        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
1032
1033        // Verify version
1034        assert_eq!(decoded.header.version, 3);
1035
1036        // Verify metadata
1037        let decoded_metadata = decoded.get_metadata().unwrap();
1038        assert_eq!(decoded_metadata.get("priority"), Some(&"high".to_string()));
1039        assert_eq!(
1040            decoded_metadata.get("timestamp"),
1041            Some(&"123456".to_string())
1042        );
1043
1044        // Verify content
1045        assert_eq!(decoded.content, transform);
1046    }
1047
1048    #[test]
1049    fn test_version3_with_extended_header_and_metadata() {
1050        let status = StatusMessage::ok("Test message");
1051        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
1052
1053        // Add both extended header and metadata
1054        msg.set_extended_header(vec![0xAA, 0xBB, 0xCC, 0xDD]);
1055        msg.add_metadata("key1".to_string(), "value1".to_string());
1056        msg.add_metadata("key2".to_string(), "value2".to_string());
1057
1058        let encoded = msg.encode().unwrap();
1059        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
1060
1061        // Verify version
1062        assert_eq!(decoded.header.version, 3);
1063
1064        // Verify extended header
1065        let expected_ext_header = vec![0xAA, 0xBB, 0xCC, 0xDD];
1066        assert_eq!(decoded.get_extended_header(), Some(expected_ext_header));
1067
1068        // Verify metadata
1069        let metadata = decoded.get_metadata().unwrap();
1070        assert_eq!(metadata.get("key1"), Some(&"value1".to_string()));
1071        assert_eq!(metadata.get("key2"), Some(&"value2".to_string()));
1072
1073        // Verify content
1074        assert_eq!(decoded.content, status);
1075    }
1076
1077    #[test]
1078    fn test_metadata_empty() {
1079        let transform = TransformMessage::identity();
1080        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
1081
1082        // Set empty metadata
1083        msg.set_metadata(HashMap::new());
1084
1085        let encoded = msg.encode().unwrap();
1086        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
1087
1088        // Should not have metadata (empty HashMap)
1089        assert_eq!(decoded.get_metadata(), None);
1090        assert_eq!(decoded.content, transform);
1091    }
1092
1093    #[test]
1094    fn test_metadata_utf8_values() {
1095        let status = StatusMessage::ok("Test");
1096        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
1097
1098        // Add UTF-8 metadata
1099        msg.add_metadata("name".to_string(), "日本語".to_string());
1100        msg.add_metadata("emoji".to_string(), "🎉✨".to_string());
1101
1102        let encoded = msg.encode().unwrap();
1103        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
1104
1105        let metadata = decoded.get_metadata().unwrap();
1106        assert_eq!(metadata.get("name"), Some(&"日本語".to_string()));
1107        assert_eq!(metadata.get("emoji"), Some(&"🎉✨".to_string()));
1108    }
1109
1110    // CRC Verification Tests
1111
1112    #[test]
1113    fn test_decode_with_crc_verification_enabled() {
1114        let transform = TransformMessage::identity();
1115        let msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
1116
1117        let encoded = msg.encode().unwrap();
1118
1119        // Should decode successfully with CRC verification
1120        let decoded = IgtlMessage::<TransformMessage>::decode_with_options(&encoded, true).unwrap();
1121        assert_eq!(decoded.content, transform);
1122    }
1123
1124    #[test]
1125    fn test_decode_with_crc_verification_disabled() {
1126        let transform = TransformMessage::identity();
1127        let msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
1128
1129        let mut encoded = msg.encode().unwrap();
1130
1131        // Corrupt the data
1132        encoded[Header::SIZE] ^= 0xFF;
1133
1134        // Should fail with CRC verification enabled
1135        let result_with_crc = IgtlMessage::<TransformMessage>::decode_with_options(&encoded, true);
1136        assert!(matches!(
1137            result_with_crc,
1138            Err(crate::error::IgtlError::CrcMismatch { .. })
1139        ));
1140
1141        // Should succeed with CRC verification disabled (even with corrupted data)
1142        let result_without_crc =
1143            IgtlMessage::<TransformMessage>::decode_with_options(&encoded, false);
1144        assert!(result_without_crc.is_ok());
1145    }
1146
1147    #[test]
1148    fn test_decode_default_uses_crc_verification() {
1149        let transform = TransformMessage::identity();
1150        let msg = IgtlMessage::new(transform, "TestDevice").unwrap();
1151
1152        let mut encoded = msg.encode().unwrap();
1153
1154        // Corrupt the data
1155        encoded[Header::SIZE] ^= 0xFF;
1156
1157        // Default decode() should verify CRC and fail
1158        let result = IgtlMessage::<TransformMessage>::decode(&encoded);
1159        assert!(matches!(
1160            result,
1161            Err(crate::error::IgtlError::CrcMismatch { .. })
1162        ));
1163    }
1164
1165    #[test]
1166    fn test_crc_skip_performance() {
1167        // This test demonstrates that skipping CRC works correctly
1168        let status = StatusMessage::ok("Performance test");
1169        let msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
1170
1171        let encoded = msg.encode().unwrap();
1172
1173        // Both should decode to the same content (when data is not corrupted)
1174        let decoded_with_crc =
1175            IgtlMessage::<StatusMessage>::decode_with_options(&encoded, true).unwrap();
1176        let decoded_without_crc =
1177            IgtlMessage::<StatusMessage>::decode_with_options(&encoded, false).unwrap();
1178
1179        assert_eq!(decoded_with_crc.content, decoded_without_crc.content);
1180        assert_eq!(decoded_with_crc.content, status);
1181    }
1182
1183    #[test]
1184    fn test_version3_crc_skip_with_extended_header() {
1185        let transform = TransformMessage::identity();
1186        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
1187
1188        msg.set_extended_header(vec![0x01, 0x02, 0x03, 0x04]);
1189
1190        let encoded = msg.encode().unwrap();
1191
1192        // Should work with CRC disabled
1193        let decoded =
1194            IgtlMessage::<TransformMessage>::decode_with_options(&encoded, false).unwrap();
1195        assert_eq!(decoded.header.version, 3);
1196        let expected = vec![0x01, 0x02, 0x03, 0x04];
1197        assert_eq!(decoded.get_extended_header(), Some(expected));
1198        assert_eq!(decoded.content, transform);
1199    }
1200
1201    #[test]
1202    fn test_version3_crc_skip_with_metadata() {
1203        let status = StatusMessage::ok("Test");
1204        let mut msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
1205
1206        msg.add_metadata("key1".to_string(), "value1".to_string());
1207        msg.add_metadata("key2".to_string(), "value2".to_string());
1208
1209        let encoded = msg.encode().unwrap();
1210
1211        // Should work with CRC disabled
1212        let decoded = IgtlMessage::<StatusMessage>::decode_with_options(&encoded, false).unwrap();
1213        assert_eq!(decoded.header.version, 3);
1214
1215        let metadata = decoded.get_metadata().unwrap();
1216        assert_eq!(metadata.get("key1"), Some(&"value1".to_string()));
1217        assert_eq!(metadata.get("key2"), Some(&"value2".to_string()));
1218        assert_eq!(decoded.content, status);
1219    }
1220
1221    #[test]
1222    fn test_extended_header_detection_by_field_not_version() {
1223        // This test verifies that Extended Header presence is determined by the
1224        // extended_header_size field, NOT the version field.
1225
1226        use crate::protocol::crc::calculate_crc;
1227
1228        let transform = TransformMessage::identity();
1229        let mut msg = IgtlMessage::new(transform.clone(), "TestDevice").unwrap();
1230
1231        // Add extended header
1232        let ext_header = vec![0xAA, 0xBB, 0xCC, 0xDD];
1233        msg.set_extended_header(ext_header.clone());
1234
1235        // Encode message
1236        let mut encoded = msg.encode().unwrap();
1237
1238        // MANUALLY set version to 1 in the header to simulate unreliable version info
1239        // Version field is at bytes 0-1 of the header
1240        encoded[0] = 0;
1241        encoded[1] = 1;
1242
1243        // Read body_size from the encoded header (bytes 42-49, big-endian u64)
1244        let body_size = u64::from_be_bytes([
1245            encoded[42],
1246            encoded[43],
1247            encoded[44],
1248            encoded[45],
1249            encoded[46],
1250            encoded[47],
1251            encoded[48],
1252            encoded[49],
1253        ]);
1254
1255        // Recalculate CRC (body didn't change, but we need consistent CRC)
1256        let body_start = Header::SIZE;
1257        let body_end = body_start + (body_size as usize);
1258        let body = &encoded[body_start..body_end];
1259        let new_crc = calculate_crc(body);
1260
1261        // Update CRC in header (bytes 50-57)
1262        encoded[50..58].copy_from_slice(&new_crc.to_be_bytes());
1263
1264        // Decode - should still detect extended header based on extended_header_size field
1265        let decoded = IgtlMessage::<TransformMessage>::decode(&encoded).unwrap();
1266
1267        // Version field should be 1 (as we set it)
1268        assert_eq!(decoded.header.version, 1);
1269
1270        // But extended header should still be detected and decoded correctly!
1271        assert_eq!(decoded.get_extended_header(), Some(ext_header.clone()));
1272        assert_eq!(decoded.content, transform);
1273    }
1274
1275    #[test]
1276    fn test_no_extended_header_with_zero_size_field() {
1277        // Test that extended_header_size = 0 means no extended header,
1278        // regardless of version
1279
1280        use crate::protocol::crc::calculate_crc;
1281
1282        let status = StatusMessage::ok("Test");
1283        let msg = IgtlMessage::new(status.clone(), "TestDevice").unwrap();
1284
1285        // Create a message with extended_header_size = 0 but version = 3
1286        let content_bytes = status.encode_content().unwrap();
1287
1288        // Build body with ext_header_size = 0
1289        let mut body = Vec::new();
1290        body.extend_from_slice(&0u16.to_be_bytes()); // ext_header_size = 0
1291        body.extend_from_slice(&content_bytes);
1292
1293        let crc = calculate_crc(&body);
1294
1295        // Build header with version 3
1296        let mut header = msg.header.clone();
1297        header.version = 3;
1298        header.body_size = body.len() as u64;
1299        header.crc = crc;
1300
1301        // Combine header + body
1302        let mut encoded = Vec::new();
1303        encoded.extend_from_slice(&header.encode());
1304        encoded.extend_from_slice(&body);
1305
1306        // Decode - should not have extended header
1307        let decoded = IgtlMessage::<StatusMessage>::decode(&encoded).unwrap();
1308
1309        assert_eq!(decoded.header.version, 3);
1310        assert_eq!(decoded.get_extended_header(), None);
1311        assert_eq!(decoded.content, status);
1312    }
1313}
openigtlink_rust/protocol/message.rs

openigtlink_rust/protocol/
message.rs
