rat-rdp-dvc 0.1.0

use alloc::boxed::Box;
use alloc::collections::{BTreeMap, BTreeSet};
use alloc::vec::Vec;
use core::any::TypeId;
use core::fmt;

use rat_rdp_core::{Decode as _, DecodeResult, ReadCursor, impl_as_any, invalid_field_err};
use rat_rdp_pdu::{self as pdu, PduError, decode_err, encode_err, pdu_other_err};
use rat_rdp_svc::{ChannelFlags, CompressionCondition, SvcMessage, SvcProcessor, SvcServerProcessor};
use pdu::PduResult;
use pdu::gcc::ChannelName;
use tracing::debug;

use crate::pdu::{
    CapabilitiesRequestPdu, CapsVersion, ClosePdu, CreateRequestPdu, CreationStatus, DrdynvcClientPdu,
    DrdynvcServerPdu, SoftSyncChannelList, SoftSyncRequestPdu, SoftSyncTunnelType,
};
use crate::{CompleteData, DvcProcessor, DynamicChannelMut, DynamicChannelRef, encode_dvc_messages};

pub trait DvcServerProcessor: DvcProcessor {}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum ChannelState {
    Pending,
    /// `Create Request` has been sent; awaiting `Create Response` from the client.
    Creation,
    Opened,
    CreationFailed(u32),
}

enum SoftSyncState {
    Idle,
    Active {
        requested_tunnels: BTreeSet<SoftSyncTunnelType>,
        response_received: bool,
    },
}

struct DynamicChannel {
    state: ChannelState,
    processor: Box<dyn DvcServerProcessor>,
    complete_data: CompleteData,
    channel_id: u32,
}

impl Drop for DynamicChannel {
    fn drop(&mut self) {
        if self.state == ChannelState::Opened {
            self.processor.close(self.channel_id);
        }
    }
}

struct DynamicChannelAllocator {
    dynamic_channels: BTreeMap<u32, DynamicChannel>,
    next_channel_id: u32,
}

impl<'a> IntoIterator for &'a DynamicChannelAllocator {
    type Item = (&'a u32, &'a DynamicChannel);

    type IntoIter = alloc::collections::btree_map::Iter<'a, u32, DynamicChannel>;

    fn into_iter(self) -> Self::IntoIter {
        self.dynamic_channels.iter()
    }
}

impl<'a> IntoIterator for &'a mut DynamicChannelAllocator {
    type Item = (&'a u32, &'a mut DynamicChannel);
    type IntoIter = alloc::collections::btree_map::IterMut<'a, u32, DynamicChannel>;
    fn into_iter(self) -> Self::IntoIter {
        self.dynamic_channels.iter_mut()
    }
}

impl DynamicChannelAllocator {
    fn new() -> Self {
        Self {
            dynamic_channels: BTreeMap::new(),
            next_channel_id: 0,
        }
    }

    fn reserve_channel(&mut self) -> u32 {
        let channel_id = self.next_channel_id;
        self.next_channel_id = self
            .next_channel_id
            .checked_add(1)
            .expect("dynamic channels reaches `u32::MAX`");
        channel_id
    }

    fn insert_channel<T>(&mut self, processor: T, state: ChannelState) -> u32
    where
        T: DvcServerProcessor + 'static,
    {
        let channel_id = self.reserve_channel();
        self.insert_channel_with_id(processor, state, channel_id);
        channel_id
    }

    fn insert_channel_with_id<T>(&mut self, processor: T, state: ChannelState, channel_id: u32)
    where
        T: DvcServerProcessor + 'static,
    {
        self.dynamic_channels
            .insert(channel_id, DynamicChannel::new(processor, channel_id, state));
    }

    fn get(&self, channel_id: u32) -> Option<&DynamicChannel> {
        self.dynamic_channels.get(&channel_id)
    }

    fn get_mut(&mut self, channel_id: u32) -> Option<&mut DynamicChannel> {
        self.dynamic_channels.get_mut(&channel_id)
    }

    fn remove(&mut self, channel_id: u32) -> Option<DynamicChannel> {
        self.dynamic_channels.remove(&channel_id)
    }
}

impl DynamicChannel {
    fn new<T>(processor: T, channel_id: u32, state: ChannelState) -> Self
    where
        T: DvcServerProcessor + 'static,
    {
        Self {
            state,
            processor: Box::new(processor),
            complete_data: CompleteData::new(),
            channel_id,
        }
    }

