netconf_rust/

codec.rs

use bytes::{Buf, Bytes, BytesMut};
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#";

/// 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,
    chunked_buf: BytesMut,
}
/// tokio_util codec for NETCONF message framing.
///
/// Implements the `Decoder` trait from tokio_util, which defines how to
/// extract discrete messages from a continuous TCP byte stream.
///
/// We don't call `decode()` directly. Instead, we pass this codec to
/// `FramedRead`, which handles the I/O loop:
///
///   let reader = FramedRead::new(read_half, NetconfCodec::new(...));
///   while let Some(msg) = reader.next().await { ... }
///
/// `FramedRead` reads bytes from the network into an internal buffer,
/// calls our `decode()` to check for a complete message, and yields
/// messages as a Stream. We only need to implement the framing logic --
/// "is there a complete message in this buffer?" -- and `FramedRead`
/// handles everything else (buffer management, read loops, EOF).
///
/// decode() is called repeatedly. It owns the internal buffer. When
/// decode() returns Ok(None), FramedRead reads more bytes from the
/// underlying AsyncRead into the buffer and calls decode() again.
/// When decode() returns Ok(Some(item)), the item is yielded to the
/// consumer. The key invariant is: decode() must consume
/// (via split_to() or advance()) exactly the bytes belonging to the
/// returned message. Any leftover bytes in src remain for the next
/// decode() call.
///
/// Returns `Bytes` (not `String`) because UTF-8 validation is deferred
/// to the XML parser. This avoids a redundant O(n) validation pass at
/// the framing layer
impl NetconfCodec {
    pub fn new(framing_mode: FramingMode, config: CodecConfig) -> Self {
        Self {
            framing_mode,
            config,
            chunked_buf: BytesMut::new(),
        }
    }

    pub fn set_mode(&mut self, framing_mode: FramingMode) {
        self.framing_mode = framing_mode;
        self.chunked_buf.clear();
    }
    pub fn framing_mode(&self) -> FramingMode {
        self.framing_mode
    }

    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(())
    }

    /// Scans src for ]]>]]> and returns message bytes if found otherwise we need more data.
    /// We assume that the marker only appears at the message boundry because
    /// 1.  XML itself forbids ]]> inside content. so ]]>]]> can never appear in well-formed XML.
    /// 2.  RFC 4742 Section 4.2 defines the framing: the message is terminated by ]]>]]>.
    ///     The content should always be valid XML, and valid XML can't contain ]]>
    ///     However, in Netconf 1.0, there is actually no guarantee. This fact relies on the netconf server actually
    ///     returning well formed XML. If it doesn't then our code will find the first occurrence
    ///     the marker, split the message in the wrong place and then return an error.
    fn decode_eom(&self, src: &mut BytesMut) -> Result<Option<Bytes>, FramingError> {
        // The smallest possible complete message is EOM_MARKER, so if src is smaller than
        // that, then we need more data
        if src.len() < EOM_LEN {
            return Ok(None);
        }
        if let Some(eom_pos_start) = find_subsequence(src, EOM_MARKER) {
            let msg_len = eom_pos_start;
            self.check_size(msg_len)?;
            let msg = src.split_to(msg_len).freeze();
            src.advance(EOM_LEN);
            Ok(Some(msg))
        } else {
            self.check_size(src.len())?;
            Ok(None)
        }
    }

    fn decode_chunked(&mut self, src: &mut BytesMut) -> Result<Option<Bytes>, FramingError> {
        // The smallest possible complete message is CHUNKED_EOM_MARKER. Otherwise more
        // data is expected.
        loop {
            if src.len() < CHUNKED_EOM_MARKER_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.
            // We know the first 2 bytes are valid already
            if src[2] == b'#' {
                if src[3] != b'\n' {
                    return Err(FramingError::InvalidHeader {
                        expected: "\\n##\\n",
                        got: src[..4].to_vec(),
                    });
                }

                // We found end of message. Consume \n##\n
                src.advance(CHUNKED_EOM_MARKER_LEN);
                let msg = self.chunked_buf.split().freeze();
                return Ok(Some(msg));
            }

            // Otherwise not end of message, so we have 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.
                    // But sanity-check the chuck size
                    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 {
                return Ok(None); // Need more data for the chunk
            }
            self.check_size(self.chunked_buf.len() + chunk_size)?;

            // Consume header
            src.advance(header_len);

            // Consume the chunk of data
            self.chunked_buf.extend_from_slice(&src[..chunk_size]);
            src.advance(chunk_size);
        }
    }
}

