netconf_rust/

codec.rs

use bytes::{Buf, Bytes, BytesMut};
use log::{debug, trace};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio_util::codec::{Decoder, Encoder};

use crate::error::FramingError;

const EOM_MARKER: &[u8] = b"]]>]]>";
const EOM_LEN: usize = EOM_MARKER.len();

const CHUNKED_EOM_MARKER: &[u8] = b"\n##\n";
const CHUNKED_EOM_MARKER_LEN: usize = CHUNKED_EOM_MARKER.len();

const CHUNKED_HEADER_START: &[u8] = b"\n#";

/// A decoded frame from the NETCONF codec.
///
/// Instead of returning complete messages, the decoder yields individual
/// chunks as they arrive. This enables streaming: the reader task can
/// forward chunks directly to consumers without buffering the entire
/// message.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DecodedFrame {
    /// A piece of the current NETCONF message.
    Chunk(Bytes),
    /// The current message is complete.
    EndOfMessage,
}

/// Netconf defines 2 framing modes.
/// 1. Netconf 1.0 defines End of Message (EOM) in RFC 4742, where each message is followed by
///    the literal sequence `]]>]]>`. For example, `<rpc-reply message-id="1"><ok/></rpc-reply>]]>]]>`
/// 2. Netconf 1.1 defines Chunked in RFC 6242, where each message is sent with length-prefixed
///    chunks. The chunk starts with \n#{num_of_bytes_in_msg}\n. \n##\n defines end of message
///    for example `\n#28\n<rpc-reply><ok/></rpc-re\n#6\nply/>\n##\n`
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FramingMode {
    /// NETCONF 1.0 (RFC 4742): messages terminated by `]]>]]`
    EndOfMessage,
    /// NETCONF 1.1 (RFC 6242): length-prefixed chunked framing
    Chunked,
}

#[derive(Default, Debug, Clone, Copy)]
pub struct CodecConfig {
    pub max_message_size: Option<usize>, // None = unlimited
}

pub struct NetconfCodec {
    framing_mode: FramingMode,
    config: CodecConfig,
    /// Tracks whether we need to yield `EndOfMessage` on the next `decode()`
    /// call. In EOM mode, when the `]]>]]>` marker is found, the final data
    /// is yielded as a `Chunk` first, and `EndOfMessage` follows on the next
    /// call. This two-step yield is needed because `decode()` can only return
    /// one item at a time.
    eom_complete: bool,
    /// Cumulative bytes yielded for the current message. Reset on
    /// `EndOfMessage`. Used to enforce `max_message_size` in both EOM
    /// and chunked modes.
    message_bytes: usize,
    /// When true, `decode_eof` silently discards leftover bytes instead of
    /// returning an error. Set by the reader loop after a `close-session` RPC
    /// is in flight so that harmless holdback bytes don't surface as a
    /// transport error.
    closing: bool,
}

/// tokio_util codec for NETCONF message framing.
///
/// Implements the `Decoder` trait from tokio_util. Instead of returning
/// complete messages, the decoder yields [`DecodedFrame`] items:
/// individual `Chunk`s as data arrives, followed by `EndOfMessage` when
/// the message boundary is reached.
///
/// This chunk-level output enables the reader task to route messages
/// without buffering: normal RPCs accumulate chunks into a `BytesMut`,
/// while streaming RPCs forward each chunk directly to the consumer.
///
/// The `Encoder` side is unchanged — it frames complete XML documents
/// with the appropriate delimiter/chunking.
impl NetconfCodec {
    pub fn new(framing_mode: FramingMode, config: CodecConfig) -> Self {
        Self {
            framing_mode,
            config,
            eom_complete: false,
            message_bytes: 0,
            closing: false,
        }
    }

    pub fn set_mode(&mut self, framing_mode: FramingMode) {
        self.framing_mode = framing_mode;
        self.eom_complete = false;
        self.message_bytes = 0;
    }
    pub fn framing_mode(&self) -> FramingMode {
        self.framing_mode
    }
    /// Mark the codec as closing so that `decode_eof` silently discards
    /// leftover holdback bytes instead of treating them as an error.
    pub fn set_closing(&mut self) {
        self.closing = true;
    }