    fn processor_type_id(&self) -> TypeId {
        self.processor.as_any().type_id()
    }
}
/// DRDYNVC Static Virtual Channel (the Remote Desktop Protocol: Dynamic Virtual Channel Extension)
///
/// It adds support for dynamic virtual channels (DVC).
pub struct DrdynvcServer {
    dynamic_channels: DynamicChannelAllocator,
    type_id_to_channel_id: BTreeMap<TypeId, u32>,
    soft_sync_state: SoftSyncState,
    outgoing_tunnel_channels: BTreeMap<u32, SoftSyncTunnelType>,
    incoming_tunnel_channels: BTreeMap<u32, SoftSyncTunnelType>,
}

impl fmt::Debug for DrdynvcServer {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "DrdynvcServer([")?;

        for (i, (id, channel)) in self.dynamic_channels.into_iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "{}:{} ({:?})", id, channel.processor.channel_name(), channel.state)?;
        }

        write!(f, "])")
    }
}

impl DrdynvcServer {
    pub const NAME: ChannelName = ChannelName::from_static(b"drdynvc\0");

    pub fn new() -> Self {
        Self {
            dynamic_channels: DynamicChannelAllocator::new(),
            type_id_to_channel_id: BTreeMap::new(),
            soft_sync_state: SoftSyncState::Idle,
            outgoing_tunnel_channels: BTreeMap::new(),
            incoming_tunnel_channels: BTreeMap::new(),
        }
    }

    pub fn get_channel_id_by_type<T>(&self) -> Option<u32>
    where
        T: DvcServerProcessor + 'static,
    {
        self.type_id_to_channel_id.get(&TypeId::of::<T>()).copied()
    }

    /// Returns `true` if the DVC channel with the given ID has completed
    /// its creation handshake and is in the `Opened` state.
    pub fn is_channel_opened(&self, channel_id: u32) -> bool {
        self.dynamic_channels
            .get(channel_id)
            .is_some_and(|c| c.state == ChannelState::Opened)
    }

    /// Registers a dynamic channel with the server.
    ///
    /// # Panics
    ///
    /// Panics if the number of registered dynamic channels reaches `u32::MAX`.
    #[must_use]
    pub fn with_dynamic_channel<T>(mut self, channel: T) -> Self
    where
        T: DvcServerProcessor + 'static,
    {
        let channel_id = self.dynamic_channels.insert_channel(channel, ChannelState::Pending);
        self.type_id_to_channel_id.insert(TypeId::of::<T>(), channel_id);
        self
    }

    fn channel_by_id(&mut self, id: u32) -> DecodeResult<&mut DynamicChannel> {
        self.dynamic_channels
            .get_mut(id)
            .ok_or_else(|| invalid_field_err!("DRDYNVC", "", "invalid channel id"))
    }

