netconf_rust/

session.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::{Arc, Mutex};

use bytes::{Bytes, BytesMut};
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf};
use tokio_stream::StreamExt;
use tokio_util::codec::{Encoder, FramedRead};

use crate::codec::{CodecConfig, FramingMode, NetconfCodec};
use crate::hello::ServerHello;
use crate::message::{self, DataPayload, RpcReply, RpcReplyBody, ServerMessage};
use crate::stream::NetconfStream;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum SessionState {
    /// Hello exchange complete, ready for RPCs
    Ready = 0,
    /// A 'close-session' RPC has been sent, awaiting reply
    Closing = 1,
    /// Session terminated gracefully or with error
    Closed = 2,
}

impl SessionState {
    fn from_u8(v: u8) -> Self {
        match v {
            0 => Self::Ready,
            1 => Self::Closing,
            _ => Self::Closed,
        }
    }
}

impl std::fmt::Display for SessionState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Ready => write!(f, "Ready"),
            Self::Closing => write!(f, "Closing"),
            Self::Closed => write!(f, "Closed"),
        }
    }
}

/// Configuration for a NETCONF session.
#[derive(Debug, Clone, Default)]
pub struct SessionConfig {
    /// Codec configuration (message size limits, etc.).
    pub codec: CodecConfig,
}

#[derive(Debug, Clone, Copy)]
pub enum Datastore {
    Running,
    Candidate,
    Startup,
}

impl Datastore {
    fn as_xml(&self) -> &'static str {
        match self {
            Datastore::Running => "<running/>",
            Datastore::Candidate => "<candidate/>",
            Datastore::Startup => "<startup/>",
        }
    }
}

/// A handle to a pending RPC reply.
///
/// Created by [`Session::rpc_send()`], this allows sending multiple RPCs
/// before awaiting any replies (pipelining).
pub struct RpcFuture {
    rx: tokio::sync::oneshot::Receiver<crate::Result<RpcReply>>,
    msg_id: u32,
}

impl RpcFuture {
    /// The message-id of this RPC.
    pub fn message_id(&self) -> u32 {
        self.msg_id
    }

    /// Await the RPC reply.
    pub async fn response(self) -> crate::Result<RpcReply> {
        self.rx.await.map_err(|_| crate::Error::SessionClosed)?
    }
}

struct SessionInner {
    /// Pending RPC replies: message-id → oneshot sender.
    pending: Mutex<HashMap<u32, tokio::sync::oneshot::Sender<crate::Result<RpcReply>>>>,
    /// Session lifecycle state, stored as AtomicU8 for lock-free access
    /// across the session and reader task.
    ///
    /// State transitions:
    ///   Ready → Closing → Closed   (graceful: user calls close_session)
    ///   Ready → Closed             (abrupt: reader hits EOF or error)
    ///
    /// - Ready:   normal operation, RPCs can be sent
    /// - Closing: close-session RPC in flight, new RPCs rejected
    /// - Closed:  session terminated, all operations fail
    state: AtomicU8,
}

impl SessionInner {
    fn state(&self) -> SessionState {
        SessionState::from_u8(self.state.load(Ordering::Acquire))
    }

    fn set_state(&self, state: SessionState) {
        self.state.store(state as u8, Ordering::Release);
    }
}

/// A NETCONF session with pipelining support.
///
/// Architecture:
/// - **Writer**: the session holds the write half of the stream and a codec
///   for encoding outgoing RPCs.
/// - **Reader task**: a background tokio task reads framed messages from
///   the read half, classifies them (`<rpc-reply>` vs `<notification>`),
///   and routes them to the correct handler.
/// - **Pipelining**: [`rpc_send()`](Session::rpc_send) writes an RPC and
///   returns an [`RpcFuture`] without waiting for the reply. Multiple RPCs
///   can be in flight simultaneously.
pub struct Session {
    /// Write half of the split stream, used to send RPCs to the server.
    writer: WriteHalf<NetconfStream>,

    /// Codec used to frame outgoing messages (EOM or chunked).
    /// Kept separate from FramedRead because the write half isn't wrapped
    /// in FramedWrite — we encode manually and write_all to the writer.
    write_codec: NetconfCodec,

    /// Shared state between this session and the background reader task.
    /// Contains the pending RPC map, notification channel, and session state.
    inner: Arc<SessionInner>,

