netconf_rust/

hello.rs

use std::io::Cursor;

use quick_xml::events::{BytesDecl, BytesEnd, BytesStart, BytesText, Event};
use quick_xml::{Reader, Writer};
use tokio::io::{AsyncRead, AsyncWrite};

use crate::capabilities;
use crate::codec::{self, FramingMode};

#[derive(Debug)]
pub struct ServerHello {
    pub session_id: u32,
    pub capabilities: Vec<String>,
}

pub async fn exchange<S: AsyncRead + AsyncWrite + Unpin>(
    stream: &mut S,
    max_message_size: Option<usize>,
) -> crate::Result<(ServerHello, FramingMode)> {
    let server_hello_xml = codec::read_eom_message(stream, max_message_size).await?;
    let server_hello = parse_server_hello(&server_hello_xml)?;

    let client_hello_xml = build_client_hello();
    codec::write_eom_message(stream, &client_hello_xml).await?;

    let framing_mode = if server_hello
        .capabilities
        .iter()
        .any(|c| c == capabilities::BASE_1_1)
    {
        FramingMode::Chunked
    } else {
        FramingMode::EndOfMessage
    };
    Ok((server_hello, framing_mode))
}

// Definitely would have been way easier with basic string stuff but i wanted to try out the API
// and this basically stays the same anyways.
fn build_client_hello() -> String {
    let mut writer = Writer::new_with_indent(Cursor::new(Vec::new()), b' ', 2);

    writer
        .write_event(Event::Decl(BytesDecl::new("1.0", Some("UTF-8"), None)))
        .unwrap();

    let mut hello = BytesStart::new("hello");
    hello.push_attribute(("xmlns", "urn:ietf:params:xml:ns:netconf:base:1.0"));
    writer.write_event(Event::Start(hello)).unwrap();

    writer
        .write_event(Event::Start(BytesStart::new("capabilities")))
        .unwrap();

    for cap in capabilities::client_capabilities() {
        writer
            .create_element("capability")
            .write_text_content(BytesText::new(cap))
            .unwrap();
    }

    writer
        .write_event(Event::End(BytesEnd::new("capabilities")))
        .unwrap();
    writer
        .write_event(Event::End(BytesEnd::new("hello")))
        .unwrap();
    String::from_utf8(writer.into_inner().into_inner()).unwrap()
}