    /// Returns a typed accessor for an active server DVC by channel ID.
    pub fn dvc_by_id<T: DvcServerProcessor>(&self, id: u32) -> Option<DynamicChannelRef<'_, T>> {
        let channel = self.dynamic_channels.get(id)?;
        if channel.state != ChannelState::Opened {
            return None;
        }
        channel
            .processor
            .as_any()
            .downcast_ref()
            .map(|p| DynamicChannelRef::new(id, p))
    }

    /// Returns a mutable typed accessor for an active server DVC by channel ID.
    pub fn dvc_by_id_mut<T: DvcServerProcessor>(&mut self, id: u32) -> Option<DynamicChannelMut<'_, T>> {
        let channel = self.dynamic_channels.get_mut(id)?;
        if channel.state != ChannelState::Opened {
            return None;
        }
        channel
            .processor
            .as_any_mut()
            .downcast_mut()
            .map(|p| DynamicChannelMut::new(id, p))
    }

    /// Creates a new DVC, returns CreateRequest PDU to send to client.
    ///
    /// # Panics
    ///
    /// Panics if the number of registered dynamic channels reaches `u32::MAX`.
    pub fn create_channel<T>(&mut self, channel: T) -> PduResult<SvcMessage>
    where
        T: DvcServerProcessor + 'static,
    {
        let channel_id = self.dynamic_channels.reserve_channel();
        self.create_channel_with_id(channel, channel_id)
    }

    /// Creates a new DVC using a processor built with its assigned channel ID.
    ///
    /// The next channel ID is reserved and passed to `build`, allowing the
    /// processor or one of its dependencies to use the ID during construction.
    ///
    /// # Panics
    ///
    /// Panics if the number of registered dynamic channels reaches `u32::MAX`.
    pub fn create_channel_with<T, E, F>(&mut self, build: F) -> Result<SvcMessage, E>
    where
        T: DvcServerProcessor + 'static,
        E: From<PduError>,
        F: FnOnce(u32) -> Result<T, E>,
    {
        let channel_id = self.dynamic_channels.reserve_channel();
        let channel = build(channel_id)?;
        self.create_channel_with_id(channel, channel_id).map_err(E::from)
    }

    fn create_channel_with_id<T>(&mut self, channel: T, channel_id: u32) -> PduResult<SvcMessage>
    where
        T: DvcServerProcessor + 'static,
    {
        let channel_name = channel.channel_name().into();
        let req = DrdynvcServerPdu::Create(CreateRequestPdu::new(channel_id, channel_name));
        let svc_msg = as_svc_msg_with_flag(req)?;
        self.dynamic_channels
            .insert_channel_with_id(channel, ChannelState::Creation, channel_id);
        Ok(svc_msg)
    }

    fn remove_by_channel_id(&mut self, id: u32) -> Option<DynamicChannel> {
        self.dynamic_channels.remove(id).inspect(|dvc| {
            let type_id = dvc.processor_type_id();

            // Only matters for pre-registered channels
            if let alloc::collections::btree_map::Entry::Occupied(entry) = self.type_id_to_channel_id.entry(type_id)
                && entry.get() == &id
            {
                entry.remove();
            }
        })
    }

    pub fn close_channel(&mut self, channel_id: u32) -> Option<SvcMessage> {
        self.remove_by_channel_id(channel_id)?;
        self.outgoing_tunnel_channels.remove(&channel_id);
        self.incoming_tunnel_channels.remove(&channel_id);
        Some(
            SvcMessage::from(DrdynvcServerPdu::Close(ClosePdu::new(channel_id)))
                .with_flags(ChannelFlags::SHOW_PROTOCOL),
        )
    }

    /// Creates a Soft-Sync request that moves the supplied channels to reliable UDP.
    ///
    /// This API emits exactly one `ReliableUdp` channel list and maps every supplied
    /// channel to that list. A future multi-tunnel request API must establish an
    /// explicit response-routing mapping before it is exposed.
    pub fn request_reliable_udp(&mut self, channel_ids: Vec<u32>) -> PduResult<SvcMessage> {
        if channel_ids.is_empty() {
            return Err(pdu_other_err!("soft-sync requires at least one dynamic channel"));
        }
        if !matches!(self.soft_sync_state, SoftSyncState::Idle) {
            return Err(pdu_other_err!("soft-sync has already been requested"));
        }

        let mut selected_channels = BTreeMap::new();
        for channel_id in &channel_ids {
            if !self.is_channel_opened(*channel_id) {
                return Err(pdu_other_err!(
                    "Soft-Sync requested for a dynamic channel that is not open"
                ));
            }
            if selected_channels
                .insert(*channel_id, SoftSyncTunnelType::RELIABLE_UDP)
                .is_some()
            {
                return Err(pdu_other_err!("soft-sync channel list contains a duplicate channel ID"));
            }
        }

        let request = SoftSyncRequestPdu::new(alloc::vec![SoftSyncChannelList::new(
            SoftSyncTunnelType::RELIABLE_UDP,
            channel_ids,
        )]);
        let message = as_svc_msg_with_flag(DrdynvcServerPdu::SoftSyncRequest(request))?;
        self.outgoing_tunnel_channels = selected_channels;
        self.soft_sync_state = SoftSyncState::Active {
            requested_tunnels: BTreeSet::from([SoftSyncTunnelType::RELIABLE_UDP]),
            response_received: false,
        };
        Ok(message)
    }

    /// Returns whether server-to-client data for `channel_id` must be sent through a tunnel.
    pub fn tunnel_for_outgoing_channel(&self, channel_id: u32) -> Option<SoftSyncTunnelType> {
        self.outgoing_tunnel_channels.get(&channel_id).copied()
    }

    /// Returns whether the client has acknowledged the Soft-Sync request over TCP.
    pub const fn soft_sync_response_received(&self) -> bool {
        matches!(
            self.soft_sync_state,
            SoftSyncState::Active {
                response_received: true,
                ..
            }
        )
    }

    /// Processes raw DRDYNVC data received through an established multitransport tunnel.
    pub fn process_tunnel(&mut self, payload: &[u8]) -> PduResult<Vec<SvcMessage>> {
        let pdu = decode_dvc_message(payload).map_err(|e| decode_err!(e))?;
        let DrdynvcClientPdu::Data(data) = pdu else {
            return Err(pdu_other_err!("only DVC data is permitted on a multitransport tunnel"));
        };
        if !self.incoming_tunnel_channels.contains_key(&data.channel_id()) {
            return Err(pdu_other_err!(
                "received tunneled data for a channel not selected by Soft-Sync"
            ));
        }
        self.process_data(data)
    }

    fn process_data(&mut self, data: crate::pdu::DrdynvcDataPdu) -> PduResult<Vec<SvcMessage>> {
        let channel_id = data.channel_id();
        let c = self.channel_by_id(channel_id).map_err(|e| decode_err!(e))?;
        if c.state != ChannelState::Opened {
            debug!(?channel_id, ?c.state, "Invalid channel state");
            return Err(pdu_other_err!("invalid channel state"));
        }
        let mut resp = Vec::new();
        if let Some(complete) = c.complete_data.process_data(data).map_err(|e| decode_err!(e))? {
            let msg = c.processor.process(channel_id, &complete)?;
            resp.extend(encode_dvc_messages(channel_id, msg, ChannelFlags::SHOW_PROTOCOL).map_err(|e| encode_err!(e))?);
        }
        Ok(resp)
    }

    fn process_soft_sync_response(&mut self, response: crate::pdu::SoftSyncResponsePdu) -> PduResult<()> {
        let SoftSyncState::Active {
            requested_tunnels,
            response_received,
        } = &mut self.soft_sync_state
        else {
            return Err(pdu_other_err!("received unexpected Soft-Sync response"));
        };
        if *response_received {
            return Err(pdu_other_err!("received duplicate Soft-Sync response"));
        }
        for tunnel_type in response.tunnels_to_switch() {
            if !requested_tunnels.contains(tunnel_type) {
                return Err(pdu_other_err!("soft-sync response selected an unrequested tunnel"));
            }
        }
        self.incoming_tunnel_channels = self
            .outgoing_tunnel_channels
            .iter()
            .filter(|(_, tunnel_type)| response.tunnels_to_switch().contains(tunnel_type))
            .map(|(channel_id, tunnel_type)| (*channel_id, *tunnel_type))
            .collect();
        *response_received = true;
        Ok(())
    }
}