    /// The server's hello response, containing its capabilities and session ID.
    server_hello: ServerHello,

    /// Negotiated framing mode (EOM for 1.0-only servers, chunked for 1.1).
    framing: FramingMode,

    /// Holds the SSH handle alive. Dropping this tears down the SSH connection,
    /// which would invalidate the stream. Never accessed, just kept alive.
    _keep_alive: Option<Box<dyn std::any::Any + Send>>,

    /// Handle to the background reader task. Aborted on drop to ensure
    /// the task doesn't outlive the session.
    _reader_handle: tokio::task::JoinHandle<()>,
}

impl Drop for Session {
    fn drop(&mut self) {
        self._reader_handle.abort();
    }
}

impl Session {
    /// Connect to a NETCONF server over SSH with password authentication.
    pub async fn connect(
        host: &str,
        port: u16,
        username: &str,
        password: &str,
    ) -> crate::Result<Self> {
        Self::connect_with_config(host, port, username, password, SessionConfig::default()).await
    }

    /// Connect with custom configuration.
    pub async fn connect_with_config(
        host: &str,
        port: u16,
        username: &str,
        password: &str,
        config: SessionConfig,
    ) -> crate::Result<Self> {
        let (mut stream, keep_alive) =
            crate::transport::connect(host, port, username, password).await?;
        let (server_hello, framing) =
            crate::hello::exchange(&mut stream, config.codec.max_message_size).await?;
        Self::build(stream, Some(keep_alive), server_hello, framing, config)
    }