    fn check_size(&self, size: usize) -> Result<(), FramingError> {
        if let Some(max_size) = self.config.max_message_size
            && size > max_size
        {
            return Err(FramingError::MessageTooLarge {
                limit: max_size,
                received: size,
            });
        }
        Ok(())
    }

    /// Decode EOM-framed data into chunks.
    ///
    /// Yields available data as `Chunk` frames, holding back `EOM_LEN - 1`
    /// bytes to handle `]]>]]>` markers split across reads. When the marker
    /// is found, the final data before it is yielded as a `Chunk`, and
    /// `EndOfMessage` follows on the next `decode()` call via the
    /// `eom_complete` flag.
    fn decode_eom(&mut self, src: &mut BytesMut) -> Result<Option<DecodedFrame>, FramingError> {
        // Two-step yield: if we just found the EOM marker and yielded
        // the final Chunk, now yield EndOfMessage.
        if self.eom_complete {
            self.eom_complete = false;
            self.message_bytes = 0;
            return Ok(Some(DecodedFrame::EndOfMessage));
        }

        if src.is_empty() {
            return Ok(None);
        }

        // Search for the EOM marker.
        if let Some(pos) = memchr::memmem::find(src, EOM_MARKER) {
            if pos > 0 {
                // Yield data before the marker as a Chunk, schedule EndOfMessage.
                let data = src.split_to(pos);
                src.advance(EOM_LEN);
                self.message_bytes += data.len();
                self.check_size(self.message_bytes)?;
                self.eom_complete = true;
                debug!("eom: final chunk ({} bytes), marker found", data.len());
                Ok(Some(DecodedFrame::Chunk(data.freeze())))
            } else {
                // Marker at the start — empty message.
                src.advance(EOM_LEN);
                self.message_bytes = 0;
                debug!("eom: empty message");
                Ok(Some(DecodedFrame::EndOfMessage))
            }
        } else {
            // No marker found. Yield everything except the holdback
            // (EOM_LEN - 1 bytes) since the marker could be split across reads.
            let holdback = EOM_LEN - 1;
            if src.len() <= holdback {
                return Ok(None);
            }
            let safe_len = src.len() - holdback;
            let chunk = src.split_to(safe_len);
            self.message_bytes += chunk.len();
            self.check_size(self.message_bytes)?;
            trace!(
                "eom: yielding chunk ({} bytes), holding back {}",
                chunk.len(),
                holdback
            );
            Ok(Some(DecodedFrame::Chunk(chunk.freeze())))
        }
    }

