diff --git a/crates/ironrdp-dvc/src/client.rs b/crates/ironrdp-dvc/src/client.rs index c38330720..22ca79bfb 100644 --- a/crates/ironrdp-dvc/src/client.rs +++ b/crates/ironrdp-dvc/src/client.rs @@ -161,6 +161,10 @@ impl DrdynvcClient { self.dynamic_channels.get_by_channel_id(channel_id) } + pub fn get_dvc_by_channel_id_mut(&mut self, channel_id: u32) -> Option<&mut DynamicVirtualChannel> { + self.dynamic_channels.get_by_channel_id_mut(channel_id) + } + fn create_capabilities_response(&mut self, server_version: CapsVersion) -> SvcMessage { let caps_response = DrdynvcClientPdu::Capabilities(CapabilitiesResponsePdu::new(server_version)); debug!("Send DVC Capabilities Response PDU: {caps_response:?}"); diff --git a/crates/ironrdp-dvc/src/lib.rs b/crates/ironrdp-dvc/src/lib.rs index 6126653b0..9fd4527b2 100644 --- a/crates/ironrdp-dvc/src/lib.rs +++ b/crates/ironrdp-dvc/src/lib.rs @@ -163,5 +163,53 @@ impl DynamicVirtualChannel { } } +#[derive(Debug, Clone, Copy)] +pub struct DynamicChannelRef<'a, T> { + channel_id: DynamicChannelId, + processor: &'a T, +} + +impl DynamicChannelRef<'_, T> { + pub fn channel_id(&self) -> DynamicChannelId { + self.channel_id + } +} + +impl<'a, T: DvcProcessor> DynamicChannelRef<'a, T> { + fn new(channel_id: u32, processor: &'a T) -> Self { + Self { channel_id, processor } + } + + pub fn processor(&self) -> &'a T { + self.processor + } +} + +#[derive(Debug)] +pub struct DynamicChannelMut<'a, T> { + channel_id: DynamicChannelId, + processor: &'a mut T, +} + +impl DynamicChannelMut<'_, T> { + pub fn channel_id(&self) -> DynamicChannelId { + self.channel_id + } +} + +impl<'a, T: DvcProcessor> DynamicChannelMut<'a, T> { + fn new(channel_id: u32, processor: &'a mut T) -> Self { + Self { channel_id, processor } + } + + pub fn processor(&self) -> &T { + self.processor + } + + pub fn processor_mut(&mut self) -> &mut T { + self.processor + } +} + pub type DynamicChannelName = String; pub type DynamicChannelId = u32; diff --git a/crates/ironrdp-dvc/src/server.rs b/crates/ironrdp-dvc/src/server.rs index 21e905e68..b3d54bba9 100644 --- a/crates/ironrdp-dvc/src/server.rs +++ b/crates/ironrdp-dvc/src/server.rs @@ -14,7 +14,7 @@ use tracing::debug; use crate::pdu::{ CapabilitiesRequestPdu, CapsVersion, ClosePdu, CreateRequestPdu, CreationStatus, DrdynvcClientPdu, DrdynvcServerPdu, }; -use crate::{CompleteData, DvcProcessor, encode_dvc_messages}; +use crate::{CompleteData, DvcProcessor, DynamicChannelMut, DynamicChannelRef, encode_dvc_messages}; pub trait DvcServerProcessor: DvcProcessor {} @@ -186,6 +186,30 @@ impl DrdynvcServer { .ok_or_else(|| invalid_field_err!("DRDYNVC", "", "invalid channel id")) } + pub fn dvc_by_id(&self, id: u32) -> Option> { + 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)) + } + + pub fn dvc_by_id_mut(&mut self, id: u32) -> Option> { + 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