refactor(server): split UpdateEncoder

The bitmap encoder dispatching code was becoming convoluted and the same
struct was handling PduEncoding and various bitmap encoding handling.
Instead, split UpdateEncoder in different types and concerns.

Signed-off-by: Marc-André Lureau <[email protected]>

Author: elmarco

crates/ironrdp-server/src/encoder/mod.rs

@@ -28,44 +28,25 @@ const MAX_FASTPATH_UPDATE_SIZE: usize = 16_374;
 const FASTPATH_HEADER_SIZE: usize = 6;
 
 pub(crate) struct UpdateEncoder {
-    buffer: Vec<u8>,
-    bitmap: BitmapEncoder,
-    remotefx: Option<(RfxEncoder, u8)>,
-    update: for<'a> fn(&'a mut UpdateEncoder, BitmapUpdate) -> Result<UpdateFragmenter<'a>>,
+    pdu_encoder: PduEncoder,
+    bitmap_updater: BitmapUpdater,
 }
 
 impl UpdateEncoder {
     pub(crate) fn new(surface_flags: CmdFlags, remotefx: Option<(EntropyBits, u8)>) -> Self {
-        let update = if !surface_flags.contains(CmdFlags::SET_SURFACE_BITS) {
-            Self::bitmap_update
+        let pdu_encoder = PduEncoder::new();
+        let bitmap_updater = if !surface_flags.contains(CmdFlags::SET_SURFACE_BITS) {
+            BitmapUpdater::Bitmap(BitmapHandler::new())
         } else if remotefx.is_some() {
-            Self::remotefx_update
+            let (algo, id) = remotefx.unwrap();
+            BitmapUpdater::RemoteFx(RemoteFxHandler::new(algo, id))
         } else {
-            Self::none_update
+            BitmapUpdater::None(NoneHandler)
         };
 
         Self {
-            buffer: vec![0; 16384],
-            bitmap: BitmapEncoder::new(),
-            remotefx: remotefx.map(|(algo, id)| (RfxEncoder::new(algo), id)),
-            update,
-        }
-    }
-
-    fn encode_pdu(&mut self, pdu: impl Encode) -> Result<usize> {
-        loop {
-            let mut cursor = WriteCursor::new(self.buffer.as_mut_slice());
-            match pdu.encode(&mut cursor) {
-                Err(e) => match e.kind() {
-                    ironrdp_core::EncodeErrorKind::NotEnoughBytes { .. } => {
-                        self.buffer.resize(self.buffer.len() * 2, 0);
-                        debug!("encoder buffer resized to: {}", self.buffer.len() * 2);
-                    }
-
-                    _ => Err(e).context("PDU encode error")?,
-                },
-                Ok(()) => return Ok(cursor.pos()),
-            }
+            pdu_encoder,
+            bitmap_updater,
         }
     }
 
@@ -88,8 +69,8 @@ impl UpdateEncoder {
             xor_bpp: 32,
             color_pointer,
         };
-        let len = self.encode_pdu(ptr)?;
-        Ok(UpdateFragmenter::new(UpdateCode::NewPointer, &self.buffer[..len]))
+        let buf = self.pdu_encoder.encode(ptr)?;
+        Ok(UpdateFragmenter::new(UpdateCode::NewPointer, buf))
     }
 
     pub(crate) fn color_pointer(&mut self, ptr: ColorPointer) -> Result<UpdateFragmenter<'_>> {
@@ -105,8 +86,8 @@ impl UpdateEncoder {
             xor_mask: &ptr.xor_mask,
             and_mask: &ptr.and_mask,
         };
-        let len = self.encode_pdu(ptr)?;
-        Ok(UpdateFragmenter::new(UpdateCode::ColorPointer, &self.buffer[..len]))
+        let buf = self.pdu_encoder.encode(ptr)?;
+        Ok(UpdateFragmenter::new(UpdateCode::ColorPointer, buf))
     }
 
     #[allow(clippy::unused_self)]
@@ -120,31 +101,77 @@ impl UpdateEncoder {
     }
 
     pub(crate) fn pointer_position(&mut self, pos: PointerPositionAttribute) -> Result<UpdateFragmenter<'_>> {