    /// Decode chunked-framed data into individual chunks.
    ///
    /// Each RFC 6242 chunk `\n#<size>\n<data>` yields a `Chunk(data)`.
    /// The end marker `\n##\n` yields `EndOfMessage`. Unlike the previous
    /// implementation, chunks are NOT accumulated — each one is yielded
    /// individually, enabling streaming.
    fn decode_chunked(&mut self, src: &mut BytesMut) -> Result<Option<DecodedFrame>, FramingError> {
        // The smallest possible frame is \n##\n (4 bytes).
        if src.len() < CHUNKED_EOM_MARKER_LEN {
            trace!(
                "chunked: buffer too small ({} bytes), need more data",
                src.len()
            );
            return Ok(None);
        }

        // Every chunk or end marker starts with \n#
        if src[0..2] != *CHUNKED_HEADER_START {
            return Err(FramingError::InvalidHeader {
                expected: "\\n#",
                got: src[..2].to_vec(),
            });
        }

        // Check for end of chunks marker: \n##\n.
        if src[2] == b'#' {
            if src[3] != b'\n' {
                return Err(FramingError::InvalidHeader {
                    expected: "\\n##\\n",
                    got: src[..4].to_vec(),
                });
            }
            src.advance(CHUNKED_EOM_MARKER_LEN);
            self.message_bytes = 0;
            debug!("chunked: end of message");
            return Ok(Some(DecodedFrame::EndOfMessage));
        }

        // Parse chunk header \n#<size>\n
        let header_start = 2; // skip \n#
        let header_end = match src[header_start..].iter().position(|&b| b == b'\n') {
            Some(pos_end_of_header) => header_start + pos_end_of_header,
            None => {
                // Header not yet complete - need more data.
                // Sanity-check the header length.
                if src.len() > 20 {
                    return Err(FramingError::InvalidChunkSize(
                        String::from_utf8_lossy(&src[header_start..]).into_owned(),
                    ));
                }
                return Ok(None);
            }
        };

        // Extract the chunk size from the header and parse into usize
        let size_str = &src[header_start..header_end];
        let chunk_size: usize = std::str::from_utf8(size_str)
            .map_err(|_| {
                FramingError::InvalidChunkSize(String::from_utf8_lossy(size_str).into_owned())
            })?
            .parse()
            .map_err(|_| {
                FramingError::InvalidChunkSize(String::from_utf8_lossy(size_str).into_owned())
            })?;

        if chunk_size == 0 {
            return Err(FramingError::InvalidChunkSize("0".into()));
        }

        // Total header length: \n# + size digits + \n
        let header_len = header_end + 1; // +1 for the trailing \n

        // Check if the full chunk (header + data) is available
        let total_chunk_len = header_len + chunk_size;
        if src.len() < total_chunk_len {
            trace!(
                "chunked: need {} more bytes for chunk (have {}, need {})",
                total_chunk_len - src.len(),
                src.len(),
                total_chunk_len
            );
            return Ok(None);
        }
        self.message_bytes += chunk_size;
        self.check_size(self.message_bytes)?;

        trace!("chunked: yielding chunk ({} bytes)", chunk_size);

        // Consume header
        src.advance(header_len);

        // Extract and yield the chunk data
        let chunk_data = src.split_to(chunk_size);
        Ok(Some(DecodedFrame::Chunk(chunk_data.freeze())))
    }
}

pub(crate) async fn read_eom_message<R: AsyncRead + Unpin>(
    reader: &mut R,
    max_size: Option<usize>,
) -> crate::Result<String> {
    let mut buf = Vec::with_capacity(4096);
    let mut tmp = [0u8; 4096];

    loop {
        let read_bytes = reader.read(&mut tmp).await?;

        if read_bytes == 0 {
            debug!("read_eom: unexpected EOF after {} bytes", buf.len());
            return Err(FramingError::UnexpectedEof.into());
        }
        buf.extend_from_slice(&tmp[..read_bytes]);
        trace!(
            "read_eom: read {} bytes, buffer now {} bytes",
            read_bytes,
            buf.len()
        );
        if let Some(limit) = max_size
            && buf.len() > limit + EOM_LEN
        {
            return Err(FramingError::MessageTooLarge {
                limit,
                received: buf.len(),
            }
            .into());
        }
        if let Some(pos) = memchr::memmem::find(&buf, EOM_MARKER) {
            buf.truncate(pos);
            debug!("read_eom: complete message ({} bytes)", buf.len());
            return String::from_utf8(buf).map_err(|_| FramingError::InvalidUtf8.into());
        }
    }
}

pub(crate) async fn write_eom_message<W: AsyncWrite + Unpin>(
    writer: &mut W,
    message: &str,
) -> crate::Result<()> {
    writer.write_all(message.as_bytes()).await?;
    writer.write_all(EOM_MARKER).await?;
    writer.flush().await?;
    Ok(())
}

/// Extract the `message-id` attribute value from a NETCONF XML fragment.
///
/// Performs a byte-level scan for `message-id="<digits>"` and returns the
/// parsed `u32`. This avoids pulling in an XML parser for the hot path
/// in the reader task, where we only need the message-id from the first
/// chunk to route the message.
pub(crate) fn extract_message_id_from_bytes(bytes: &[u8]) -> Option<u32> {
    const PATTERN: &[u8] = b"message-id=\"";
    let pos = memchr::memmem::find(bytes, PATTERN)?;
    let start = pos + PATTERN.len();
    let remaining = bytes.get(start..)?;
    let end = memchr::memchr(b'"', remaining)?;
    let id_bytes = &remaining[..end];
    std::str::from_utf8(id_bytes).ok()?.parse().ok()
}

