netconf_rust/

session.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

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

use crate::codec::{FramingMode, NetconfCodec};
use crate::config::Config;
use crate::error::TransportError;
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"),
        }
    }
}

/// Reason the session disconnected.
///
/// Delivered via [`Session::disconnected()`] when the background reader
/// task detects that the connection is no longer alive.
#[derive(Debug, Clone)]
pub enum DisconnectReason {
    /// The remote end closed the connection cleanly (TCP FIN / EOF).
    Eof,
    /// A transport error severed the connection.
    ///
    /// Contains the error's display string.
    TransportError(String),
    /// The [`Session`] was dropped without calling [`Session::close_session()`].
    Dropped,
}

impl std::fmt::Display for DisconnectReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Eof => write!(f, "connection closed by remote"),
            Self::TransportError(e) => write!(f, "transport error: {e}"),
            Self::Dropped => write!(f, "session dropped"),
        }
    }
}

#[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,
    rpc_timeout: Option<Duration>,
}

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

    /// Await the RPC reply.
    ///
    /// If an `rpc_timeout` was configured, this will fail with
    /// [`TransportError::Timeout`]
    /// if the server does not reply in time.
    pub async fn response(self) -> crate::Result<RpcReply> {
        let result = match self.rpc_timeout {
            Some(duration) => tokio::time::timeout(duration, self.rx)
                .await
                .map_err(|_| crate::Error::Transport(TransportError::Timeout(duration)))?,
            None => self.rx.await,
        };
        result.map_err(|_| crate::Error::SessionClosed)?
    }

    /// Await the RPC reply with an explicit timeout, ignoring the
    /// session-level `rpc_timeout`.
    ///
    /// Fails with [`TransportError::Timeout`]
    /// if the server does not reply within `timeout`.
    pub async fn response_with_timeout(self, timeout: Duration) -> crate::Result<RpcReply> {
        let result = tokio::time::timeout(timeout, self.rx)
            .await
            .map_err(|_| crate::Error::Transport(TransportError::Timeout(timeout)))?;
        result.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,
    /// Sender for the disconnect notification. The reader_loop sends the
    /// reason just before exiting. Receivers (from `Session::disconnected()`)
    /// wake up immediately.
    disconnect_tx: tokio::sync::watch::Sender<Option<DisconnectReason>>,
    /// Anchor instant for computing `last_rpc_at` from `last_rpc_nanos`.
    created_at: Instant,
    /// Nanoseconds elapsed since `created_at` when the last RPC reply was
    /// successfully routed. 0 means no reply has been received yet.
    last_rpc_nanos: AtomicU64,
}

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

    fn drain_pending(&self) -> usize {
        let mut pending = self.pending.lock().unwrap();
        let count = pending.len();
        for (_, tx) in pending.drain() {
            let _ = tx.send(Err(crate::Error::SessionClosed));
        }
        count
    }
}

/// Write half + codec, behind a `tokio::sync::Mutex` so Session methods
/// can take `&self`. The Mutex is held only during encode + write + flush.
struct WriterState {
    writer: WriteHalf<NetconfStream>,
    codec: NetconfCodec,
}

/// 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 state behind an async Mutex. Held only for the duration of
    /// encoding and flushing a single message.
    writer: tokio::sync::Mutex<WriterState>,

    /// 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,

    /// Timeout applied to each RPC response wait.
    rpc_timeout: Option<Duration>,

    /// Receiver for disconnect notifications. Cloned for each call to
    /// `disconnected()`, allowing multiple independent subscribers.
    disconnect_rx: tokio::sync::watch::Receiver<Option<DisconnectReason>>,

    /// Instant when the session was established (hello exchange complete).
    connected_since: Instant,

    /// Handle to the background reader task. Aborted on drop to ensure
    /// the task doesn't outlive the session. Declared before `_keep_alive`
    /// so it is dropped first — the reader task must be aborted before
    /// the SSH connection is torn down.
    _reader_handle: tokio::task::JoinHandle<()>,

    /// 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 + Sync>>,
}