// find ]]>]]> and return the start byte index
// TODO: Potential optimization: why scan the thing again?
fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option<usize> {
    haystack
        .windows(needle.len())
        .position(|window| window == needle)
}

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 {
            return Err(FramingError::UnexpectedEof.into());
        }
        buf.extend_from_slice(&tmp[..read_bytes]);
        if let Some(limit) = max_size
            && buf.len() > limit + EOM_LEN
        {
            return Err(FramingError::MessageTooLarge {
                limit,
                received: buf.len(),
            }
            .into());
        }
        if buf.len() >= EOM_LEN && buf[buf.len() - EOM_LEN..] == *EOM_MARKER {
            buf.truncate(buf.len() - EOM_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(())
}

/// <https://docs.rs/tokio-util/latest/tokio_util/codec/trait.Decoder.html>
impl Decoder for NetconfCodec {
    type Item = Bytes;
    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),
        }
    }
}

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

    fn encode(&mut self, item: Bytes, dst: &mut BytesMut) -> Result<(), Self::Error> {
        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::*;

    // ── 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/>]]>]]>"[..]);
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(result, Some(Bytes::from_static(b"<rpc-reply/>")));
        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/>"[..]);
        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/>]]>"[..]);
        assert_eq!(codec.decode(&mut buf).unwrap(), None);
        // Now complete the marker
        buf.extend_from_slice(b"]]>");
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(result, Some(Bytes::from_static(b"<ok/>")));
    }

    #[test]
    fn eom_decode_empty_message() {
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, CodecConfig::default());
        let mut buf = BytesMut::from(&b"]]>]]>"[..]);
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(result, Some(Bytes::from_static(b"")));
    }

    #[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(Bytes::from_static(b"<a/>"))
        );
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(Bytes::from_static(b"<b/>"))
        );
    }

    #[test]
    fn eom_decode_size_limit_exceeded() {
        let config = CodecConfig {
            max_message_size: Some(5),
        };
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, config);
        let mut buf = BytesMut::from(&b"<too-large/>]]>]]>"[..]);
        let err = codec.decode(&mut buf).unwrap_err();
        assert!(matches!(err, FramingError::MessageTooLarge { .. }));
    }

    #[test]
    fn eom_decode_size_limit_ok() {
        let config = CodecConfig {
            max_message_size: Some(100),
        };
        let mut codec = NetconfCodec::new(FramingMode::EndOfMessage, config);
        let mut buf = BytesMut::from(&b"<ok/>]]>]]>"[..]);
        assert!(codec.decode(&mut buf).unwrap().is_some());
    }

    // ── 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"[..]);
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(result, Some(Bytes::from_static(b"<data/>")));
        assert!(buf.is_empty());
    }

    #[test]
    fn chunked_decode_multiple_chunks() {
        let mut codec = NetconfCodec::new(FramingMode::Chunked, CodecConfig::default());
        let mut buf = BytesMut::from(&b"\n#5\nHello\n#6\n World\n##\n"[..]);
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(result, Some(Bytes::from_static(b"Hello World")));
    }

    #[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");
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(result, Some(Bytes::from_static(b"Hello Wrld")));
    }

    #[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();
        assert_eq!(result.unwrap().len(), 10000);
    }

    #[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 { .. }));
    }

    // ── 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 decoded = codec.decode(&mut buf).unwrap().unwrap();
        assert_eq!(decoded, 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 decoded = codec.decode(&mut buf).unwrap().unwrap();
        assert_eq!(decoded, 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();
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(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();
        assert_eq!(
            codec.decode(&mut buf).unwrap(),
            Some(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());
    }
}