ironrdp_dvc/
lib.rs

1#![cfg_attr(doc, doc = include_str!("../README.md"))]
2#![doc(html_logo_url = "https://cdnweb.devolutions.net/images/projects/devolutions/logos/devolutions-icon-shadow.svg")]
3#![cfg_attr(not(feature = "std"), no_std)]
4
5extern crate alloc;
6
7use alloc::boxed::Box;
8use alloc::collections::BTreeMap;
9use alloc::string::String;
10use alloc::vec::Vec;
11use core::any::TypeId;
12
13use pdu::DrdynvcDataPdu;
14
15use crate::alloc::borrow::ToOwned as _;
16// Re-export ironrdp_pdu crate for convenience
17#[rustfmt::skip] // do not re-order this pub use
18pub use ironrdp_pdu;
19use ironrdp_core::{assert_obj_safe, cast_length, encode_vec, other_err, AsAny, Encode, EncodeResult};
20use ironrdp_pdu::{decode_err, pdu_other_err, PduResult};
21use ironrdp_svc::SvcMessage;
22
23mod complete_data;
24use complete_data::CompleteData;
25
26mod client;
27pub use client::*;
28
29mod server;
30pub use server::*;
31
32pub mod pdu;
33
34/// Represents a message that, when encoded, forms a complete PDU for a given dynamic virtual channel.
35/// This means a message that is ready to be wrapped in [`pdu::DataFirstPdu`] and [`pdu::DataPdu`] PDUs
36/// (being split into multiple of such PDUs if necessary).
37pub trait DvcEncode: Encode + Send {}
38pub type DvcMessage = Box<dyn DvcEncode>;
39
40/// A type that is a Dynamic Virtual Channel (DVC)
41///
42/// Dynamic virtual channels may be created at any point during the RDP session.
43/// The Dynamic Virtual Channel APIs exist to address limitations of Static Virtual Channels:
44///   - Limited number of channels
45///   - Packet reconstruction
46pub trait DvcProcessor: AsAny + Send {
47    /// The name of the channel, e.g. "Microsoft::Windows::RDS::DisplayControl"
48    fn channel_name(&self) -> &str;
49
50    /// Returns any messages that should be sent immediately
51    /// upon the channel being created.
52    fn start(&mut self, channel_id: u32) -> PduResult<Vec<DvcMessage>>;
53
54    fn process(&mut self, channel_id: u32, payload: &[u8]) -> PduResult<Vec<DvcMessage>>;
55
56    fn close(&mut self, _channel_id: u32) {}
57}
58
59assert_obj_safe!(DvcProcessor);
60
61pub fn encode_dvc_messages(
62    channel_id: u32,
63    messages: Vec<DvcMessage>,
64    flags: ironrdp_svc::ChannelFlags,
65) -> EncodeResult<Vec<SvcMessage>> {
66    let mut res = Vec::new();
67    for msg in messages {
68        let total_length = msg.size();
69        let needs_splitting = total_length >= DrdynvcDataPdu::MAX_DATA_SIZE;
70
71        let msg = encode_vec(msg.as_ref())?;
72        let mut off = 0;
73
74        while off < total_length {
75            let first = off == 0;
76            let remaining_length = total_length.checked_sub(off).unwrap();
77            let size = core::cmp::min(remaining_length, DrdynvcDataPdu::MAX_DATA_SIZE);
78            let end = off
79                .checked_add(size)
80                .ok_or_else(|| other_err!("encode_dvc_messages", "overflow occurred"))?;
81
82            let pdu = if needs_splitting && first {
83                DrdynvcDataPdu::DataFirst(pdu::DataFirstPdu::new(
84                    channel_id,
85                    cast_length!("total_length", total_length)?,
86                    msg[off..end].to_vec(),
87                ))
88            } else {
89                DrdynvcDataPdu::Data(pdu::DataPdu::new(channel_id, msg[off..end].to_vec()))
90            };
91
92            let svc = SvcMessage::from(pdu).with_flags(flags);
93
94            res.push(svc);
95            off = end;
96        }
97    }
98
99    Ok(res)
100}
101
102pub struct DynamicVirtualChannel {
103    channel_processor: Box<dyn DvcProcessor + Send>,
104    complete_data: CompleteData,
105    /// The channel ID assigned by the server.
106    ///
107    /// This field is `None` until the server assigns a channel ID.
108    channel_id: Option<DynamicChannelId>,
109}
110
111impl DynamicVirtualChannel {
112    fn new<T: DvcProcessor + 'static>(handler: T) -> Self {
113        Self {
114            channel_processor: Box::new(handler),
115            complete_data: CompleteData::new(),
116            channel_id: None,
117        }
118    }
119
120    pub fn is_open(&self) -> bool {
121        self.channel_id.is_some()
122    }
123
124    pub fn channel_id(&self) -> Option<DynamicChannelId> {
125        self.channel_id
126    }
127
128    pub fn channel_processor_downcast_ref<T: DvcProcessor>(&self) -> Option<&T> {
129        self.channel_processor.as_any().downcast_ref()
130    }
131
132    fn start(&mut self) -> PduResult<Vec<DvcMessage>> {
133        if let Some(channel_id) = self.channel_id {
134            self.channel_processor.start(channel_id)
135        } else {
136            Err(pdu_other_err!("DynamicVirtualChannel::start", "channel ID not set"))
137        }
138    }
139
140    fn process(&mut self, pdu: DrdynvcDataPdu) -> PduResult<Vec<DvcMessage>> {
141        let channel_id = pdu.channel_id();
142        let complete_data = self.complete_data.process_data(pdu).map_err(|e| decode_err!(e))?;
143        if let Some(complete_data) = complete_data {
144            self.channel_processor.process(channel_id, &complete_data)
145        } else {
146            Ok(Vec::new())
147        }
148    }
149
150    fn channel_name(&self) -> &str {
151        self.channel_processor.channel_name()
152    }
153}
154
155struct DynamicChannelSet {
156    channels: BTreeMap<DynamicChannelName, DynamicVirtualChannel>,
157    name_to_channel_id: BTreeMap<DynamicChannelName, DynamicChannelId>,
158    channel_id_to_name: BTreeMap<DynamicChannelId, DynamicChannelName>,
159    type_id_to_name: BTreeMap<TypeId, DynamicChannelName>,
160}
161
162impl DynamicChannelSet {
163    #[inline]
164    fn new() -> Self {
165        Self {
166            channels: BTreeMap::new(),
167            name_to_channel_id: BTreeMap::new(),
168            channel_id_to_name: BTreeMap::new(),
169            type_id_to_name: BTreeMap::new(),
170        }
171    }
172
173    fn insert<T: DvcProcessor + 'static>(&mut self, channel: T) -> Option<DynamicVirtualChannel> {
174        let name = channel.channel_name().to_owned();
175        self.type_id_to_name.insert(TypeId::of::<T>(), name.clone());
176        self.channels.insert(name, DynamicVirtualChannel::new(channel))
177    }
178
179    fn attach_channel_id(&mut self, name: DynamicChannelName, id: DynamicChannelId) -> Option<DynamicChannelId> {
180        self.channel_id_to_name.insert(id, name.clone());
181        self.name_to_channel_id.insert(name.clone(), id);
182        let dvc = self.get_by_channel_name_mut(&name)?;
183        let old_id = dvc.channel_id;
184        dvc.channel_id = Some(id);
185        old_id
186    }
187
188    fn get_by_type_id(&self, type_id: TypeId) -> Option<&DynamicVirtualChannel> {
189        self.type_id_to_name
190            .get(&type_id)
191            .and_then(|name| self.channels.get(name))
192    }
193
194    fn get_by_channel_name(&self, name: &DynamicChannelName) -> Option<&DynamicVirtualChannel> {
195        self.channels.get(name)
196    }
197
198    fn get_by_channel_name_mut(&mut self, name: &DynamicChannelName) -> Option<&mut DynamicVirtualChannel> {
199        self.channels.get_mut(name)
200    }
201
202    fn get_by_channel_id(&self, id: DynamicChannelId) -> Option<&DynamicVirtualChannel> {
203        self.channel_id_to_name
204            .get(&id)
205            .and_then(|name| self.channels.get(name))
206    }
207
208    fn get_by_channel_id_mut(&mut self, id: DynamicChannelId) -> Option<&mut DynamicVirtualChannel> {
209        self.channel_id_to_name
210            .get(&id)
211            .and_then(|name| self.channels.get_mut(name))
212    }
213
214    fn remove_by_channel_id(&mut self, id: DynamicChannelId) -> Option<DynamicChannelId> {
215        if let Some(name) = self.channel_id_to_name.remove(&id) {
216            return self.name_to_channel_id.remove(&name);
217            // Channels are retained in the `self.channels` and `self.type_id_to_name` map to allow potential
218            // dynamic re-addition by the server.
219        }
220        None
221    }
222
223    #[inline]
224    fn values(&self) -> impl Iterator<Item = &DynamicVirtualChannel> {
225        self.channels.values()
226    }
227}
228
229pub type DynamicChannelName = String;
230pub type DynamicChannelId = u32;