-        let len = self.encode_pdu(pos)?;
-        Ok(UpdateFragmenter::new(UpdateCode::PositionPointer, &self.buffer[..len]))
+        let buf = self.pdu_encoder.encode(pos)?;
+        Ok(UpdateFragmenter::new(UpdateCode::PositionPointer, buf))
     }
 
     pub(crate) fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter<'_>> {
-        let update = self.update;
-
-        update(self, bitmap)
+        self.bitmap_updater.handle(bitmap, &mut self.pdu_encoder)
     }
 
     pub(crate) fn fragmenter_from_owned(&self, res: UpdateFragmenterOwned) -> UpdateFragmenter<'_> {
         UpdateFragmenter {
             code: res.code,
             index: res.index,
-            data: &self.buffer[0..res.len],
+            data: &self.pdu_encoder.buffer[0..res.len],
         }
     }
+}
 
-    fn bitmap_update(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter<'_>> {
+enum BitmapUpdater {
+    None(NoneHandler),
+    Bitmap(BitmapHandler),
+    RemoteFx(RemoteFxHandler),
+}
+
+impl BitmapUpdater {
+    fn handle<'a>(&mut self, bitmap: BitmapUpdate, encoder: &'a mut PduEncoder) -> Result<UpdateFragmenter<'a>> {
+        match self {
+            Self::None(up) => up.handle(bitmap, encoder),
+            Self::Bitmap(up) => up.handle(bitmap, encoder),
+            Self::RemoteFx(up) => up.handle(bitmap, encoder),
+        }
+    }
+}
+
+trait BitmapUpdateHandler {
+    fn handle<'a>(&mut self, bitmap: BitmapUpdate, encoder: &'a mut PduEncoder) -> Result<UpdateFragmenter<'a>>;
+}
+
+struct NoneHandler;
+
+impl BitmapUpdateHandler for NoneHandler {
+    fn handle<'a>(&mut self, bitmap: BitmapUpdate, encoder: &'a mut PduEncoder) -> Result<UpdateFragmenter<'a>> {
+        let stride = usize::from(bitmap.format.bytes_per_pixel()) * usize::from(bitmap.width.get());
+        let mut data = Vec::with_capacity(stride * usize::from(bitmap.height.get()));
+        for row in bitmap.data.chunks(bitmap.stride).rev() {
+            data.extend_from_slice(&row[..stride]);
+        }
+
+        encoder.set_surface(bitmap, CodecId::None as u8, &data)
+    }
+}
+
+struct BitmapHandler {
+    bitmap: BitmapEncoder,
+}
+
+impl BitmapHandler {
+    fn new() -> Self {
+        Self {
+            bitmap: BitmapEncoder::new(),
+        }
+    }
+}
+
+impl BitmapUpdateHandler for BitmapHandler {
+    fn handle<'a>(&mut self, bitmap: BitmapUpdate, encoder: &'a mut PduEncoder) -> Result<UpdateFragmenter<'a>> {
         let len = loop {
-            match self.bitmap.encode(&bitmap, self.buffer.as_mut_slice()) {
+            match self.bitmap.encode(&bitmap, encoder.buffer.as_mut_slice()) {
                 Err(e) => match e.kind() {
                     ironrdp_core::EncodeErrorKind::NotEnoughBytes { .. } => {
-                        self.buffer.resize(self.buffer.len() * 2, 0);
-                        debug!("encoder buffer resized to: {}", self.buffer.len() * 2);
+                        encoder.buffer.resize(encoder.buffer.len() * 2, 0);
+                        debug!("encoder buffer resized to: {}", encoder.buffer.len() * 2);
                     }
 
                     _ => Err(e).context("bitmap encode error")?,
@@ -153,7 +180,71 @@ impl UpdateEncoder {
             }
         };
 
-        Ok(UpdateFragmenter::new(UpdateCode::Bitmap, &self.buffer[..len]))
+        Ok(UpdateFragmenter::new(UpdateCode::Bitmap, &encoder.buffer[..len]))
+    }
+}
+
+struct RemoteFxHandler {
+    remotefx: RfxEncoder,
+    codec_id: u8,
+}
+
+impl RemoteFxHandler {
+    fn new(algo: EntropyBits, codec_id: u8) -> Self {
+        Self {
+            remotefx: RfxEncoder::new(algo),
+            codec_id,
+        }
+    }
+}
+
+impl BitmapUpdateHandler for RemoteFxHandler {
+    fn handle<'a>(&mut self, bitmap: BitmapUpdate, encoder: &'a mut PduEncoder) -> Result<UpdateFragmenter<'a>> {
+        let mut buffer = vec![0; bitmap.data.len()];
+        let len = loop {
+            match self.remotefx.encode(&bitmap, buffer.as_mut_slice()) {
+                Err(e) => match e.kind() {
+                    ironrdp_core::EncodeErrorKind::NotEnoughBytes { .. } => {
+                        buffer.resize(buffer.len() * 2, 0);
+                        debug!("encoder buffer resized to: {}", buffer.len() * 2);
+                    }
+
+                    _ => Err(e).context("RemoteFX encode error")?,
+                },
+                Ok(len) => break len,
+            }
+        };
+
+        encoder.set_surface(bitmap, self.codec_id, &buffer[..len])
+    }
+}
+
+struct PduEncoder {
+    buffer: Vec<u8>,
+}
+
+impl PduEncoder {
+    fn new() -> Self {
+        Self { buffer: vec![0; 16384] }
+    }
+
+    fn encode(&mut self, pdu: impl Encode) -> Result<&[u8]> {
+        let pos = loop {
+            let mut cursor = WriteCursor::new(self.buffer.as_mut_slice());
+            match pdu.encode(&mut cursor) {
+                Err(e) => match e.kind() {
+                    ironrdp_core::EncodeErrorKind::NotEnoughBytes { .. } => {
+                        self.buffer.resize(self.buffer.len() * 2, 0);
+                        debug!("encoder buffer resized to: {}", self.buffer.len() * 2);
+                    }
+
+                    _ => Err(e).context("PDU encode error")?,
+                },
+                Ok(()) => break cursor.pos(),
+            }
+        };
+
+        Ok(&self.buffer[..pos])
     }
 
     fn set_surface(&mut self, bitmap: BitmapUpdate, codec_id: u8, data: &[u8]) -> Result<UpdateFragmenter<'_>> {
@@ -176,39 +267,8 @@ impl UpdateEncoder {
             extended_bitmap_data,
         };
         let cmd = SurfaceCommand::SetSurfaceBits(pdu);
-        let len = self.encode_pdu(cmd)?;
-        Ok(UpdateFragmenter::new(UpdateCode::SurfaceCommands, &self.buffer[..len]))
-    }
-
-    fn remotefx_update(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter<'_>> {
-        let (remotefx, codec_id) = self.remotefx.as_mut().unwrap();
-        let codec_id = *codec_id;
-        let mut buffer = vec![0; bitmap.data.len()];
-        let len = loop {
-            match remotefx.encode(&bitmap, buffer.as_mut_slice()) {
-                Err(e) => match e.kind() {
-                    ironrdp_core::EncodeErrorKind::NotEnoughBytes { .. } => {
-                        buffer.resize(buffer.len() * 2, 0);
-                        debug!("encoder buffer resized to: {}", self.buffer.len() * 2);
-                    }
-
-                    _ => Err(e).context("RemoteFX encode error")?,
-                },
-                Ok(len) => break len,
-            }
-        };
-
-        self.set_surface(bitmap, codec_id, &buffer[..len])
-    }
-
-    fn none_update(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter<'_>> {
-        let stride = usize::from(bitmap.format.bytes_per_pixel()) * usize::from(bitmap.width.get());
-        let mut data = Vec::with_capacity(stride * usize::from(bitmap.height.get()));
-        for row in bitmap.data.chunks(bitmap.stride).rev() {
-            data.extend_from_slice(&row[..stride]);
-        }
-
-        self.set_surface(bitmap, CodecId::None as u8, &data)
+        let buf = self.encode(cmd)?;
+        Ok(UpdateFragmenter::new(UpdateCode::SurfaceCommands, buf))
     }
 }