/// <https://docs.rs/tokio-util/latest/tokio_util/codec/trait.Decoder.html>
impl Decoder for NetconfCodec {
    type Item = DecodedFrame;
    type Error = FramingError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        match self.framing_mode {
            FramingMode::EndOfMessage => self.decode_eom(src),
            FramingMode::Chunked => self.decode_chunked(src),
        }
    }

    /// Handle EOF from the underlying transport.
    ///
    /// When `closing` is set (a `close-session` RPC is in flight), leftover
    /// holdback bytes are harmless artifacts of the EOM split-marker guard
    /// and are silently discarded. Otherwise we fall back to the default
    /// behaviour: any remaining bytes at EOF indicate a transport problem.
    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        match self.decode(buf)? {
            Some(frame) => Ok(Some(frame)),
            None => {
                if !buf.is_empty() {
                    if self.closing {
                        debug!(
                            "decode_eof: discarding {} leftover bytes (session closing)",
                            buf.len()
                        );
                        buf.clear();
                    } else {
                        return Err(FramingError::Io(std::io::Error::other(format!(
                            "bytes remaining on stream ({} bytes)",
                            buf.len()
                        ))));
                    }
                }
                Ok(None)
            }
        }
    }
}

impl Encoder<Bytes> for NetconfCodec {
    type Error = FramingError;