pub(crate) fn parse_server_hello(xml: &str) -> crate::Result<ServerHello> {
    let mut reader = Reader::from_str(xml);
    reader.config_mut().trim_text(true);

    let mut capabilities = Vec::new();
    let mut session_id: Option<u32> = None;

    // use these to track with section of the response we are in
    let mut in_capability = false;
    let mut in_session_id = false;

    loop {
        match reader.read_event() {
            Ok(Event::Start(e)) | Ok(Event::Empty(e)) => {
                let local_name = e.local_name();
                match local_name.as_ref() {
                    b"capability" => in_capability = true,
                    b"session-id" => in_session_id = true,
                    _ => {}
                }
            }
            Ok(Event::Text(e)) => {
                let text = e
                    .xml_content()
                    .map_err(|err| {
                        crate::error::Error::HelloFailed(format!("XML unescape error: {err}"))
                    })?
                    .to_string();

                if in_capability {
                    if !text.is_empty() {
                        capabilities.push(text);
                    }
                } else if in_session_id {
                    session_id = Some(text.parse::<u32>().map_err(|e| {
                        crate::Error::HelloFailed(format!("invalid session-id '{text}': {e}"))
                    })?);
                }
            }
            Ok(Event::End(e)) => {
                let local_name = e.local_name();
                match local_name.as_ref() {
                    b"capability" => in_capability = false,
                    b"session-id" => in_session_id = false,
                    _ => {}
                }
            }
            Ok(Event::Eof) => break,
            Err(e) => return Err(crate::Error::HelloFailed(format!("XML parse error: {e}"))),
            _ => {}
        }
    }

    let session_id =
        session_id.ok_or_else(|| crate::Error::HelloFailed("missing session-id".into()))?;

    if capabilities.is_empty() {
        return Err(crate::Error::HelloFailed("no capabilities found".into()));
    }
    Ok(ServerHello {
        session_id,
        capabilities,
    })
}

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

    #[test]
    fn test_parse_server_hello() {
        let xml = r#"<?xml version="1.0" encoding="UTF-8"?>
<hello xmlns="urn:ietf:params:xml:ns:netconf:base:1.0">
  <capabilities>
    <capability>urn:ietf:params:netconf:base:1.0</capability>
    <capability>urn:ietf:params:netconf:base:1.1</capability>
    <capability>urn:ietf:params:netconf:capability:candidate:1.0</capability>
  </capabilities>
  <session-id>42</session-id>
</hello>"#;
        let hello = parse_server_hello(xml).unwrap();
        assert_eq!(hello.session_id, 42);
        assert_eq!(hello.capabilities.len(), 3);
        assert!(
            hello
                .capabilities
                .contains(&"urn:ietf:params:netconf:base:1.0".to_string())
        );
        assert!(
            hello
                .capabilities
                .contains(&"urn:ietf:params:netconf:base:1.1".to_string())
        );
    }

    #[test]
    fn test_parse_hello_1_0_only() {
        let xml = r#"<hello xmlns="urn:ietf:params:xml:ns:netconf:base:1.0">
  <capabilities>
    <capability>urn:ietf:params:netconf:base:1.0</capability>
  </capabilities>
  <session-id>1</session-id>
</hello>"#;
        let hello = parse_server_hello(xml).unwrap();
        assert_eq!(hello.session_id, 1);
        assert_eq!(hello.capabilities.len(), 1);
    }

    #[test]
    fn test_parse_hello_missing_session_id() {
        let xml = r#"<hello xmlns="urn:ietf:params:xml:ns:netconf:base:1.0">
  <capabilities>
    <capability>urn:ietf:params:netconf:base:1.0</capability>
  </capabilities>
</hello>"#;
        let result = parse_server_hello(xml);
        assert!(result.is_err());
    }

    #[test]
    fn test_parse_hello_no_capabilities() {
        let xml = r#"<hello xmlns="urn:ietf:params:xml:ns:netconf:base:1.0">
  <capabilities/>
  <session-id>1</session-id>
</hello>"#;
        let result = parse_server_hello(xml);
        assert!(result.is_err());
    }

    #[test]
    fn test_build_client_hello() {
        let hello = build_client_hello();
        assert!(hello.contains("urn:ietf:params:netconf:base:1.0"));
        assert!(hello.contains("urn:ietf:params:netconf:base:1.1"));
        assert!(hello.contains("<hello"));
        assert!(hello.contains("</hello>"));
        assert!(hello.contains("<capability>"));
        assert!(hello.contains("</capability>"));
    }

    #[test]
    fn test_build_client_hello_is_valid_xml() {
        let hello = build_client_hello();
        // Verify it parses without error
        let mut reader = Reader::from_str(&hello);
        reader.config_mut().trim_text(true);
        let mut caps = Vec::new();
        let mut in_cap = false;
        loop {
            match reader.read_event() {
                Ok(Event::Start(e)) if e.local_name().as_ref() == b"capability" => in_cap = true,
                Ok(Event::Text(e)) if in_cap => {
                    caps.push(
                        e.xml_content()
                            .map_err(|e| e.to_string())
                            .unwrap()
                            .to_string(),
                    );
                }
                Ok(Event::End(e)) if e.local_name().as_ref() == b"capability" => in_cap = false,
                Ok(Event::Eof) => break,
                Err(e) => panic!("invalid XML: {e}"),
                _ => {}
            }
        }
        assert!(!caps.is_empty());
        assert!(caps.contains(&"urn:ietf:params:netconf:base:1.0".to_string()));
        assert!(caps.contains(&"urn:ietf:params:netconf:base:1.1".to_string()));
    }

    #[tokio::test]
    async fn test_exchange_negotiates_chunked() {
        use tokio::io::AsyncWriteExt;
        use tokio::io::duplex;

        let (mut client, mut server) = duplex(4096);

        let server_hello = r#"<?xml version="1.0" encoding="UTF-8"?>
<hello xmlns="urn:ietf:params:xml:ns:netconf:base:1.0">
  <capabilities>
    <capability>urn:ietf:params:netconf:base:1.0</capability>
    <capability>urn:ietf:params:netconf:base:1.1</capability>
  </capabilities>
  <session-id>99</session-id>
</hello>"#;

        tokio::spawn(async move {
            use tokio::io::AsyncReadExt;
            server
                .write_all(format!("{server_hello}]]>]]>").as_bytes())
                .await
                .unwrap();
            let mut buf = vec![0u8; 4096];
            let _ = server.read(&mut buf).await;
        });

        let (hello, mode) = exchange(&mut client, None).await.unwrap();
        assert_eq!(hello.session_id, 99);
        assert_eq!(mode, FramingMode::Chunked);
    }

    #[tokio::test]
    async fn test_exchange_negotiates_eom() {
        use tokio::io::AsyncWriteExt;
        use tokio::io::duplex;

        let (mut client, mut server) = duplex(4096);

        let server_hello = r#"<?xml version="1.0" encoding="UTF-8"?>
<hello xmlns="urn:ietf:params:xml:ns:netconf:base:1.0">
  <capabilities>
    <capability>urn:ietf:params:netconf:base:1.0</capability>
  </capabilities>
  <session-id>10</session-id>
</hello>"#;

        tokio::spawn(async move {
            use tokio::io::AsyncReadExt;
            server
                .write_all(format!("{server_hello}]]>]]>").as_bytes())
                .await
                .unwrap();
            let mut buf = vec![0u8; 4096];
            let _ = server.read(&mut buf).await;
        });

        let (hello, mode) = exchange(&mut client, None).await.unwrap();
        assert_eq!(hello.session_id, 10);
        assert_eq!(mode, FramingMode::EndOfMessage);
    }
}