Skip to main content

consortium/
communication.rs

1//! Inter-node messaging protocol.
2//!
3//! Rust implementation of `ClusterShell.Communication`.
4//!
5//! Messages are the fundamental units of communication between nodes in the
6//! propagation tree. They are serialized to/from XML for transmission over
7//! SSH channels. Each message has a type identifier, a unique message ID,
8//! optional attributes, and an optional payload (base64-encoded bytes).
9//!
10//! The [`Channel`] struct manages bidirectional communication using an
11//! [`XmlReader`] for parsing incoming XML and XML generation for outgoing
12//! messages.
13
14use std::collections::{HashMap, VecDeque};
15use std::sync::atomic::{AtomicU64, Ordering};
16
17use thiserror::Error;
18
19// ============================================================================
20// Constants
21// ============================================================================
22
23/// XML character encoding
24pub const ENCODING: &str = "utf-8";
25
26/// Default base64 line length for large payloads
27pub const DEFAULT_B64_LINE_LENGTH: usize = 65536;
28
29// Channel stream name constants
30/// Stream name for the writer side of a channel
31pub const SNAME_WRITER: &str = "ch-writer";
32/// Stream name for the reader side of a channel
33pub const SNAME_READER: &str = "ch-reader";
34/// Stream name for the error side of a channel
35pub const SNAME_ERROR: &str = "ch-error";
36
37// Message type identifier constants
38const IDENT_GEN: &str = "GEN";
39const IDENT_CFG: &str = "CFG";
40const IDENT_CTL: &str = "CTL";
41const IDENT_ACK: &str = "ACK";
42const IDENT_ERR: &str = "ERR";
43const IDENT_OUT: &str = "OUT";
44const IDENT_SER: &str = "SER";
45const IDENT_RET: &str = "RET";
46const IDENT_TIM: &str = "TIM";
47const IDENT_RTR: &str = "RTR";
48const IDENT_CHA: &str = "CHA";
49const IDENT_END: &str = "END";
50
51// ============================================================================
52// Errors
53// ============================================================================
54
55/// Errors raised during message processing.
56#[derive(Error, Debug)]
57pub enum MessageProcessingError {
58    /// An unknown or invalid message type was encountered.
59    #[error("Unknown message type: {0}")]
60    UnknownType(String),
61
62    /// A required attribute is missing from a message.
63    #[error("Invalid message attributes: missing key \"{0}\"")]
64    MissingAttribute(String),
65
66    /// An unexpected payload was found on a message that doesn't support one.
67    #[error("Got unexpected payload for Message {0}")]
68    UnexpectedPayload(String),
69
70    /// The message payload is invalid or corrupted.
71    #[error("Message {0} has an invalid payload")]
72    InvalidPayload(String),
73
74    /// XML parse error.
75    #[error("Parse error: {0}")]
76    ParseError(String),
77
78    /// Invalid starting XML tag.
79    #[error("Invalid starting tag {0}")]
80    InvalidTag(String),
81}
82
83// ============================================================================
84// Global message ID counter
85// ============================================================================
86
87static MSG_COUNTER: AtomicU64 = AtomicU64::new(0);
88
89fn next_msg_id() -> u64 {
90    MSG_COUNTER.fetch_add(1, Ordering::Relaxed)
91}
92
93/// Reset the global message counter (for testing only).
94#[cfg(test)]
95fn reset_msg_counter() {
96    MSG_COUNTER.store(0, Ordering::Relaxed);
97}
98
99// ============================================================================
100// Minimal base64
101// ============================================================================
102
103/// Simple base64 alphabet (standard RFC 4648).
104const B64_CHARS: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
105
106/// Encode bytes to base64 string.
107pub fn base64_encode(input: &[u8]) -> String {
108    let mut out = String::with_capacity((input.len() + 2) / 3 * 4);
109    for chunk in input.chunks(3) {
110        let b0 = chunk[0] as u32;
111        let b1 = if chunk.len() > 1 { chunk[1] as u32 } else { 0 };
112        let b2 = if chunk.len() > 2 { chunk[2] as u32 } else { 0 };
113        let triple = (b0 << 16) | (b1 << 8) | b2;
114
115        out.push(B64_CHARS[((triple >> 18) & 0x3F) as usize] as char);
116        out.push(B64_CHARS[((triple >> 12) & 0x3F) as usize] as char);
117        if chunk.len() > 1 {
118            out.push(B64_CHARS[((triple >> 6) & 0x3F) as usize] as char);
119        } else {
120            out.push('=');
121        }
122        if chunk.len() > 2 {
123            out.push(B64_CHARS[(triple & 0x3F) as usize] as char);
124        } else {
125            out.push('=');
126        }
127    }
128    out
129}
130
131/// Decode base64 string to bytes. Ignores whitespace/newlines (RFC 4648 relaxed).
132pub fn base64_decode(input: &str) -> Result<Vec<u8>, MessageProcessingError> {
133    fn b64_val(c: u8) -> Result<u8, MessageProcessingError> {
134        match c {
135            b'A'..=b'Z' => Ok(c - b'A'),
136            b'a'..=b'z' => Ok(c - b'a' + 26),
137            b'0'..=b'9' => Ok(c - b'0' + 52),
138            b'+' => Ok(62),
139            b'/' => Ok(63),
140            _ => Err(MessageProcessingError::InvalidPayload(
141                "invalid base64 character".into(),
142            )),
143        }
144    }
145
146    // Strip whitespace
147    let clean: Vec<u8> = input.bytes().filter(|b| !b.is_ascii_whitespace()).collect();
148    if clean.is_empty() {
149        return Ok(Vec::new());
150    }
151
152    let mut out = Vec::with_capacity(clean.len() * 3 / 4);
153    for chunk in clean.chunks(4) {
154        if chunk.len() < 2 {
155            return Err(MessageProcessingError::InvalidPayload(
156                "truncated base64".into(),
157            ));
158        }
159        let a = b64_val(chunk[0])?;
160        let b = b64_val(chunk[1])?;
161        out.push((a << 2) | (b >> 4));
162        if chunk.len() > 2 && chunk[2] != b'=' {
163            let c = b64_val(chunk[2])?;
164            out.push((b << 4) | (c >> 2));
165            if chunk.len() > 3 && chunk[3] != b'=' {
166                let d = b64_val(chunk[3])?;
167                out.push((c << 6) | d);
168            }
169        }
170    }
171    Ok(out)
172}
173
174// ============================================================================
175// XML escape helpers
176// ============================================================================
177
178fn xml_escape(s: &str) -> String {
179    let mut out = String::with_capacity(s.len());
180    for c in s.chars() {
181        match c {
182            '&' => out.push_str("&amp;"),
183            '<' => out.push_str("&lt;"),
184            '>' => out.push_str("&gt;"),
185            '"' => out.push_str("&quot;"),
186            '\'' => out.push_str("&apos;"),
187            _ => out.push(c),
188        }
189    }
190    out
191}
192
193fn xml_unescape(s: &str) -> String {
194    s.replace("&amp;", "&")
195        .replace("&lt;", "<")
196        .replace("&gt;", ">")
197        .replace("&quot;", "\"")
198        .replace("&apos;", "'")
199}
200
201// ============================================================================
202// Message
203// ============================================================================
204
205/// A communication message in the propagation tree protocol.
206///
207/// Messages are tagged unions representing the different message types
208/// exchanged between ClusterShell nodes. Each variant carries the fields
209/// specific to that message type.
210#[derive(Debug, Clone)]
211pub enum Message {
212    /// Generic base message.
213    General { msgid: u64, data: Option<Vec<u8>> },
214
215    /// Configuration propagation container (ident = "CFG").
216    /// Has payload. Carries serialized topology.
217    Configuration {
218        msgid: u64,
219        gateway: String,
220        data: Option<Vec<u8>>,
221    },
222
223    /// Action request routed through the tree (ident = "CTL").
224    /// Has payload. Carries serialized command data.
225    Control {
226        msgid: u64,
227        srcid: u64,
228        action: String,
229        target: String,
230        data: Option<Vec<u8>>,
231    },
232
233    /// Acknowledgement message (ident = "ACK").
234    Ack { msgid: u64, ack: u64 },
235
236    /// Error message (ident = "ERR").
237    Error { msgid: u64, reason: String },
238
239    /// Standard output container (ident = "OUT").
240    /// Has payload. Carries encoded output data.
241    StdOut {
242        msgid: u64,
243        srcid: u64,
244        nodes: String,
245        data: Option<Vec<u8>>,
246    },
247
248    /// Standard error container (ident = "SER").
249    /// Has payload. Carries encoded stderr data.
250    StdErr {
251        msgid: u64,
252        srcid: u64,
253        nodes: String,
254        data: Option<Vec<u8>>,
255    },
256
257    /// Return code container (ident = "RET").
258    Retcode {
259        msgid: u64,
260        srcid: u64,
261        retcode: i32,
262        nodes: String,
263    },
264
265    /// Timeout notification (ident = "TIM").
266    Timeout {
267        msgid: u64,
268        srcid: u64,
269        nodes: String,
270    },
271
272    /// Routing notification (ident = "RTR").
273    Routing {
274        msgid: u64,
275        srcid: u64,
276        event: String,
277        gateway: String,
278        targets: String,
279    },
280
281    /// Start of channel communication (ident = "CHA").
282    Start { msgid: u64 },
283
284    /// End of channel communication (ident = "END").
285    End { msgid: u64 },
286}
287
288impl Message {
289    // -- Constructors --------------------------------------------------------
290
291    /// Create a new General message.
292    pub fn general() -> Self {
293        Message::General {
294            msgid: next_msg_id(),
295            data: None,
296        }
297    }
298
299    /// Create a new Configuration message.
300    pub fn configuration(gateway: &str) -> Self {
301        Message::Configuration {
302            msgid: next_msg_id(),
303            gateway: gateway.to_string(),
304            data: None,
305        }
306    }
307
308    /// Create a new Control message.
309    pub fn control(srcid: u64) -> Self {
310        Message::Control {
311            msgid: next_msg_id(),
312            srcid,
313            action: String::new(),
314            target: String::new(),
315            data: None,
316        }
317    }
318
319    /// Create a new Ack message.
320    pub fn ack(ackid: u64) -> Self {
321        Message::Ack {
322            msgid: next_msg_id(),
323            ack: ackid,
324        }
325    }
326
327    /// Create a new Error message.
328    pub fn error(reason: &str) -> Self {
329        Message::Error {
330            msgid: next_msg_id(),
331            reason: reason.to_string(),
332        }
333    }
334
335    /// Create a new StdOut message.
336    pub fn stdout(nodes: &str, srcid: u64) -> Self {
337        Message::StdOut {
338            msgid: next_msg_id(),
339            srcid,
340            nodes: nodes.to_string(),
341            data: None,
342        }
343    }
344
345    /// Create a new StdErr message.
346    pub fn stderr(nodes: &str, srcid: u64) -> Self {
347        Message::StdErr {
348            msgid: next_msg_id(),
349            srcid,
350            nodes: nodes.to_string(),
351            data: None,
352        }
353    }
354
355    /// Create a new Retcode message.
356    pub fn retcode(nodes: &str, retcode: i32, srcid: u64) -> Self {
357        Message::Retcode {
358            msgid: next_msg_id(),
359            srcid,
360            retcode,
361            nodes: nodes.to_string(),
362        }
363    }
364
365    /// Create a new Timeout message.
366    pub fn timeout(nodes: &str, srcid: u64) -> Self {
367        Message::Timeout {
368            msgid: next_msg_id(),
369            srcid,
370            nodes: nodes.to_string(),
371        }
372    }
373
374    /// Create a new Routing message.
375    pub fn routing(event: &str, gateway: &str, targets: &str, srcid: u64) -> Self {
376        Message::Routing {
377            msgid: next_msg_id(),
378            srcid,
379            event: event.to_string(),
380            gateway: gateway.to_string(),
381            targets: targets.to_string(),
382        }
383    }
384
385    /// Create a Start message.
386    pub fn start() -> Self {
387        Message::Start {
388            msgid: next_msg_id(),
389        }
390    }
391
392    /// Create an End message.
393    pub fn end() -> Self {
394        Message::End {
395            msgid: next_msg_id(),
396        }
397    }
398
399    // -- Accessors -----------------------------------------------------------
400
401    /// Get the message type identifier string.
402    pub fn ident(&self) -> &str {
403        match self {
404            Message::General { .. } => IDENT_GEN,
405            Message::Configuration { .. } => IDENT_CFG,
406            Message::Control { .. } => IDENT_CTL,
407            Message::Ack { .. } => IDENT_ACK,
408            Message::Error { .. } => IDENT_ERR,
409            Message::StdOut { .. } => IDENT_OUT,
410            Message::StdErr { .. } => IDENT_SER,
411            Message::Retcode { .. } => IDENT_RET,
412            Message::Timeout { .. } => IDENT_TIM,
413            Message::Routing { .. } => IDENT_RTR,
414            Message::Start { .. } => IDENT_CHA,
415            Message::End { .. } => IDENT_END,
416        }
417    }
418
419    /// Get the message ID.
420    pub fn msgid(&self) -> u64 {
421        match self {
422            Message::General { msgid, .. }
423            | Message::Configuration { msgid, .. }
424            | Message::Control { msgid, .. }
425            | Message::Ack { msgid, .. }
426            | Message::Error { msgid, .. }
427            | Message::StdOut { msgid, .. }
428            | Message::StdErr { msgid, .. }
429            | Message::Retcode { msgid, .. }
430            | Message::Timeout { msgid, .. }
431            | Message::Routing { msgid, .. }
432            | Message::Start { msgid, .. }
433            | Message::End { msgid, .. } => *msgid,
434        }
435    }
436
437    /// Whether this message type carries a payload.
438    pub fn has_payload(&self) -> bool {
439        matches!(
440            self,
441            Message::General { .. }
442                | Message::Configuration { .. }
443                | Message::Control { .. }
444                | Message::StdOut { .. }
445                | Message::StdErr { .. }
446        )
447    }
448
449    /// Get a reference to the payload data, if any.
450    pub fn data(&self) -> Option<&[u8]> {
451        match self {
452            Message::General { data, .. }
453            | Message::Configuration { data, .. }
454            | Message::Control { data, .. }
455            | Message::StdOut { data, .. }
456            | Message::StdErr { data, .. } => data.as_deref(),
457            _ => None,
458        }
459    }
460
461    /// Set the payload data (raw bytes, will be base64-encoded on serialization).
462    pub fn set_data(&mut self, payload: Vec<u8>) {
463        match self {
464            Message::General { data, .. }
465            | Message::Configuration { data, .. }
466            | Message::Control { data, .. }
467            | Message::StdOut { data, .. }
468            | Message::StdErr { data, .. } => {
469                *data = Some(payload);
470            }
471            _ => {}
472        }
473    }
474
475    /// Encode arbitrary bytes as base64 and store as payload.
476    pub fn data_encode(&mut self, raw: &[u8]) {
477        let encoded = base64_encode(raw);
478        self.set_data(encoded.into_bytes());
479    }
480
481    /// Decode the base64 payload back to raw bytes.
482    pub fn data_decode(&self) -> Result<Vec<u8>, MessageProcessingError> {
483        match self.data() {
484            Some(d) => {
485                let s = std::str::from_utf8(d)
486                    .map_err(|_| MessageProcessingError::InvalidPayload(self.ident().into()))?;
487                base64_decode(s)
488            }
489            None => Err(MessageProcessingError::InvalidPayload(self.ident().into())),
490        }
491    }
492
493    /// Append data during incremental XML parsing.
494    pub fn data_update(&mut self, raw: &[u8]) -> Result<(), MessageProcessingError> {
495        if !self.has_payload() {
496            return Err(MessageProcessingError::UnexpectedPayload(
497                self.ident().into(),
498            ));
499        }
500        match self {
501            Message::General { data, .. }
502            | Message::Configuration { data, .. }
503            | Message::Control { data, .. }
504            | Message::StdOut { data, .. }
505            | Message::StdErr { data, .. } => {
506                if let Some(ref mut d) = data {
507                    d.extend_from_slice(raw);
508                } else {
509                    *data = Some(raw.to_vec());
510                }
511            }
512            _ => {}
513        }
514        Ok(())
515    }
516
517    /// Get the srcid for routed messages, None for non-routed.
518    pub fn srcid(&self) -> Option<u64> {
519        match self {
520            Message::Control { srcid, .. }
521            | Message::StdOut { srcid, .. }
522            | Message::StdErr { srcid, .. }
523            | Message::Retcode { srcid, .. }
524            | Message::Timeout { srcid, .. }
525            | Message::Routing { srcid, .. } => Some(*srcid),
526            _ => None,
527        }
528    }
529
530    /// Get the nodes field if present.
531    pub fn nodes(&self) -> Option<&str> {
532        match self {
533            Message::StdOut { nodes, .. }
534            | Message::StdErr { nodes, .. }
535            | Message::Retcode { nodes, .. }
536            | Message::Timeout { nodes, .. } => Some(nodes.as_str()),
537            _ => None,
538        }
539    }
540
541    // -- XML serialization ---------------------------------------------------
542
543    /// Serialize this message to XML bytes.
544    ///
545    /// Produces output like:
546    /// `<message type="CFG" msgid="0" gateway="gw1">base64payload</message>`
547    pub fn xml(&self) -> Vec<u8> {
548        let mut attrs = Vec::new();
549
550        // type and msgid are always present
551        attrs.push(format!("type=\"{}\"", xml_escape(self.ident())));
552        attrs.push(format!("msgid=\"{}\"", self.msgid()));
553
554        // variant-specific attributes
555        match self {
556            Message::Configuration { gateway, .. } => {
557                attrs.push(format!("gateway=\"{}\"", xml_escape(gateway)));
558            }
559            Message::Control {
560                srcid,
561                action,
562                target,
563                ..
564            } => {
565                attrs.push(format!("srcid=\"{}\"", srcid));
566                attrs.push(format!("action=\"{}\"", xml_escape(action)));
567                attrs.push(format!("target=\"{}\"", xml_escape(target)));
568            }
569            Message::Ack { ack, .. } => {
570                attrs.push(format!("ack=\"{}\"", ack));
571            }
572            Message::Error { reason, .. } => {
573                attrs.push(format!("reason=\"{}\"", xml_escape(reason)));
574            }
575            Message::StdOut { srcid, nodes, .. } => {
576                attrs.push(format!("srcid=\"{}\"", srcid));
577                attrs.push(format!("nodes=\"{}\"", xml_escape(nodes)));
578            }
579            Message::StdErr { srcid, nodes, .. } => {
580                attrs.push(format!("srcid=\"{}\"", srcid));
581                attrs.push(format!("nodes=\"{}\"", xml_escape(nodes)));
582            }
583            Message::Retcode {
584                srcid,
585                retcode,
586                nodes,
587                ..
588            } => {
589                attrs.push(format!("srcid=\"{}\"", srcid));
590                attrs.push(format!("retcode=\"{}\"", retcode));
591                attrs.push(format!("nodes=\"{}\"", xml_escape(nodes)));
592            }
593            Message::Timeout { srcid, nodes, .. } => {
594                attrs.push(format!("srcid=\"{}\"", srcid));
595                attrs.push(format!("nodes=\"{}\"", xml_escape(nodes)));
596            }
597            Message::Routing {
598                srcid,
599                event,
600                gateway,
601                targets,
602                ..
603            } => {
604                attrs.push(format!("srcid=\"{}\"", srcid));
605                attrs.push(format!("event=\"{}\"", xml_escape(event)));
606                attrs.push(format!("gateway=\"{}\"", xml_escape(gateway)));
607                attrs.push(format!("targets=\"{}\"", xml_escape(targets)));
608            }
609            _ => {}
610        }
611
612        let attr_str = attrs.join(" ");
613        let payload = self.data().unwrap_or(&[]);
614
615        if payload.is_empty() {
616            format!("<message {}/>", attr_str).into_bytes()
617        } else {
618            let payload_str = String::from_utf8_lossy(payload);
619            format!("<message {}>{}</message>", attr_str, payload_str).into_bytes()
620        }
621    }
622
623    /// Build a message from parsed XML attributes.
624    ///
625    /// The message type is selected based on the "type" attribute, and all
626    /// other attributes are read according to the variant's schema.
627    pub fn from_attrs(attrs: &HashMap<String, String>) -> Result<Message, MessageProcessingError> {
628        let msg_type = attrs
629            .get("type")
630            .ok_or_else(|| MessageProcessingError::MissingAttribute("type".into()))?;
631
632        let msgid: u64 = attrs
633            .get("msgid")
634            .and_then(|v| v.parse().ok())
635            .unwrap_or_else(next_msg_id);
636
637        match msg_type.as_str() {
638            IDENT_CFG => {
639                let gateway = attrs.get("gateway").cloned().unwrap_or_default();
640                Ok(Message::Configuration {
641                    msgid,
642                    gateway,
643                    data: None,
644                })
645            }
646            IDENT_CTL => {
647                let srcid = attr_u64(attrs, "srcid")?;
648                let action = attrs.get("action").cloned().unwrap_or_default();
649                let target = attrs.get("target").cloned().unwrap_or_default();
650                Ok(Message::Control {
651                    msgid,
652                    srcid,
653                    action,
654                    target,
655                    data: None,
656                })
657            }
658            IDENT_ACK => {
659                let ack = attr_u64(attrs, "ack")?;
660                Ok(Message::Ack { msgid, ack })
661            }
662            IDENT_ERR => {
663                let reason = attrs.get("reason").cloned().unwrap_or_default();
664                Ok(Message::Error { msgid, reason })
665            }
666            IDENT_OUT => {
667                let srcid = attr_u64(attrs, "srcid")?;
668                let nodes = attrs.get("nodes").cloned().unwrap_or_default();
669                Ok(Message::StdOut {
670                    msgid,
671                    srcid,
672                    nodes,
673                    data: None,
674                })
675            }
676            IDENT_SER => {
677                let srcid = attr_u64(attrs, "srcid")?;
678                let nodes = attrs.get("nodes").cloned().unwrap_or_default();
679                Ok(Message::StdErr {
680                    msgid,
681                    srcid,
682                    nodes,
683                    data: None,
684                })
685            }
686            IDENT_RET => {
687                let srcid = attr_u64(attrs, "srcid")?;
688                let retcode: i32 = attrs
689                    .get("retcode")
690                    .and_then(|v| v.parse().ok())
691                    .ok_or_else(|| MessageProcessingError::MissingAttribute("retcode".into()))?;
692                let nodes = attrs.get("nodes").cloned().unwrap_or_default();
693                Ok(Message::Retcode {
694                    msgid,
695                    srcid,
696                    retcode,
697                    nodes,
698                })
699            }
700            IDENT_TIM => {
701                let srcid = attr_u64(attrs, "srcid")?;
702                let nodes = attrs.get("nodes").cloned().unwrap_or_default();
703                Ok(Message::Timeout {
704                    msgid,
705                    srcid,
706                    nodes,
707                })
708            }
709            IDENT_RTR => {
710                let srcid = attr_u64(attrs, "srcid")?;
711                let event = attrs.get("event").cloned().unwrap_or_default();
712                let gateway = attrs.get("gateway").cloned().unwrap_or_default();
713                let targets = attrs.get("targets").cloned().unwrap_or_default();
714                Ok(Message::Routing {
715                    msgid,
716                    srcid,
717                    event,
718                    gateway,
719                    targets,
720                })
721            }
722            other => Err(MessageProcessingError::UnknownType(other.into())),
723        }
724    }
725}
726
727impl std::fmt::Display for Message {
728    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
729        write!(
730            f,
731            "Message {} (type: {}, msgid: {})",
732            self.ident(),
733            self.ident(),
734            self.msgid()
735        )
736    }
737}
738
739/// Helper: parse a u64 attribute.
740fn attr_u64(attrs: &HashMap<String, String>, key: &str) -> Result<u64, MessageProcessingError> {
741    attrs
742        .get(key)
743        .and_then(|v| v.parse().ok())
744        .ok_or_else(|| MessageProcessingError::MissingAttribute(key.into()))
745}
746
747// ============================================================================
748// XmlReader — incremental SAX-style XML parser
749// ============================================================================
750
751/// Parser states for the simple XML state machine.
752#[derive(Debug, Clone, Copy, PartialEq)]
753enum ParserState {
754    /// Outside any tag, looking for '<'.
755    Text,
756    /// Inside a tag name (after '<').
757    TagName,
758    /// Inside a closing tag (after '</').
759    CloseTagName,
760    /// Reading attributes inside an opening tag.
761    Attributes,
762    /// Inside a self-closing tag (seen '/').
763    SelfClose,
764}
765
766/// Incremental XML reader that parses fragments and produces [`Message`] instances.
767///
768/// Modeled after Python's SAX `ContentHandler`. Feed XML data incrementally
769/// via [`feed()`](XmlReader::feed) and retrieve parsed messages from the queue.
770#[derive(Debug)]
771pub struct XmlReader {
772    msg_queue: VecDeque<Message>,
773    /// Protocol version from `<channel version="...">`.
774    pub version: Option<String>,
775    /// Current message being assembled.
776    draft: Option<Message>,
777    // Parser internals
778    state: ParserState,
779    tag_buf: String,
780    attr_name: String,
781    attr_val: String,
782    attrs: HashMap<String, String>,
783    text_buf: String,
784    in_attr_value: bool,
785    attr_quote: char,
786}
787
788impl XmlReader {
789    /// Create a new empty XmlReader.
790    pub fn new() -> Self {
791        Self {
792            msg_queue: VecDeque::new(),
793            version: None,
794            draft: None,
795            state: ParserState::Text,
796            tag_buf: String::new(),
797            attr_name: String::new(),
798            attr_val: String::new(),
799            attrs: HashMap::new(),
800            text_buf: String::new(),
801            in_attr_value: false,
802            attr_quote: '"',
803        }
804    }
805
806    /// Whether a complete message is available.
807    pub fn msg_available(&self) -> bool {
808        !self.msg_queue.is_empty()
809    }
810
811    /// Pop and return the oldest message from the queue.
812    pub fn pop_msg(&mut self) -> Option<Message> {
813        self.msg_queue.pop_front()
814    }
815
816    /// Feed XML data to the parser.
817    ///
818    /// This is the incremental entry point. Each call may produce zero or more
819    /// messages that can be retrieved with [`pop_msg()`](XmlReader::pop_msg).
820    pub fn feed(&mut self, data: &str) -> Result<(), MessageProcessingError> {
821        for ch in data.chars() {
822            match self.state {
823                ParserState::Text => {
824                    if ch == '<' {
825                        // Flush any accumulated text as character data
826                        if !self.text_buf.is_empty() {
827                            let text = std::mem::take(&mut self.text_buf);
828                            self.characters(&text)?;
829                        }
830                        self.tag_buf.clear();
831                        self.state = ParserState::TagName;
832                    } else {
833                        self.text_buf.push(ch);
834                    }
835                }
836                ParserState::TagName => {
837                    if ch == '/' && self.tag_buf.is_empty() {
838                        // Closing tag: </...>
839                        self.tag_buf.clear();
840                        self.state = ParserState::CloseTagName;
841                    } else if ch == ' ' || ch == '\t' {
842                        // End of tag name, start reading attributes
843                        self.attrs.clear();
844                        self.attr_name.clear();
845                        self.state = ParserState::Attributes;
846                    } else if ch == '>' {
847                        // End of opening tag with no attributes
848                        self.attrs.clear();
849                        self.start_element(&self.tag_buf.clone(), &self.attrs.clone())?;
850                        self.state = ParserState::Text;
851                    } else if ch == '/' {
852                        self.attrs.clear();
853                        self.state = ParserState::SelfClose;
854                    } else {
855                        self.tag_buf.push(ch);
856                    }
857                }
858                ParserState::CloseTagName => {
859                    if ch == '>' {
860                        self.end_element(&self.tag_buf.clone())?;
861                        self.state = ParserState::Text;
862                    } else {
863                        self.tag_buf.push(ch);
864                    }
865                }
866                ParserState::Attributes => {
867                    if self.in_attr_value {
868                        if ch == self.attr_quote {
869                            // End of attribute value
870                            let name = std::mem::take(&mut self.attr_name);
871                            let val = xml_unescape(&std::mem::take(&mut self.attr_val));
872                            self.attrs.insert(name, val);
873                            self.in_attr_value = false;
874                        } else {
875                            self.attr_val.push(ch);
876                        }
877                    } else if ch == '=' {
878                        // About to read attribute value
879                    } else if ch == '"' || ch == '\'' {
880                        self.attr_quote = ch;
881                        self.attr_val.clear();
882                        self.in_attr_value = true;
883                    } else if ch == '>' {
884                        // End of opening tag
885                        self.start_element(&self.tag_buf.clone(), &self.attrs.clone())?;
886                        self.state = ParserState::Text;
887                    } else if ch == '/' {
888                        self.state = ParserState::SelfClose;
889                    } else if ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r' {
890                        // whitespace between attributes
891                    } else {
892                        self.attr_name.push(ch);
893                    }
894                }
895                ParserState::SelfClose => {
896                    if ch == '>' {
897                        // Self-closing tag: treat as open + close
898                        let tag = self.tag_buf.clone();
899                        let attrs = self.attrs.clone();
900                        self.start_element(&tag, &attrs)?;
901                        self.end_element(&tag)?;
902                        self.state = ParserState::Text;
903                    }
904                    // else: unexpected char, skip
905                }
906            }
907        }
908        Ok(())
909    }
910
911    fn start_element(
912        &mut self,
913        name: &str,
914        attrs: &HashMap<String, String>,
915    ) -> Result<(), MessageProcessingError> {
916        match name {
917            "channel" => {
918                self.version = attrs.get("version").cloned();
919                self.msg_queue.push_back(Message::start());
920            }
921            "message" => {
922                let msg = Message::from_attrs(attrs)?;
923                self.draft = Some(msg);
924            }
925            _ => {
926                return Err(MessageProcessingError::InvalidTag(name.into()));
927            }
928        }
929        Ok(())
930    }
931
932    fn end_element(&mut self, name: &str) -> Result<(), MessageProcessingError> {
933        match name {
934            "message" => {
935                if let Some(msg) = self.draft.take() {
936                    self.msg_queue.push_back(msg);
937                }
938            }
939            "channel" => {
940                self.msg_queue.push_back(Message::end());
941            }
942            _ => {}
943        }
944        Ok(())
945    }
946
947    fn characters(&mut self, content: &str) -> Result<(), MessageProcessingError> {
948        if let Some(ref mut draft) = self.draft {
949            let trimmed = content.trim();
950            if !trimmed.is_empty() {
951                draft.data_update(trimmed.as_bytes())?;
952            }
953        }
954        Ok(())
955    }
956}
957
958impl Default for XmlReader {
959    fn default() -> Self {
960        Self::new()
961    }
962}
963
964// ============================================================================
965// Channel
966// ============================================================================
967
968/// Trait for channel event handling (start/recv).
969///
970/// Implementors define the logic for initializing a channel and processing
971/// incoming messages.
972pub trait ChannelHandler: Send {
973    /// Called when the channel is connected and ready.
974    fn start(&mut self, channel: &mut Channel);
975
976    /// Called when a complete message has been received.
977    fn recv(&mut self, channel: &mut Channel, msg: Message);
978}
979
980/// Communication channel between nodes in the propagation tree.
981///
982/// A Channel manages XML-based message exchange over a transport (typically
983/// an SSH pipe). It uses an [`XmlReader`] for incoming messages and XML
984/// generation for outgoing messages.
985#[derive(Debug)]
986pub struct Channel {
987    /// Whether the channel has been opened (received StartMessage).
988    pub opened: bool,
989    /// Whether the channel setup is complete (received ACK to config).
990    pub setup: bool,
991    /// Whether this end initiated the channel.
992    pub initiator: bool,
993    /// The XML reader for parsing incoming data.
994    pub xml_reader: XmlReader,
995    /// Outgoing message buffer.
996    pub outbox: VecDeque<Vec<u8>>,
997}
998
999impl Channel {
1000    /// Create a new Channel.
1001    ///
1002    /// # Arguments
1003    /// * `initiator` - true if this end is initiating the channel
1004    pub fn new(initiator: bool) -> Self {
1005        Self {
1006            opened: false,
1007            setup: false,
1008            initiator,
1009            xml_reader: XmlReader::new(),
1010            outbox: VecDeque::new(),
1011        }
1012    }
1013
1014    /// Generate the XML document header.
1015    pub fn init_xml(&self) -> Vec<u8> {
1016        b"<?xml version='1.0' encoding='UTF-8'?>".to_vec()
1017    }
1018
1019    /// Generate the opening `<channel>` tag.
1020    pub fn open_xml(&self, version: &str) -> Vec<u8> {
1021        format!("<channel version=\"{}\">", xml_escape(version)).into_bytes()
1022    }
1023
1024    /// Generate the closing `</channel>` tag.
1025    pub fn close_xml(&self) -> Vec<u8> {
1026        b"</channel>".to_vec()
1027    }
1028
1029    /// Serialize a message and enqueue it for sending.
1030    pub fn send(&mut self, msg: &Message) {
1031        let mut xml = msg.xml();
1032        xml.push(b'\n');
1033        self.outbox.push_back(xml);
1034    }
1035
1036    /// Feed incoming data to the XML reader and return any parsed messages.
1037    pub fn feed(&mut self, data: &str) -> Result<Vec<Message>, MessageProcessingError> {
1038        self.xml_reader.feed(data)?;
1039        let mut messages = Vec::new();
1040        while let Some(msg) = self.xml_reader.pop_msg() {
1041            messages.push(msg);
1042        }
1043        Ok(messages)
1044    }
1045
1046    /// Take the next outgoing message from the outbox.
1047    pub fn take_outgoing(&mut self) -> Option<Vec<u8>> {
1048        self.outbox.pop_front()
1049    }
1050}
1051
1052// ============================================================================
1053// Tests
1054// ============================================================================
1055
1056#[cfg(test)]
1057mod tests {
1058    use super::*;
1059
1060    // -- Base64 tests --------------------------------------------------------
1061
1062    #[test]
1063    fn test_base64_roundtrip() {
1064        let data = b"Hello, ClusterShell!";
1065        let encoded = base64_encode(data);
1066        let decoded = base64_decode(&encoded).unwrap();
1067        assert_eq!(decoded, data);
1068    }
1069
1070    #[test]
1071    fn test_base64_empty() {
1072        assert_eq!(base64_encode(b""), "");
1073        assert_eq!(base64_decode("").unwrap(), Vec::<u8>::new());
1074    }
1075
1076    #[test]
1077    fn test_base64_padding() {
1078        // 1 byte -> 4 chars with ==
1079        let e1 = base64_encode(b"A");
1080        assert!(e1.ends_with("=="));
1081        assert_eq!(base64_decode(&e1).unwrap(), b"A");
1082
1083        // 2 bytes -> 4 chars with =
1084        let e2 = base64_encode(b"AB");
1085        assert!(e2.ends_with('=') && !e2.ends_with("=="));
1086        assert_eq!(base64_decode(&e2).unwrap(), b"AB");
1087
1088        // 3 bytes -> 4 chars, no padding
1089        let e3 = base64_encode(b"ABC");
1090        assert!(!e3.ends_with('='));
1091        assert_eq!(base64_decode(&e3).unwrap(), b"ABC");
1092    }
1093
1094    #[test]
1095    fn test_base64_with_newlines() {
1096        // base64 decoder should ignore whitespace
1097        let encoded = "SGVs\nbG8=";
1098        let decoded = base64_decode(encoded).unwrap();
1099        assert_eq!(decoded, b"Hello");
1100    }
1101
1102    // -- XML escape tests ----------------------------------------------------
1103
1104    #[test]
1105    fn test_xml_escape() {
1106        assert_eq!(xml_escape("hello"), "hello");
1107        assert_eq!(xml_escape("a&b"), "a&amp;b");
1108        assert_eq!(xml_escape("<tag>"), "&lt;tag&gt;");
1109        assert_eq!(xml_escape("x=\"y\""), "x=&quot;y&quot;");
1110    }
1111
1112    #[test]
1113    fn test_xml_unescape() {
1114        assert_eq!(xml_unescape("a&amp;b"), "a&b");
1115        assert_eq!(xml_unescape("&lt;tag&gt;"), "<tag>");
1116    }
1117
1118    // -- Message construction tests ------------------------------------------
1119
1120    #[test]
1121    fn test_message_idents() {
1122        assert_eq!(Message::general().ident(), "GEN");
1123        assert_eq!(Message::configuration("gw1").ident(), "CFG");
1124        assert_eq!(Message::control(0).ident(), "CTL");
1125        assert_eq!(Message::ack(0).ident(), "ACK");
1126        assert_eq!(Message::error("oops").ident(), "ERR");
1127        assert_eq!(Message::stdout("node1", 0).ident(), "OUT");
1128        assert_eq!(Message::stderr("node1", 0).ident(), "SER");
1129        assert_eq!(Message::retcode("node1", 0, 0).ident(), "RET");
1130        assert_eq!(Message::timeout("node1", 0).ident(), "TIM");
1131        assert_eq!(Message::routing("ev", "gw", "tgt", 0).ident(), "RTR");
1132        assert_eq!(Message::start().ident(), "CHA");
1133        assert_eq!(Message::end().ident(), "END");
1134    }
1135
1136    #[test]
1137    fn test_message_has_payload() {
1138        assert!(Message::general().has_payload());
1139        assert!(Message::configuration("gw1").has_payload());
1140        assert!(Message::control(0).has_payload());
1141        assert!(Message::stdout("n1", 0).has_payload());
1142        assert!(Message::stderr("n1", 0).has_payload());
1143        assert!(!Message::ack(0).has_payload());
1144        assert!(!Message::error("x").has_payload());
1145        assert!(!Message::retcode("n1", 0, 0).has_payload());
1146        assert!(!Message::timeout("n1", 0).has_payload());
1147        assert!(!Message::start().has_payload());
1148        assert!(!Message::end().has_payload());
1149    }
1150
1151    #[test]
1152    fn test_message_srcid() {
1153        assert_eq!(Message::control(42).srcid(), Some(42));
1154        assert_eq!(Message::stdout("n", 7).srcid(), Some(7));
1155        assert_eq!(Message::stderr("n", 8).srcid(), Some(8));
1156        assert_eq!(Message::retcode("n", 0, 9).srcid(), Some(9));
1157        assert_eq!(Message::timeout("n", 10).srcid(), Some(10));
1158        assert_eq!(Message::routing("e", "g", "t", 11).srcid(), Some(11));
1159        assert_eq!(Message::general().srcid(), None);
1160        assert_eq!(Message::ack(0).srcid(), None);
1161        assert_eq!(Message::configuration("gw").srcid(), None);
1162    }
1163
1164    #[test]
1165    fn test_message_nodes() {
1166        assert_eq!(Message::stdout("node[1-3]", 0).nodes(), Some("node[1-3]"));
1167        assert_eq!(Message::stderr("node5", 0).nodes(), Some("node5"));
1168        assert_eq!(Message::retcode("node1", 0, 0).nodes(), Some("node1"));
1169        assert_eq!(Message::timeout("node2", 0).nodes(), Some("node2"));
1170        assert_eq!(Message::general().nodes(), None);
1171        assert_eq!(Message::ack(0).nodes(), None);
1172    }
1173
1174    // -- Data encode/decode tests -------------------------------------------
1175
1176    #[test]
1177    fn test_data_encode_decode() {
1178        let mut msg = Message::configuration("gw1");
1179        let payload = b"some topology data here";
1180        msg.data_encode(payload);
1181        assert!(msg.data().is_some());
1182        let decoded = msg.data_decode().unwrap();
1183        assert_eq!(decoded, payload);
1184    }
1185
1186    #[test]
1187    fn test_data_update() {
1188        let mut msg = Message::configuration("gw");
1189        msg.data_update(b"SGVs").unwrap();
1190        msg.data_update(b"bG8=").unwrap();
1191        assert_eq!(msg.data().unwrap(), b"SGVsbG8=");
1192        let decoded = msg.data_decode().unwrap();
1193        assert_eq!(decoded, b"Hello");
1194    }
1195
1196    #[test]
1197    fn test_data_update_unexpected() {
1198        let mut msg = Message::ack(0);
1199        let result = msg.data_update(b"data");
1200        assert!(result.is_err());
1201        assert!(matches!(
1202            result.unwrap_err(),
1203            MessageProcessingError::UnexpectedPayload(_)
1204        ));
1205    }
1206
1207    // -- XML serialization tests --------------------------------------------
1208
1209    #[test]
1210    fn test_xml_ack() {
1211        let msg = Message::Ack { msgid: 5, ack: 3 };
1212        let xml = String::from_utf8(msg.xml()).unwrap();
1213        assert!(xml.contains("type=\"ACK\""));
1214        assert!(xml.contains("msgid=\"5\""));
1215        assert!(xml.contains("ack=\"3\""));
1216        assert!(xml.contains("/>") || xml.contains("</message>"));
1217    }
1218
1219    #[test]
1220    fn test_xml_configuration_with_payload() {
1221        let mut msg = Message::Configuration {
1222            msgid: 0,
1223            gateway: "gateway1".into(),
1224            data: None,
1225        };
1226        msg.data_encode(b"topology");
1227        let xml = String::from_utf8(msg.xml()).unwrap();
1228        assert!(xml.contains("type=\"CFG\""));
1229        assert!(xml.contains("gateway=\"gateway1\""));
1230        assert!(xml.contains("</message>"));
1231        // Should have base64 content between tags
1232        assert!(xml.contains(&base64_encode(b"topology")));
1233    }
1234
1235    #[test]
1236    fn test_xml_error() {
1237        let msg = Message::Error {
1238            msgid: 1,
1239            reason: "something broke".into(),
1240        };
1241        let xml = String::from_utf8(msg.xml()).unwrap();
1242        assert!(xml.contains("type=\"ERR\""));
1243        assert!(xml.contains("reason=\"something broke\""));
1244    }
1245
1246    #[test]
1247    fn test_xml_retcode() {
1248        let msg = Message::Retcode {
1249            msgid: 2,
1250            srcid: 10,
1251            retcode: -1,
1252            nodes: "node[1-5]".into(),
1253        };
1254        let xml = String::from_utf8(msg.xml()).unwrap();
1255        assert!(xml.contains("type=\"RET\""));
1256        assert!(xml.contains("retcode=\"-1\""));
1257        assert!(xml.contains("nodes=\"node[1-5]\""));
1258    }
1259
1260    #[test]
1261    fn test_xml_control() {
1262        let msg = Message::Control {
1263            msgid: 3,
1264            srcid: 1,
1265            action: "shell".into(),
1266            target: "node[1-10]".into(),
1267            data: None,
1268        };
1269        let xml = String::from_utf8(msg.xml()).unwrap();
1270        assert!(xml.contains("type=\"CTL\""));
1271        assert!(xml.contains("action=\"shell\""));
1272        assert!(xml.contains("target=\"node[1-10]\""));
1273    }
1274
1275    #[test]
1276    fn test_xml_routing() {
1277        let msg = Message::Routing {
1278            msgid: 4,
1279            srcid: 2,
1280            event: "connect".into(),
1281            gateway: "gw2".into(),
1282            targets: "node[20-30]".into(),
1283        };
1284        let xml = String::from_utf8(msg.xml()).unwrap();
1285        assert!(xml.contains("type=\"RTR\""));
1286        assert!(xml.contains("event=\"connect\""));
1287        assert!(xml.contains("gateway=\"gw2\""));
1288        assert!(xml.contains("targets=\"node[20-30]\""));
1289    }
1290
1291    // -- XmlReader tests ----------------------------------------------------
1292
1293    #[test]
1294    fn test_xmlreader_channel_start_end() {
1295        let mut reader = XmlReader::new();
1296        reader.feed("<channel version=\"1.0\"></channel>").unwrap();
1297        assert!(reader.msg_available());
1298        let m1 = reader.pop_msg().unwrap();
1299        assert_eq!(m1.ident(), "CHA");
1300        assert_eq!(reader.version.as_deref(), Some("1.0"));
1301
1302        let m2 = reader.pop_msg().unwrap();
1303        assert_eq!(m2.ident(), "END");
1304        assert!(!reader.msg_available());
1305    }
1306
1307    #[test]
1308    fn test_xmlreader_ack_message() {
1309        let mut reader = XmlReader::new();
1310        reader
1311            .feed(
1312                "<channel version=\"1.0\"><message type=\"ACK\" msgid=\"0\" ack=\"5\"/></channel>",
1313            )
1314            .unwrap();
1315
1316        let start = reader.pop_msg().unwrap();
1317        assert_eq!(start.ident(), "CHA");
1318
1319        let ack = reader.pop_msg().unwrap();
1320        assert_eq!(ack.ident(), "ACK");
1321        if let Message::Ack { ack, msgid, .. } = ack {
1322            assert_eq!(ack, 5);
1323            assert_eq!(msgid, 0);
1324        } else {
1325            panic!("expected Ack variant");
1326        }
1327
1328        let end = reader.pop_msg().unwrap();
1329        assert_eq!(end.ident(), "END");
1330    }
1331
1332    #[test]
1333    fn test_xmlreader_message_with_payload() {
1334        let payload_b64 = base64_encode(b"hello world");
1335        let xml = format!(
1336            "<channel version=\"2.0\"><message type=\"CFG\" msgid=\"1\" gateway=\"gw1\">{}</message></channel>",
1337            payload_b64
1338        );
1339
1340        let mut reader = XmlReader::new();
1341        reader.feed(&xml).unwrap();
1342
1343        let _start = reader.pop_msg().unwrap();
1344        let cfg = reader.pop_msg().unwrap();
1345        assert_eq!(cfg.ident(), "CFG");
1346        let decoded = cfg.data_decode().unwrap();
1347        assert_eq!(decoded, b"hello world");
1348
1349        if let Message::Configuration { gateway, .. } = &cfg {
1350            assert_eq!(gateway, "gw1");
1351        } else {
1352            panic!("expected Configuration variant");
1353        }
1354    }
1355
1356    #[test]
1357    fn test_xmlreader_incremental_feed() {
1358        let mut reader = XmlReader::new();
1359
1360        // Feed in chunks
1361        reader.feed("<chan").unwrap();
1362        assert!(!reader.msg_available());
1363        reader.feed("nel version=\"1.0\">").unwrap();
1364        assert!(reader.msg_available());
1365        let start = reader.pop_msg().unwrap();
1366        assert_eq!(start.ident(), "CHA");
1367
1368        reader
1369            .feed("<message type=\"ACK\" msgid=\"0\" ack=\"1\"/>")
1370            .unwrap();
1371        let ack = reader.pop_msg().unwrap();
1372        assert_eq!(ack.ident(), "ACK");
1373
1374        reader.feed("</channel>").unwrap();
1375        let end = reader.pop_msg().unwrap();
1376        assert_eq!(end.ident(), "END");
1377    }
1378
1379    #[test]
1380    fn test_xmlreader_unknown_type() {
1381        let mut reader = XmlReader::new();
1382        let result =
1383            reader.feed("<channel version=\"1\"><message type=\"BOGUS\" msgid=\"0\"/></channel>");
1384        assert!(result.is_err());
1385    }
1386
1387    #[test]
1388    fn test_xmlreader_error_message() {
1389        let mut reader = XmlReader::new();
1390        reader
1391            .feed("<channel version=\"1\"><message type=\"ERR\" msgid=\"0\" reason=\"parse failed\"/></channel>")
1392            .unwrap();
1393
1394        let _start = reader.pop_msg().unwrap();
1395        let err = reader.pop_msg().unwrap();
1396        assert_eq!(err.ident(), "ERR");
1397        if let Message::Error { reason, .. } = &err {
1398            assert_eq!(reason, "parse failed");
1399        } else {
1400            panic!("expected Error variant");
1401        }
1402    }
1403
1404    #[test]
1405    fn test_xmlreader_retcode() {
1406        let mut reader = XmlReader::new();
1407        reader
1408            .feed(
1409                "<message type=\"RET\" msgid=\"7\" srcid=\"3\" retcode=\"0\" nodes=\"node[1-5]\"/>",
1410            )
1411            .unwrap();
1412
1413        let msg = reader.pop_msg().unwrap();
1414        if let Message::Retcode {
1415            srcid,
1416            retcode,
1417            nodes,
1418            ..
1419        } = &msg
1420        {
1421            assert_eq!(*srcid, 3);
1422            assert_eq!(*retcode, 0);
1423            assert_eq!(nodes, "node[1-5]");
1424        } else {
1425            panic!("expected Retcode variant");
1426        }
1427    }
1428
1429    #[test]
1430    fn test_xmlreader_routing() {
1431        let mut reader = XmlReader::new();
1432        reader
1433            .feed("<message type=\"RTR\" msgid=\"0\" srcid=\"1\" event=\"connect\" gateway=\"gw1\" targets=\"node[1-10]\"/>")
1434            .unwrap();
1435
1436        let msg = reader.pop_msg().unwrap();
1437        if let Message::Routing {
1438            event,
1439            gateway,
1440            targets,
1441            ..
1442        } = &msg
1443        {
1444            assert_eq!(event, "connect");
1445            assert_eq!(gateway, "gw1");
1446            assert_eq!(targets, "node[1-10]");
1447        } else {
1448            panic!("expected Routing variant");
1449        }
1450    }
1451
1452    #[test]
1453    fn test_xmlreader_stdout_stderr() {
1454        let payload = base64_encode(b"output line");
1455        let xml = format!(
1456            "<message type=\"OUT\" msgid=\"0\" srcid=\"5\" nodes=\"node1\">{}</message>",
1457            payload
1458        );
1459        let mut reader = XmlReader::new();
1460        reader.feed(&xml).unwrap();
1461
1462        let msg = reader.pop_msg().unwrap();
1463        assert_eq!(msg.ident(), "OUT");
1464        assert_eq!(msg.srcid(), Some(5));
1465        assert_eq!(msg.nodes(), Some("node1"));
1466        assert_eq!(msg.data_decode().unwrap(), b"output line");
1467
1468        // StdErr
1469        let xml2 = format!(
1470            "<message type=\"SER\" msgid=\"1\" srcid=\"6\" nodes=\"node2\">{}</message>",
1471            base64_encode(b"error line")
1472        );
1473        reader.feed(&xml2).unwrap();
1474        let msg2 = reader.pop_msg().unwrap();
1475        assert_eq!(msg2.ident(), "SER");
1476        assert_eq!(msg2.data_decode().unwrap(), b"error line");
1477    }
1478
1479    #[test]
1480    fn test_xmlreader_timeout() {
1481        let mut reader = XmlReader::new();
1482        reader
1483            .feed("<message type=\"TIM\" msgid=\"0\" srcid=\"1\" nodes=\"node[5-8]\"/>")
1484            .unwrap();
1485        let msg = reader.pop_msg().unwrap();
1486        assert_eq!(msg.ident(), "TIM");
1487        assert_eq!(msg.nodes(), Some("node[5-8]"));
1488    }
1489
1490    #[test]
1491    fn test_xmlreader_control() {
1492        let payload = base64_encode(b"command data");
1493        let xml = format!(
1494            "<message type=\"CTL\" msgid=\"0\" srcid=\"1\" action=\"shell\" target=\"node[1-3]\">{}</message>",
1495            payload
1496        );
1497        let mut reader = XmlReader::new();
1498        reader.feed(&xml).unwrap();
1499
1500        let msg = reader.pop_msg().unwrap();
1501        if let Message::Control { action, target, .. } = &msg {
1502            assert_eq!(action, "shell");
1503            assert_eq!(target, "node[1-3]");
1504        } else {
1505            panic!("expected Control variant");
1506        }
1507        assert_eq!(msg.data_decode().unwrap(), b"command data");
1508    }
1509
1510    // -- XML roundtrip tests ------------------------------------------------
1511
1512    #[test]
1513    fn test_xml_roundtrip_ack() {
1514        let original = Message::Ack { msgid: 10, ack: 7 };
1515        let xml = String::from_utf8(original.xml()).unwrap();
1516        let mut reader = XmlReader::new();
1517        reader.feed(&xml).unwrap();
1518        let parsed = reader.pop_msg().unwrap();
1519        assert_eq!(parsed.ident(), "ACK");
1520        if let Message::Ack { msgid, ack, .. } = parsed {
1521            assert_eq!(msgid, 10);
1522            assert_eq!(ack, 7);
1523        }
1524    }
1525
1526    #[test]
1527    fn test_xml_roundtrip_cfg_with_data() {
1528        let mut original = Message::Configuration {
1529            msgid: 0,
1530            gateway: "my-gateway".into(),
1531            data: None,
1532        };
1533        original.data_encode(b"topology bytes here");
1534        let xml = String::from_utf8(original.xml()).unwrap();
1535
1536        let mut reader = XmlReader::new();
1537        reader.feed(&xml).unwrap();
1538        let parsed = reader.pop_msg().unwrap();
1539        assert_eq!(parsed.ident(), "CFG");
1540        if let Message::Configuration { gateway, .. } = &parsed {
1541            assert_eq!(gateway, "my-gateway");
1542        }
1543        assert_eq!(parsed.data_decode().unwrap(), b"topology bytes here");
1544    }
1545
1546    #[test]
1547    fn test_xml_roundtrip_error() {
1548        let original = Message::Error {
1549            msgid: 3,
1550            reason: "test <error> & reason".into(),
1551        };
1552        let xml = String::from_utf8(original.xml()).unwrap();
1553        let mut reader = XmlReader::new();
1554        reader.feed(&xml).unwrap();
1555        let parsed = reader.pop_msg().unwrap();
1556        if let Message::Error { reason, .. } = &parsed {
1557            assert_eq!(reason, "test <error> & reason");
1558        } else {
1559            panic!("expected Error variant");
1560        }
1561    }
1562
1563    #[test]
1564    fn test_xml_roundtrip_retcode() {
1565        let original = Message::Retcode {
1566            msgid: 5,
1567            srcid: 2,
1568            retcode: 127,
1569            nodes: "node[1-100]".into(),
1570        };
1571        let xml = String::from_utf8(original.xml()).unwrap();
1572        let mut reader = XmlReader::new();
1573        reader.feed(&xml).unwrap();
1574        let parsed = reader.pop_msg().unwrap();
1575        if let Message::Retcode { retcode, nodes, .. } = &parsed {
1576            assert_eq!(*retcode, 127);
1577            assert_eq!(nodes, "node[1-100]");
1578        } else {
1579            panic!("expected Retcode variant");
1580        }
1581    }
1582
1583    // -- Channel tests ------------------------------------------------------
1584
1585    #[test]
1586    fn test_channel_creation() {
1587        let ch = Channel::new(true);
1588        assert!(!ch.opened);
1589        assert!(!ch.setup);
1590        assert!(ch.initiator);
1591    }
1592
1593    #[test]
1594    fn test_channel_init_xml() {
1595        let ch = Channel::new(false);
1596        let xml = ch.init_xml();
1597        assert_eq!(
1598            String::from_utf8(xml).unwrap(),
1599            "<?xml version='1.0' encoding='UTF-8'?>"
1600        );
1601    }
1602
1603    #[test]
1604    fn test_channel_open_close_xml() {
1605        let ch = Channel::new(true);
1606        let open = String::from_utf8(ch.open_xml("3.0")).unwrap();
1607        assert_eq!(open, "<channel version=\"3.0\">");
1608        let close = String::from_utf8(ch.close_xml()).unwrap();
1609        assert_eq!(close, "</channel>");
1610    }
1611
1612    #[test]
1613    fn test_channel_send() {
1614        let mut ch = Channel::new(true);
1615        let msg = Message::ack(42);
1616        ch.send(&msg);
1617        let outgoing = ch.take_outgoing().unwrap();
1618        let xml = String::from_utf8(outgoing).unwrap();
1619        assert!(xml.contains("type=\"ACK\""));
1620        assert!(xml.ends_with('\n'));
1621    }
1622
1623    #[test]
1624    fn test_channel_feed_and_receive() {
1625        let mut ch = Channel::new(false);
1626        let messages = ch
1627            .feed(
1628                "<channel version=\"1.0\"><message type=\"ACK\" msgid=\"0\" ack=\"1\"/></channel>",
1629            )
1630            .unwrap();
1631        assert_eq!(messages.len(), 3); // Start, ACK, End
1632        assert_eq!(messages[0].ident(), "CHA");
1633        assert_eq!(messages[1].ident(), "ACK");
1634        assert_eq!(messages[2].ident(), "END");
1635    }
1636
1637    #[test]
1638    fn test_channel_outbox_ordering() {
1639        let mut ch = Channel::new(true);
1640        ch.send(&Message::ack(1));
1641        ch.send(&Message::ack(2));
1642        ch.send(&Message::ack(3));
1643
1644        let m1 = ch.take_outgoing().unwrap();
1645        let m2 = ch.take_outgoing().unwrap();
1646        let m3 = ch.take_outgoing().unwrap();
1647        assert!(ch.take_outgoing().is_none());
1648
1649        // They should come out in order
1650        assert!(String::from_utf8(m1).unwrap().contains("ACK"));
1651        assert!(String::from_utf8(m2).unwrap().contains("ACK"));
1652        assert!(String::from_utf8(m3).unwrap().contains("ACK"));
1653    }
1654
1655    // -- Message Display test -----------------------------------------------
1656
1657    #[test]
1658    fn test_message_display() {
1659        let msg = Message::ack(0);
1660        let display = format!("{}", msg);
1661        assert!(display.contains("ACK"));
1662    }
1663
1664    // -- Edge case tests ----------------------------------------------------
1665
1666    #[test]
1667    fn test_from_attrs_missing_type() {
1668        let attrs = HashMap::new();
1669        assert!(Message::from_attrs(&attrs).is_err());
1670    }
1671
1672    #[test]
1673    fn test_from_attrs_missing_srcid() {
1674        let mut attrs = HashMap::new();
1675        attrs.insert("type".into(), "CTL".into());
1676        attrs.insert("msgid".into(), "0".into());
1677        // missing srcid
1678        assert!(Message::from_attrs(&attrs).is_err());
1679    }
1680
1681    #[test]
1682    fn test_data_decode_no_data() {
1683        let msg = Message::configuration("gw");
1684        let result = msg.data_decode();
1685        assert!(result.is_err());
1686    }
1687
1688    #[test]
1689    fn test_xml_special_characters_roundtrip() {
1690        let msg = Message::Error {
1691            msgid: 0,
1692            reason: "node<1>&\"2\"".into(),
1693        };
1694        let xml = String::from_utf8(msg.xml()).unwrap();
1695        // Should be escaped in XML
1696        assert!(xml.contains("&lt;"));
1697        assert!(xml.contains("&amp;"));
1698
1699        // Parse it back
1700        let mut reader = XmlReader::new();
1701        reader.feed(&xml).unwrap();
1702        let parsed = reader.pop_msg().unwrap();
1703        if let Message::Error { reason, .. } = &parsed {
1704            assert_eq!(reason, "node<1>&\"2\"");
1705        }
1706    }
1707}