    /// Create a session from an existing stream (useful for testing).
    pub async fn from_stream<S: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
        stream: S,
    ) -> crate::Result<Self> {
        Self::from_stream_with_config(stream, SessionConfig::default()).await
    }

    /// Create a session from an existing stream with custom configuration.
    pub async fn from_stream_with_config<S: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
        mut stream: S,
        config: SessionConfig,
    ) -> crate::Result<Self> {
        let (server_hello, framing) =
            crate::hello::exchange(&mut stream, config.codec.max_message_size).await?;
        let boxed: NetconfStream = Box::new(stream);
        Self::build(boxed, None, server_hello, framing, config)
    }

    fn build(
        stream: NetconfStream,
        keep_alive: Option<Box<dyn std::any::Any + Send>>,
        server_hello: ServerHello,
        framing: FramingMode,
        config: SessionConfig,
    ) -> crate::Result<Self> {
        let (read_half, write_half) = tokio::io::split(stream);

        let read_codec = NetconfCodec::new(framing, config.codec);
        let write_codec = NetconfCodec::new(framing, config.codec);
        let reader = FramedRead::new(read_half, read_codec);

        let inner = Arc::new(SessionInner {
            pending: Mutex::new(HashMap::new()),
            state: AtomicU8::new(SessionState::Ready as u8),
        });

        let reader_inner = Arc::clone(&inner);
        let reader_handle = tokio::spawn(reader_loop(reader, reader_inner));

        Ok(Self {
            writer: write_half,
            write_codec,
            inner,
            server_hello,
            framing,
            _keep_alive: keep_alive,
            _reader_handle: reader_handle,
        })
    }

    pub fn session_id(&self) -> u32 {
        self.server_hello.session_id
    }

    pub fn server_capabilities(&self) -> &[String] {
        &self.server_hello.capabilities
    }

    pub fn framing_mode(&self) -> FramingMode {
        self.framing
    }

    pub fn state(&self) -> SessionState {
        self.inner.state()
    }

    fn check_state(&self) -> crate::Result<()> {
        let state = self.inner.state();
        if state != SessionState::Ready {
            return Err(crate::Error::InvalidState(state.to_string()));
        }
        Ok(())
    }

    /// Encode a message with the negotiated framing (EOM or chunked) and
    /// write it to the stream. We encode manually rather than using
    /// FramedWrite because NETCONF requires complete XML documents —
    /// the server doesn't process anything until the message delimiter
    /// arrives. A simple encode + write_all + flush is sufficient since
    /// we always write whole messages. The kernel's TCP send buffer
    /// handles chunking large writes over the wire, and the single
    /// Session owner means there's no need for Sink-based backpressure.
    async fn send_encoded(&mut self, xml: &str) -> crate::Result<()> {
        let mut buf = BytesMut::new();
        self.write_codec
            .encode(Bytes::from(xml.to_string()), &mut buf)?;
        self.writer.write_all(&buf).await?;
        self.writer.flush().await?;
        Ok(())
    }

    /// Send a raw RPC and return a future for the reply (pipelining).
    ///
    /// This writes the RPC to the server immediately but does not wait
    /// for the reply. Call [`RpcFuture::response()`] to await the reply.
    /// Multiple RPCs can be pipelined by calling this repeatedly before
    /// awaiting any of them.
    pub async fn rpc_send(&mut self, inner_xml: &str) -> crate::Result<RpcFuture> {
        self.check_state()?;
        let (msg_id, xml) = message::build_rpc(inner_xml);
        let (tx, rx) = tokio::sync::oneshot::channel();

        self.inner.pending.lock().unwrap().insert(msg_id, tx);

        if let Err(e) = self.send_encoded(&xml).await {
            self.inner.pending.lock().unwrap().remove(&msg_id);
            return Err(e);
        }
        Ok(RpcFuture { rx, msg_id })
    }

    /// Send a raw RPC and wait for the reply.
    pub async fn rpc_raw(&mut self, inner_xml: &str) -> crate::Result<RpcReply> {
        let future = self.rpc_send(inner_xml).await?;
        future.response().await
    }

    /// Internal rpc send that skips state check. Only used for sending close-session.
    async fn rpc_send_unchecked(&mut self, inner_xml: &str) -> crate::Result<RpcFuture> {
        let (msg_id, xml) = message::build_rpc(inner_xml);
        let (tx, rx) = tokio::sync::oneshot::channel();

        self.inner.pending.lock().unwrap().insert(msg_id, tx);

        if let Err(e) = self.send_encoded(&xml).await {
            self.inner.pending.lock().unwrap().remove(&msg_id);
            return Err(e);
        }

        Ok(RpcFuture { rx, msg_id })
    }

    /// Retrieve configuration from a datastore.
    pub async fn get_config(
        &mut self,
        source: Datastore,
        filter: Option<&str>,
    ) -> crate::Result<String> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!(
            "<get-config><source>{}</source>{filter_xml}</get-config>",
            source.as_xml()
        );
        let reply = self.rpc_raw(&inner).await?;
        reply_to_data(reply)
    }

    /// Retrieve configuration as a zero-copy `DataPayload`.
    ///
    /// Same as `get_config()` but returns a `DataPayload` instead of `String`,
    /// avoiding a copy of the response body. Use `payload.as_str()` for a
    /// zero-copy `&str` view, or `payload.reader()` for streaming XML events.
    pub async fn get_config_payload(
        &mut self,
        source: Datastore,
        filter: Option<&str>,
    ) -> crate::Result<DataPayload> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!(
            "<get-config><source>{}</source>{filter_xml}</get-config>",
            source.as_xml()
        );
        let reply = self.rpc_raw(&inner).await?;
        reply.into_data()
    }

    /// Retrieve running configuration and state data.
    pub async fn get(&mut self, filter: Option<&str>) -> crate::Result<String> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!("<get>{filter_xml}</get>");
        let reply = self.rpc_raw(&inner).await?;
        reply_to_data(reply)
    }

    /// Retrieve running configuration and state data as a zero-copy `DataPayload`.
    ///
    /// Same as `get()` but returns a `DataPayload` instead of `String`,
    /// avoiding a copy of the response body.
    pub async fn get_payload(&mut self, filter: Option<&str>) -> crate::Result<DataPayload> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!("<get>{filter_xml}</get>");
        let reply = self.rpc_raw(&inner).await?;
        reply.into_data()
    }

    /// Edit the configuration of a target datastore.
    pub async fn edit_config(&mut self, target: Datastore, config: &str) -> crate::Result<()> {
        let inner = format!(
            "<edit-config><target>{}</target><config>{config}</config></edit-config>",
            target.as_xml()
        );
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Lock a datastore
    pub async fn lock(&mut self, target: Datastore) -> crate::Result<()> {
        let inner = format!("<lock><target>{}</target></lock>", target.as_xml());
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Unlock a datastore.
    pub async fn unlock(&mut self, target: Datastore) -> crate::Result<()> {
        let inner = format!("<unlock><target>{}</target></unlock>", target.as_xml());
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Commit the candidate configuration to running.
    pub async fn commit(&mut self) -> crate::Result<()> {
        let reply = self.rpc_raw("<commit/>").await?;
        reply_to_ok(reply)
    }

    /// Gracefully close the NETCONF session.
    pub async fn close_session(&mut self) -> crate::Result<()> {
        self.check_state()?;
        self.inner.set_state(SessionState::Closing);
        let result = self.rpc_send_unchecked("<close-session/>").await;
        match result {
            Ok(future) => {
                let reply = future.response().await;
                self.inner.set_state(SessionState::Closed);
                reply_to_ok(reply?)
            }
            Err(e) => {
                self.inner.set_state(SessionState::Closed);
                Err(e)
            }
        }
    }

    /// Force-close another NETCONF session.
    pub async fn kill_session(&mut self, session_id: u32) -> crate::Result<()> {
        let inner = format!("<kill-session><session-id>{session_id}</session-id></kill-session>");
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }
}

/// This loop is the only thing reading from the SSH stream. It runs in a background tokio task.
/// The session's main API (the writer side) never reads — it only writes RPCs and waits on oneshot channels.
/// This separation is what makes pipelining work.
async fn reader_loop(
    mut reader: FramedRead<ReadHalf<NetconfStream>, NetconfCodec>,
    inner: Arc<SessionInner>,
) {
    // FramedRead wraps the ReadHalf with the NetconfCodec decoder.
    // Each .next() reads bytes from SSH, feeds them to decode(), and yields a
    // complete framed message as Bytes.
    while let Some(result) = reader.next().await {
        match result {
            // Takes the raw Bytes, validates UTF-8, finds the root element, parses it.
            Ok(bytes) => match message::classify_message(bytes) {
                Ok(ServerMessage::RpcReply(reply)) => {
                    // When rpc_send sends an RPC, it inserts a oneshot sender into the pending map
                    // keyed by message-id. Here, it is removed and used to send the reply
                    // through the channel. The caller who's .awaiting the RpcFuture receives it.
                    let tx = {
                        let mut pending = inner.pending.lock().unwrap();
                        pending.remove(&reply.message_id)
                    };
                    if let Some(tx) = tx {
                        // we can ignore if the receiver was dropped. Nothing to do about it.
                        let _ = tx.send(Ok(reply));
                    } else {
                        log::warn!("received reply for unknown message-id {}", reply.message_id);
                    }
                }
                Err(e) => {
                    log::warn!("failed to classify message: {e}");
                }
            },
            // The stream broke (EOF or Error). Every pending RPC would hang forever waiting for a reply that's never coming.
            // So we drain the map and send SessionClosed to each one, unblocking all waiters.
            Err(e) => {
                log::error!("reader error: {e}");
                let mut pending = inner.pending.lock().unwrap();
                for (_, tx) in pending.drain() {
                    let _ = tx.send(Err(crate::Error::SessionClosed));
                }
                break;
            }
        }
    }
    // Drain any remaining pending RPCs — this handles clean EOF where the while loop exits
    // via None (stream closed) rather than Err (which drains inline above).
    {
        let mut pending = inner.pending.lock().unwrap();
        for (_, tx) in pending.drain() {
            let _ = tx.send(Err(crate::Error::SessionClosed));
        }
    }
    // Mark the session as closed so future rpc_send calls fail immediately with InvalidState
    // instead of inserting into the pending map and hanging.
    inner.set_state(SessionState::Closed);
}

fn reply_to_data(reply: RpcReply) -> crate::Result<String> {
    match reply.body {
        RpcReplyBody::Data(payload) => Ok(payload.into_string()),
        RpcReplyBody::Ok => Ok(String::new()),
        RpcReplyBody::Error(errors) => Err(crate::Error::Rpc {
            message_id: reply.message_id,
            error: errors
                .first()
                .map(|e| e.error_message.clone())
                .unwrap_or_default(),
        }),
    }
}

fn reply_to_ok(reply: RpcReply) -> crate::Result<()> {
    match reply.body {
        RpcReplyBody::Ok => Ok(()),
        RpcReplyBody::Data(_) => Ok(()),
        RpcReplyBody::Error(errors) => Err(crate::Error::Rpc {
            message_id: reply.message_id,
            error: errors
                .first()
                .map(|e| e.error_message.clone())
                .unwrap_or_default(),
        }),
    }
}