impl_as_any!(DrdynvcServer);

impl Default for DrdynvcServer {
    fn default() -> Self {
        Self::new()
    }
}

impl SvcProcessor for DrdynvcServer {
    fn channel_name(&self) -> ChannelName {
        DrdynvcServer::NAME
    }

    fn compression_condition(&self) -> CompressionCondition {
        CompressionCondition::WhenRdpDataIsCompressed
    }

    fn start(&mut self) -> PduResult<Vec<SvcMessage>> {
        let cap = CapabilitiesRequestPdu::new(CapsVersion::V2, None);
        let req = DrdynvcServerPdu::Capabilities(cap);
        let msg = as_svc_msg_with_flag(req)?;
        Ok(alloc::vec![msg])
    }

    fn process(&mut self, payload: &[u8]) -> PduResult<Vec<SvcMessage>> {
        let pdu = decode_dvc_message(payload).map_err(|e| decode_err!(e))?;
        let mut resp = Vec::new();

        match pdu {
            DrdynvcClientPdu::Capabilities(caps_resp) => {
                debug!("Got DVC Capabilities Response PDU: {caps_resp:?}");
                for (id, c) in &mut self.dynamic_channels {
                    if c.state != ChannelState::Pending {
                        continue;
                    }
                    let req = DrdynvcServerPdu::Create(CreateRequestPdu::new(*id, c.processor.channel_name().into()));
                    c.state = ChannelState::Creation;
                    resp.push(as_svc_msg_with_flag(req)?);
                }
            }
            DrdynvcClientPdu::Create(create_resp) => {
                debug!("Got DVC Create Response PDU: {create_resp:?}");
                let id = create_resp.channel_id();
                let c = self.channel_by_id(id).map_err(|e| decode_err!(e))?;
                if c.state != ChannelState::Creation {
                    return Err(pdu_other_err!("invalid channel state"));
                }
                if create_resp.creation_status() != CreationStatus::OK {
                    c.state = ChannelState::CreationFailed(create_resp.creation_status().into());
                    return Ok(resp);
                }
                c.state = ChannelState::Opened;
                let msg = c.processor.start(create_resp.channel_id())?;
                resp.extend(encode_dvc_messages(id, msg, ChannelFlags::SHOW_PROTOCOL).map_err(|e| encode_err!(e))?);
            }
            DrdynvcClientPdu::Close(close) => {
                debug!("Got DVC Close PDU: {close:?}");
                let channel_id = close.channel_id();
                self.remove_by_channel_id(channel_id);
            }
            DrdynvcClientPdu::Data(data) => {
                if self.incoming_tunnel_channels.contains_key(&data.channel_id()) {
                    return Err(pdu_other_err!("received TCP data for a channel selected by Soft-Sync"));
                }
                resp.extend(self.process_data(data)?);
            }
            DrdynvcClientPdu::SoftSyncResponse(response) => {
                debug!("Got DVC Soft-Sync Response PDU: {response:?}");
                self.process_soft_sync_response(response)?;
            }
        }

        Ok(resp)
    }
}

