From 704084952f4513fd434fb5ecd2fd79f51c8cf408 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Lureau?= Date: Thu, 28 Mar 2024 21:07:00 +0400 Subject: [PATCH] refactor(server): convert encoder to use anyhow::Result MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This is a bit more idiomatic than returning None for errors, and allows to push up the reporting. Signed-off-by: Marc-André Lureau --- crates/ironrdp-server/src/encoder/mod.rs | 44 ++++++++---------------- crates/ironrdp-server/src/server.rs | 26 ++++++++------ 2 files changed, 31 insertions(+), 39 deletions(-) diff --git a/crates/ironrdp-server/src/encoder/mod.rs b/crates/ironrdp-server/src/encoder/mod.rs index 9b46854c..1afba2dc 100644 --- a/crates/ironrdp-server/src/encoder/mod.rs +++ b/crates/ironrdp-server/src/encoder/mod.rs @@ -1,6 +1,7 @@ pub(crate) mod bitmap; pub(crate) mod rfx; +use anyhow::{Context, Result}; use std::{cmp, mem}; use ironrdp_pdu::cursor::WriteCursor; @@ -31,7 +32,7 @@ pub(crate) struct UpdateEncoder { buffer: Vec, bitmap: BitmapEncoder, remotefx: Option<(RfxEncoder, u8)>, - update: for<'a> fn(&'a mut UpdateEncoder, BitmapUpdate) -> Option>, + update: for<'a> fn(&'a mut UpdateEncoder, BitmapUpdate) -> Result>, } impl UpdateEncoder { @@ -52,7 +53,7 @@ impl UpdateEncoder { } } - fn encode_pdu(&mut self, pdu: impl PduEncode) -> Option { + fn encode_pdu(&mut self, pdu: impl PduEncode) -> Result { loop { let mut cursor = WriteCursor::new(self.buffer.as_mut_slice()); match pdu.encode(&mut cursor) { @@ -62,23 +63,20 @@ impl UpdateEncoder { debug!("encoder buffer resized to: {}", self.buffer.len() * 2); } - _ => { - debug!("encode error: {:?}", e); - return None; - } + _ => Err(e).context("PDU encode error")?, }, - Ok(()) => return Some(cursor.pos()), + Ok(()) => return Ok(cursor.pos()), } } } - pub(crate) fn bitmap(&mut self, bitmap: BitmapUpdate) -> Option> { + pub(crate) fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result> { let update = self.update; update(self, bitmap) } - fn bitmap_update(&mut self, bitmap: BitmapUpdate) -> Option> { + fn bitmap_update(&mut self, bitmap: BitmapUpdate) -> Result> { let len = loop { match self.bitmap.encode(&bitmap, self.buffer.as_mut_slice()) { Err(e) => match e.kind() { @@ -87,19 +85,16 @@ impl UpdateEncoder { debug!("encoder buffer resized to: {}", self.buffer.len() * 2); } - _ => { - debug!("bitmap encode error: {:?}", e); - return None; - } + _ => Err(e).context("bitmap encode error")?, }, Ok(len) => break len, } }; - return Some(UpdateFragmenter::new(UpdateCode::Bitmap, &self.buffer[..len])); + Ok(UpdateFragmenter::new(UpdateCode::Bitmap, &self.buffer[..len])) } - fn set_surface(&mut self, bitmap: BitmapUpdate, codec_id: u8, data: Vec) -> Option> { + fn set_surface(&mut self, bitmap: BitmapUpdate, codec_id: u8, data: Vec) -> Result> { let destination = ExclusiveRectangle { left: bitmap.left, top: bitmap.top, @@ -119,28 +114,19 @@ impl UpdateEncoder { extended_bitmap_data, }; let cmd = SurfaceCommand::SetSurfaceBits(pdu); - let Some(len) = self.encode_pdu(cmd) else { - return None; - }; - - Some(UpdateFragmenter::new(UpdateCode::SurfaceCommands, &self.buffer[..len])) + let len = self.encode_pdu(cmd)?; + Ok(UpdateFragmenter::new(UpdateCode::SurfaceCommands, &self.buffer[..len])) } - fn remotefx_update(&mut self, bitmap: BitmapUpdate) -> Option> { + fn remotefx_update(&mut self, bitmap: BitmapUpdate) -> Result> { let (remotefx, codec_id) = self.remotefx.as_mut().unwrap(); let codec_id = *codec_id; - let data = match remotefx.encode(&bitmap) { - Ok(data) => data, - Err(e) => { - debug!("remotefx encode error: {:?}", e); - return None; - } - }; + let data = remotefx.encode(&bitmap).context("RemoteFX encoding")?; self.set_surface(bitmap, codec_id, data) } - fn none_update(&mut self, mut bitmap: BitmapUpdate) -> Option> { + fn none_update(&mut self, mut bitmap: BitmapUpdate) -> Result> { let data = match bitmap.order { PixelOrder::BottomToTop => mem::take(&mut bitmap.data), PixelOrder::TopToBottom => { diff --git a/crates/ironrdp-server/src/server.rs b/crates/ironrdp-server/src/server.rs index 8146bbc7..c2100367 100644 --- a/crates/ironrdp-server/src/server.rs +++ b/crates/ironrdp-server/src/server.rs @@ -347,20 +347,26 @@ impl RdpServer { Some(update) = display_updates.next_update() => { let fragmenter = match update { - DisplayUpdate::Bitmap(bitmap) => encoder.bitmap(bitmap) + DisplayUpdate::Bitmap(bitmap) => encoder.bitmap(bitmap), }; - if let Some(mut fragmenter) = fragmenter { - if fragmenter.size_hint() > buffer.len() { - buffer.resize(fragmenter.size_hint(), 0); + let mut fragmenter = match fragmenter { + Ok(fragmenter) => fragmenter, + Err(error) => { + error!(?error, "Error during update encoding"); + break; } + }; - while let Some(len) = fragmenter.next(&mut buffer) { - if let Err(error) = framed.write_all(&buffer[..len]).await { - error!(?error, "Write display update error"); - break; - }; - } + if fragmenter.size_hint() > buffer.len() { + buffer.resize(fragmenter.size_hint(), 0); + } + + while let Some(len) = fragmenter.next(&mut buffer) { + if let Err(error) = framed.write_all(&buffer[..len]).await { + error!(?error, "Write display update error"); + break; + }; } } }