impl Drop for Session {
    fn drop(&mut self) {
        // 1. Drain pending RPCs → waiters get deterministic SessionClosed
        let drained = self.inner.drain_pending();
        if drained > 0 {
            debug!(
                "session {}: drop: drained {drained} pending RPCs",
                self.server_hello.session_id
            );
        }
        // 2. Mark session closed
        self.inner.set_state(SessionState::Closed);
        // 3. Notify disconnect subscribers (only if reader_loop hasn't already)
        self.inner.disconnect_tx.send_if_modified(|current| {
            if current.is_none() {
                *current = Some(DisconnectReason::Dropped);
                true
            } else {
                false
            }
        });
        // 4. Abort reader task (its cleanup is now a no-op)
        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, Config::default()).await
    }

    /// Connect with custom configuration.
    pub async fn connect_with_config(
        host: &str,
        port: u16,
        username: &str,
        password: &str,
        config: Config,
    ) -> crate::Result<Self> {
        let (mut stream, keep_alive) =
            crate::transport::connect(host, port, username, password, &config).await?;
        let (server_hello, framing) = exchange_hello(&mut stream, &config).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, Config::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: Config,
    ) -> crate::Result<Self> {
        let (server_hello, framing) = exchange_hello(&mut stream, &config).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 + Sync>>,
        server_hello: ServerHello,
        framing: FramingMode,
        config: Config,
    ) -> crate::Result<Self> {
        debug!(
            "session {}: building (framing={:?}, capabilities={})",
            server_hello.session_id,
            framing,
            server_hello.capabilities.len()
        );
        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 (disconnect_tx, disconnect_rx) = tokio::sync::watch::channel(None);

        let inner = Arc::new(SessionInner {
            pending: Mutex::new(HashMap::new()),
            state: AtomicU8::new(SessionState::Ready as u8),
            disconnect_tx,
            created_at: Instant::now(),
            last_rpc_nanos: AtomicU64::new(0),
        });

        let reader_inner = Arc::clone(&inner);
        let session_id = server_hello.session_id;
        let reader_handle = tokio::spawn(async move {
            reader_loop(reader, reader_inner, session_id).await;
        });

        Ok(Self {
            writer: tokio::sync::Mutex::new(WriterState {
                writer: write_half,
                codec: write_codec,
            }),
            inner,
            server_hello,
            framing,
            rpc_timeout: config.rpc_timeout,
            disconnect_rx,
            connected_since: Instant::now(),
            _reader_handle: reader_handle,
            _keep_alive: keep_alive,
        })
    }

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

    /// Returns a future that completes when the session disconnects.
    ///
    /// Can be called multiple times — each call clones an internal
    /// `watch::Receiver`, so multiple tasks can independently await
    /// the same disconnect event. If the session is already disconnected
    /// when called, returns immediately.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use netconf_rust::Session;
    /// # async fn example(session: &Session) {
    /// let reason = session.disconnected().await;
    /// println!("session died: {reason}");
    /// # }
    /// ```
    pub async fn disconnected(&self) -> DisconnectReason {
        let mut rx = self.disconnect_rx.clone();
        // Check if already disconnected (late subscriber).
        if let Some(reason) = rx.borrow_and_update().clone() {
            return reason;
        }
        // Wait for the reader_loop to send the reason. If the sender
        // is dropped (Session dropped, reader aborted), treat as Dropped.
        loop {
            if rx.changed().await.is_err() {
                return DisconnectReason::Dropped;
            }
            if let Some(reason) = rx.borrow_and_update().clone() {
                return reason;
            }
        }
    }

    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.
    ///
    /// The internal Mutex is held only for encode + write + flush,
    /// NOT for the RPC response wait.
    async fn send_encoded(&self, xml: &str) -> crate::Result<()> {
        let mut buf = BytesMut::new();
        let mut state = self.writer.lock().await;
        state.codec.encode(Bytes::from(xml.to_string()), &mut buf)?;
        trace!(
            "session {}: writing {} bytes to stream",
            self.server_hello.session_id,
            buf.len()
        );
        state.writer.write_all(&buf).await?;
        state.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(&self, inner_xml: &str) -> crate::Result<RpcFuture> {
        self.check_state()?;
        let (msg_id, xml) = message::build_rpc(inner_xml);
        debug!(
            "session {}: sending rpc message-id={} ({} bytes)",
            self.server_hello.session_id,
            msg_id,
            xml.len()
        );
        trace!(
            "session {}: rpc content: {}",
            self.server_hello.session_id, 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 {
            debug!(
                "session {}: send failed for message-id={}: {}",
                self.server_hello.session_id, msg_id, e
            );
            self.inner.pending.lock().unwrap().remove(&msg_id);
            return Err(e);
        }
        Ok(RpcFuture {
            rx,
            msg_id,
            rpc_timeout: self.rpc_timeout,
        })
    }

    /// Send a raw RPC and wait for the reply.
    pub async fn rpc_raw(&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(&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,
            rpc_timeout: self.rpc_timeout,
        })
    }

    /// Retrieve configuration from a datastore.
    pub async fn get_config(
        &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(
        &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(&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(&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(&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(&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(&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(&self) -> crate::Result<()> {
        let reply = self.rpc_raw("<commit/>").await?;
        reply_to_ok(reply)
    }

    /// Gracefully close the NETCONF session.
    pub async fn close_session(&self) -> crate::Result<()> {
        // Atomically transition Ready → Closing. If another caller already
        // moved us out of Ready, we fail immediately.
        let prev = self.inner.state.compare_exchange(
            SessionState::Ready as u8,
            SessionState::Closing as u8,
            Ordering::AcqRel,
            Ordering::Acquire,
        );
        if let Err(current) = prev {
            let state = SessionState::from_u8(current);
            return Err(crate::Error::InvalidState(state.to_string()));
        }
        debug!("session {}: closing", self.server_hello.session_id);
        let result = self.rpc_send_unchecked("<close-session/>").await;
        match result {
            Ok(future) => {
                let reply = future.response().await;
                self.inner.set_state(SessionState::Closed);
                debug!(
                    "session {}: closed gracefully",
                    self.server_hello.session_id
                );
                reply_to_ok(reply?)
            }
            Err(e) => {
                self.inner.set_state(SessionState::Closed);
                debug!(
                    "session {}: close failed: {}",
                    self.server_hello.session_id, e
                );
                Err(e)
            }
        }
    }

    /// Force-close another NETCONF session.
    pub async fn kill_session(&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)
    }

    /// Return a wrapper that applies `timeout` to every RPC sent through it,
    /// overriding the session-level `rpc_timeout`.
    pub fn with_timeout(&self, timeout: Duration) -> SessionWithTimeout<'_> {
        SessionWithTimeout {
            session: self,
            timeout,
        }
    }

    /// Number of RPCs that have been sent but not yet replied to.
    pub fn pending_rpc_count(&self) -> usize {
        self.inner.pending.lock().unwrap().len()
    }

    /// The [`Instant`] when the most recent RPC reply was received, or
    /// `None` if no reply has been received yet.
    pub fn last_rpc_at(&self) -> Option<Instant> {
        let nanos = self.inner.last_rpc_nanos.load(Ordering::Acquire);
        if nanos == 0 {
            None
        } else {
            Some(self.inner.created_at + Duration::from_nanos(nanos))
        }
    }

    /// The [`Instant`] when the session was established (hello exchange
    /// complete and reader task started).
    pub fn connected_since(&self) -> Instant {
        self.connected_since
    }
}

/// A wrapper around [`Session`] that applies a per-call timeout to every RPC.
///
/// Created by [`Session::with_timeout()`]. Each method sends the RPC via
/// the underlying session and awaits the reply with the configured timeout.
pub struct SessionWithTimeout<'a> {
    session: &'a Session,
    timeout: Duration,
}

impl SessionWithTimeout<'_> {
    /// Send a raw RPC and wait for the reply with the configured timeout.
    pub async fn rpc_raw(&self, inner_xml: &str) -> crate::Result<RpcReply> {
        let future = self.session.rpc_send(inner_xml).await?;
        future.response_with_timeout(self.timeout).await
    }

    /// Retrieve configuration from a datastore.
    pub async fn get_config(
        &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`.
    pub async fn get_config_payload(
        &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(&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`.
    pub async fn get_payload(&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(&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(&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(&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(&self) -> crate::Result<()> {
        let reply = self.rpc_raw("<commit/>").await?;
        reply_to_ok(reply)
    }
}

/// Perform the NETCONF hello exchange, optionally with a timeout.
async fn exchange_hello<S: AsyncRead + AsyncWrite + Unpin>(
    stream: &mut S,
    config: &Config,
) -> crate::Result<(ServerHello, FramingMode)> {
    let fut = crate::hello::exchange(stream, config.codec.max_message_size);
    match config.hello_timeout {
        Some(duration) => tokio::time::timeout(duration, fut)
            .await
            .map_err(|_| crate::Error::Transport(TransportError::Timeout(duration)))?,
        None => fut.await,
    }
}

/// 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>,
    session_id: u32,
) {
    debug!("session {}: reader loop started", session_id);
    // Track the disconnect reason so we can notify subscribers when we exit.
    let mut disconnect_reason = DisconnectReason::Eof;
    // 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) => {
                trace!(
                    "session {}: received frame ({} bytes)",
                    session_id,
                    bytes.len()
                );
                match message::classify_message(bytes) {
                    Ok(ServerMessage::RpcReply(reply)) => {
                        debug!(
                            "session {}: received rpc-reply message-id={}",
                            session_id, reply.message_id
                        );
                        // 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 {
                            let nanos = inner.created_at.elapsed().as_nanos() as u64;
                            inner.last_rpc_nanos.store(nanos, Ordering::Release);
                            // we can ignore if the receiver was dropped. Nothing to do about it.
                            let _ = tx.send(Ok(reply));
                        } else {
                            warn!(
                                "session {}: received reply for unknown message-id {}",
                                session_id, reply.message_id
                            );
                        }
                    }
                    Err(e) => {
                        warn!("session {}: failed to classify message: {e}", session_id);
                    }
                }
            }
            // 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) => {
                debug!("session {}: reader error: {e}", session_id);
                disconnect_reason = DisconnectReason::TransportError(e.to_string());
                let drained = inner.drain_pending();
                if drained > 0 {
                    debug!(
                        "session {}: drained {} pending RPCs after error",
                        session_id, drained
                    );
                }
                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 drained = inner.drain_pending();
        if drained > 0 {
            debug!(
                "session {}: drained {} pending RPCs on stream close",
                session_id, drained
            );
        }
    }
    // 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);
    // Notify all disconnect subscribers. Late callers of disconnected() will
    // see the value immediately via the watch channel's last-value semantics.
    let _ = inner.disconnect_tx.send(Some(disconnect_reason));
    debug!("session {}: reader loop ended", session_id);
}

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