impl SvcServerProcessor for DrdynvcServer {}

fn decode_dvc_message(user_data: &[u8]) -> DecodeResult<DrdynvcClientPdu> {
    DrdynvcClientPdu::decode(&mut ReadCursor::new(user_data))
}

fn as_svc_msg_with_flag(pdu: DrdynvcServerPdu) -> PduResult<SvcMessage> {
    Ok(SvcMessage::from(pdu).with_flags(ChannelFlags::SHOW_PROTOCOL))
}

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

    struct TestDvc;

    impl_as_any!(TestDvc);

    impl DvcProcessor for TestDvc {
        fn channel_name(&self) -> &str {
            "test"
        }

        fn start(&mut self, _channel_id: u32) -> PduResult<Vec<crate::DvcMessage>> {
            Ok(Vec::new())
        }

        fn process(&mut self, _channel_id: u32, _payload: &[u8]) -> PduResult<Vec<crate::DvcMessage>> {
            Ok(Vec::new())
        }
    }

    impl DvcServerProcessor for TestDvc {}

    #[test]
    fn soft_sync_rejects_tunnel_data_until_the_client_responds() {
        let mut server = DrdynvcServer::new();
        let channel_id = server.dynamic_channels.insert_channel(TestDvc, ChannelState::Opened);

        server.request_reliable_udp(alloc::vec![channel_id]).unwrap();
        assert_eq!(
            server.tunnel_for_outgoing_channel(channel_id),
            Some(SoftSyncTunnelType::RELIABLE_UDP)
        );

        let tunnel_data = rat_rdp_core::encode_vec(&DrdynvcClientPdu::Data(crate::pdu::DrdynvcDataPdu::Data(
            crate::pdu::DataPdu::new(channel_id, Vec::new()),
        )))
        .unwrap();
        assert!(server.process_tunnel(&tunnel_data).is_err());

        server
            .process_soft_sync_response(crate::pdu::SoftSyncResponsePdu::new(alloc::vec![
                SoftSyncTunnelType::RELIABLE_UDP,
            ]))
            .unwrap();

        assert!(server.soft_sync_response_received());
        assert!(server.process_tunnel(&tunnel_data).is_ok());

        server.close_channel(channel_id).unwrap();
        assert!(server.request_reliable_udp(alloc::vec![channel_id]).is_err());
    }

    #[test]
    fn soft_sync_accepts_a_response_after_the_selected_channel_closes() {
        let mut server = DrdynvcServer::new();
        let channel_id = server.dynamic_channels.insert_channel(TestDvc, ChannelState::Opened);

        server.request_reliable_udp(alloc::vec![channel_id]).unwrap();
        server.close_channel(channel_id).unwrap();

        server
            .process_soft_sync_response(crate::pdu::SoftSyncResponsePdu::new(alloc::vec![
                SoftSyncTunnelType::RELIABLE_UDP,
            ]))
            .unwrap();

        assert!(server.soft_sync_response_received());
        assert!(server.request_reliable_udp(alloc::vec![channel_id]).is_err());
    }
}