    fn encode(&mut self, item: Bytes, dst: &mut BytesMut) -> Result<(), Self::Error> {
        debug!(
            "encode: framing={:?}, message={} bytes",
            self.framing_mode,
            item.len()
        );
        trace!(
            "encode: message preview: {:?}",
            String::from_utf8_lossy(&item[..item.len().min(200)])
        );
        match self.framing_mode {
            FramingMode::EndOfMessage => {
                dst.reserve(item.len() + EOM_LEN);
                dst.extend_from_slice(&item);
                dst.extend_from_slice(EOM_MARKER);
            }
            FramingMode::Chunked => {
                let header = format!("\n#{}\n", item.len());
                dst.reserve(header.len() + item.len() + CHUNKED_EOM_MARKER_LEN);
                dst.extend_from_slice(header.as_bytes());
                dst.extend_from_slice(&item);
                dst.extend_from_slice(CHUNKED_EOM_MARKER);
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Helper: decode all frames until EndOfMessage and concatenate the chunks.
    fn collect_message(codec: &mut NetconfCodec, buf: &mut BytesMut) -> Bytes {
        let mut result = BytesMut::new();
        loop {
            match codec.decode(buf).unwrap() {
                Some(DecodedFrame::Chunk(chunk)) => result.extend_from_slice(&chunk),
                Some(DecodedFrame::EndOfMessage) => break,
                None => panic!("unexpected None before EndOfMessage"),
            }
        }
        result.freeze()
    }

    // ── EOM Decoder tests ───────────────────────────────────────────────

    #[test]
    fn eom_decode_complete_message() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let mut buf = BytesMut::from(&b"<rpc-reply/>]]>]]>"[..]);

        // First decode: yields the data before the marker as a Chunk
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(
            result,
            Some(DecodedFrame::Chunk(Bytes::from_static(b"<rpc-reply/>")))
        );

        // Second decode: yields EndOfMessage
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(result, Some(DecodedFrame::EndOfMessage));

        assert!(buf.is_empty());
    }

    #[test]
    fn eom_decode_incomplete_message() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let mut buf = BytesMut::from(&b"<rpc-reply/>"[..]);

        // 12 bytes available, holdback is 5. Safe to emit: 7 bytes.
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(
            result,
            Some(DecodedFrame::Chunk(Bytes::from_static(b"<rpc-re")))
        );

        // Remaining 5 bytes <= holdback, need more data.
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(result, None);
    }

    #[test]
    fn eom_decode_partial_marker() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let mut buf = BytesMut::from(&b"<ok/>]]>"[..]);

        // 8 bytes, holdback 5, safe 3 — yield first chunk
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"<ok")))
        );
        // Remaining 5 bytes <= holdback
        assert_eq!(codec.decode(&mut buf).unwrap(), None);

        // Complete the marker
        buf.extend_from_slice(b"]]>");
        // Now buf is b"/>]]>]]>" — marker found at pos 2
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"/>")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
    }

    #[test]
    fn eom_decode_empty_message() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let mut buf = BytesMut::from(&b"]]>]]>"[..]);

        // Marker at position 0 — empty message, yield EndOfMessage directly.
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
    }

    #[test]
    fn eom_decode_two_messages() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let mut buf = BytesMut::from(&b"<a/>]]>]]><b/>]]>]]>"[..]);

        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"<a/>")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"<b/>")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
    }

    #[test]
    fn eom_decode_within_holdback() {
        // Buffer smaller than holdback — need more data.
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let mut buf = BytesMut::from(&b"<ok>"[..]);
        assert_eq!(codec.decode(&mut buf).unwrap(), None);
    }

    #[test]
    fn eom_decode_large_message_streams_chunks() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let data = "x".repeat(1000);
        let mut buf = BytesMut::new();
        buf.extend_from_slice(data.as_bytes());
        buf.extend_from_slice(EOM_MARKER);

        // Collect all chunks — should reconstruct the original data.
        let collected = collect_message(&mut codec, &mut buf);
        assert_eq!(collected, Bytes::from(data));
    }

    // ── Chunked Decoder tests ───────────────────────────────────────────

    #[test]
    fn chunked_decode_single_chunk() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::from(&b"\n#7\n<data/>\n##\n"[..]);

        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"<data/>")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
        assert!(buf.is_empty());
    }

    #[test]
    fn chunked_decode_multiple_chunks_yields_individually() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::from(&b"\n#5\nHello\n#6\n World\n##\n"[..]);

        // Each chunk is yielded individually (not accumulated).
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"Hello")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b" World")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
    }

    #[test]
    fn chunked_decode_incomplete_header() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::from(&b"\n#"[..]);
        assert_eq!(codec.decode(&mut buf).unwrap(), None);
    }

    #[test]
    fn chunked_decode_incomplete_data() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::from(&b"\n#10\nHello"[..]);
        assert_eq!(codec.decode(&mut buf).unwrap(), None);

        // Complete the data + end marker
        buf.extend_from_slice(b" Wrld\n##\n");
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"Hello Wrld")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
    }

    #[test]
    fn chunked_decode_large_chunk() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let data = "x".repeat(10000);
        let mut buf = BytesMut::new();
        buf.extend_from_slice(format!("\n#{}\n", data.len()).as_bytes());
        buf.extend_from_slice(data.as_bytes());
        buf.extend_from_slice(b"\n##\n");

        let result = codec.decode(&mut buf).unwrap();
        match result {
            Some(DecodedFrame::Chunk(chunk)) => assert_eq!(chunk.len(), 10000),
            other => panic!("expected Chunk, got {:?}", other),
        }
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
    }

    #[test]
    fn chunked_decode_invalid_header() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::from(&b"\n#abc\n"[..]);
        let err = codec.decode(&mut buf).unwrap_err();
        assert!(matches!(err, FramingError::InvalidChunkSize(_)));
    }

    #[test]
    fn chunked_decode_zero_chunk_size() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::from(&b"\n#0\n\n##\n"[..]);
        let err = codec.decode(&mut buf).unwrap_err();
        assert!(matches!(err, FramingError::InvalidChunkSize(_)));
    }

    #[test]
    fn chunked_decode_size_limit() {
        let config = CodecConfig {
            max_message_size: Some(5),
        };
        let mut codec = NetconfCodec::new(FramingMode::Chunked, config);
        let mut buf = BytesMut::from(&b"\n#10\n0123456789\n##\n"[..]);
        let err = codec.decode(&mut buf).unwrap_err();
        assert!(matches!(err, FramingError::MessageTooLarge { .. }));
    }

    #[test]
    fn chunked_decode_size_limit_cumulative() {
        let config = CodecConfig {
            max_message_size: Some(10),
        };
        let mut codec = NetconfCodec::new(FramingMode::Chunked, config);
        // Three 5-byte chunks: cumulative 15 bytes exceeds the 10-byte limit
        let mut buf = BytesMut::from(&b"\n#5\naaaaa\n#5\nbbbbb\n#5\nccccc\n##\n"[..]);

        // First chunk (5 bytes) — under limit
        assert!(codec.decode(&mut buf).unwrap().is_some());
        // Second chunk (cumulative 10 bytes) — at limit
        assert!(codec.decode(&mut buf).unwrap().is_some());
        // Third chunk (cumulative 15 bytes) — exceeds limit
        let err = codec.decode(&mut buf).unwrap_err();
        assert!(matches!(err, FramingError::MessageTooLarge { .. }));
    }

    #[test]
    fn eom_decode_size_limit() {
        let config = CodecConfig {
            max_message_size: Some(10),
        };
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, config);
        let mut buf = BytesMut::new();
        // 20 bytes of data exceeds the 10-byte limit
        buf.extend_from_slice(b"01234567890123456789]]>]]>");
        let err = codec.decode(&mut buf).unwrap_err();
        assert!(matches!(err, FramingError::MessageTooLarge { .. }));
    }

    #[test]
    fn eom_decode_size_limit_cumulative() {
        let config = CodecConfig {
            max_message_size: Some(10),
        };
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, config);

        // First batch: 8 bytes of data, no marker yet. Holdback keeps 5,
        // so 3 bytes are yielded — under the limit.
        let mut buf = BytesMut::from(&b"01234567"[..]);
        assert!(codec.decode(&mut buf).unwrap().is_some()); // yields 3 bytes

        // Second batch: more data + marker. Cumulative now exceeds 10.
        buf.extend_from_slice(b"89ABCDEF]]>]]>");
        let err = codec.decode(&mut buf).unwrap_err();
        assert!(matches!(err, FramingError::MessageTooLarge { .. }));
    }

    #[test]
    fn eom_decode_size_limit_resets_between_messages() {
        let config = CodecConfig {
            max_message_size: Some(10),
        };
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, config);

        // First message: 5 bytes — under the limit
        let mut buf = BytesMut::from(&b"hello]]>]]>"[..]);
        let collected = collect_message(&mut codec, &mut buf);
        assert_eq!(collected, Bytes::from_static(b"hello"));

        // Second message: also 5 bytes — should succeed (counter was reset)
        buf.extend_from_slice(b"world]]>]]>");
        let collected = collect_message(&mut codec, &mut buf);
        assert_eq!(collected, Bytes::from_static(b"world"));
    }

    // ── Encoder tests ───────────────────────────────────────────────────

    #[test]
    fn eom_encode() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let mut buf = BytesMut::new();
        codec
            .encode(Bytes::from_static(b"<ok/>"), &mut buf)
            .unwrap();
        assert_eq!(&buf[..], b"<ok/>]]>]]>");
    }

    #[test]
    fn chunked_encode() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::new();
        codec
            .encode(Bytes::from_static(b"<ok/>"), &mut buf)
            .unwrap();
        assert_eq!(&buf[..], b"\n#5\n<ok/>\n##\n");
    }

    // ── Roundtrip tests ─────────────────────────────────────────────────

    #[test]
    fn eom_roundtrip() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let original = Bytes::from_static(b"<rpc message-id=\"1\"><get/></rpc>");
        let mut buf = BytesMut::new();
        codec.encode(original.clone(), &mut buf).unwrap();

        let collected = collect_message(&mut codec, &mut buf);
        assert_eq!(collected, original);
    }

    #[test]
    fn chunked_roundtrip() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let original = Bytes::from_static(b"<rpc message-id=\"1\"><get/></rpc>");
        let mut buf = BytesMut::new();
        codec.encode(original.clone(), &mut buf).unwrap();

        let collected = collect_message(&mut codec, &mut buf);
        assert_eq!(collected, original);
    }

    #[test]
    fn mode_switch() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());

        // Encode/decode in EOM mode
        let mut buf = BytesMut::new();
        codec
            .encode(Bytes::from_static(b"hello"), &mut buf)
            .unwrap();
        let collected = collect_message(&mut codec, &mut buf);
        assert_eq!(collected, Bytes::from_static(b"hello"));

        // Switch to chunked
        codec.set_mode(FramingMode::Chunked);

        let mut buf = BytesMut::new();
        codec
            .encode(Bytes::from_static(b"world"), &mut buf)
            .unwrap();
        let collected = collect_message(&mut codec, &mut buf);
        assert_eq!(collected, Bytes::from_static(b"world"));
    }

    // ── EOM hello helpers tests ─────────────────────────────────────────

    #[tokio::test]
    async fn eom_helper_roundtrip() {
        let (mut client, mut server) = tokio::io::duplex(4096);

        let msg = "<hello/>";
        tokio::spawn(async move {
            write_eom_message(&mut server, msg).await.unwrap();
        });

        let received = read_eom_message(&mut client, None).await.unwrap();
        assert_eq!(received, msg);
    }

    #[tokio::test]
    async fn eom_helper_size_limit() {
        let (mut client, mut server) = tokio::io::duplex(4096);

        let msg = "x".repeat(1000);
        tokio::spawn(async move {
            write_eom_message(&mut server, &msg).await.unwrap();
        });

        let result = read_eom_message(&mut client, Some(10)).await;
        assert!(result.is_err());
    }

    // ── extract_message_id_from_bytes tests ─────────────────────────────

    #[test]
    fn extract_message_id_basic() {
        let xml =
            b"<rpc-reply message-id=\"42\" xmlns=\"urn:ietf:params:xml:ns:netconf:base:1.0\">";
        assert_eq!(extract_message_id_from_bytes(xml), Some(42));
    }

    #[test]
    fn extract_message_id_large_id() {
        let xml = b"<rpc-reply message-id=\"4294967295\">";
        assert_eq!(extract_message_id_from_bytes(xml), Some(4294967295));
    }

    #[test]
    fn extract_message_id_with_xml_decl() {
        let xml = b"<?xml version=\"1.0\"?><rpc-reply message-id=\"7\">";
        assert_eq!(extract_message_id_from_bytes(xml), Some(7));
    }

    #[test]
    fn extract_message_id_missing() {
        let xml = b"<rpc-reply xmlns=\"urn:ietf:params:xml:ns:netconf:base:1.0\">";
        assert_eq!(extract_message_id_from_bytes(xml), None);
    }

    #[test]
    fn extract_message_id_non_numeric() {
        let xml = b"<rpc-reply message-id=\"abc\">";
        assert_eq!(extract_message_id_from_bytes(xml), None);
    }

    #[test]
    fn extract_message_id_empty_bytes() {
        assert_eq!(extract_message_id_from_bytes(b""), None);
    }

    #[test]
    fn extract_message_id_partial_header() {
        // Only have the start of the attribute, no closing quote
        let xml = b"<rpc-reply message-id=\"12";
        assert_eq!(extract_message_id_from_bytes(xml), None);
    }

    // ── Chunk-level decode tests (new) ──────────────────────────────────

    #[test]
    fn eom_incremental_arrival() {
        // Simulate data arriving in small pieces.
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let mut buf = BytesMut::new();

        // First batch: just a few bytes (within holdback)
        buf.extend_from_slice(b"<ok");
        assert_eq!(codec.decode(&mut buf).unwrap(), None);

        // Second batch: more data, now exceeds holdback
        buf.extend_from_slice(b"/>");
        // buf = b"<ok/>" (5 bytes), exactly holdback — still None
        assert_eq!(codec.decode(&mut buf).unwrap(), None);

        // Third batch: data + marker
        buf.extend_from_slice(b"]]>]]>");
        // buf = b"<ok/>]]>]]>" — marker found at pos 5
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"<ok/>")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
    }

    #[test]
    fn chunked_two_messages_interleaved() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::new();
        buf.extend_from_slice(b"\n#3\naaa\n##\n\n#3\nbbb\n##\n");

        // First message
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"aaa")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );

        // Second message
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"bbb")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
    }

    #[test]
    fn chunked_three_chunks_in_one_message() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::from(&b"\n#1\na\n#1\nb\n#1\nc\n##\n"[..]);

        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"a")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"b")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::Chunk(Bytes::from_static(b"c")))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(DecodedFrame::EndOfMessage)
        );
    }
}