From fcb390140d096a32ebc1d3a87cc782e573c91f33 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc-Andr=C3=A9=20Lureau?= <marcandre.lureau@redhat.com>
Date: Tue, 25 Mar 2025 18:26:00 +0400
Subject: [PATCH] refactor(server): introduce UpdateEncoderCodecs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This will simplify setting up the UpdateEncoder with further codecs.

Signed-off-by: Marc-André Lureau <marcandre.lureau@redhat.com>
---
 crates/ironrdp-server/src/encoder/mod.rs | 38 +++++++++++++++++++-----
 crates/ironrdp-server/src/server.rs      | 10 +++----
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/crates/ironrdp-server/src/encoder/mod.rs b/crates/ironrdp-server/src/encoder/mod.rs
index b1b7a4e4..919f8903 100644
--- a/crates/ironrdp-server/src/encoder/mod.rs
+++ b/crates/ironrdp-server/src/encoder/mod.rs
@@ -29,6 +29,27 @@ enum CodecId {
     None = 0x0,
 }
 
+#[derive(Debug)]
+pub(crate) struct UpdateEncoderCodecs {
+    remotefx: Option<(EntropyBits, u8)>,
+}
+
+impl UpdateEncoderCodecs {
+    pub(crate) fn new() -> Self {
+        Self { remotefx: None }
+    }
+
+    pub(crate) fn set_remotefx(&mut self, remotefx: Option<(EntropyBits, u8)>) {
+        self.remotefx = remotefx
+    }
+}
+
+impl Default for UpdateEncoderCodecs {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 pub(crate) struct UpdateEncoder {
     desktop_size: DesktopSize,
     framebuffer: Option<Framebuffer>,
@@ -44,14 +65,17 @@ impl fmt::Debug for UpdateEncoder {
 }
 
 impl UpdateEncoder {
-    pub(crate) fn new(desktop_size: DesktopSize, surface_flags: CmdFlags, remotefx: Option<(EntropyBits, u8)>) -> Self {
-        let bitmap_updater = if !surface_flags.contains(CmdFlags::SET_SURFACE_BITS) {
-            BitmapUpdater::Bitmap(BitmapHandler::new())
-        } else if remotefx.is_some() {
-            let (algo, id) = remotefx.unwrap();
-            BitmapUpdater::RemoteFx(RemoteFxHandler::new(algo, id, desktop_size))
+    pub(crate) fn new(desktop_size: DesktopSize, surface_flags: CmdFlags, codecs: UpdateEncoderCodecs) -> Self {
+        let bitmap_updater = if surface_flags.contains(CmdFlags::SET_SURFACE_BITS) {
+            let mut bitmap = BitmapUpdater::None(NoneHandler);
+
+            if let Some((algo, id)) = codecs.remotefx {
+                bitmap = BitmapUpdater::RemoteFx(RemoteFxHandler::new(algo, id, desktop_size));
+            }
+
+            bitmap
         } else {
-            BitmapUpdater::None(NoneHandler)
+            BitmapUpdater::Bitmap(BitmapHandler::new())
         };
 
         Self {
diff --git a/crates/ironrdp-server/src/server.rs b/crates/ironrdp-server/src/server.rs
index fabf9d6e..49fa849d 100644
--- a/crates/ironrdp-server/src/server.rs
+++ b/crates/ironrdp-server/src/server.rs
@@ -30,7 +30,7 @@ use {ironrdp_dvc as dvc, ironrdp_rdpsnd as rdpsnd};
 
 use crate::clipboard::CliprdrServerFactory;
 use crate::display::{DisplayUpdate, RdpServerDisplay};
-use crate::encoder::UpdateEncoder;
+use crate::encoder::{UpdateEncoder, UpdateEncoderCodecs};
 use crate::handler::RdpServerInputHandler;
 use crate::{builder, capabilities, SoundServerFactory};
 
@@ -663,7 +663,7 @@ impl RdpServer {
             }
         }
 
-        let mut rfxcodec = None;
+        let mut update_codecs = UpdateEncoderCodecs::new();
         let mut surface_flags = CmdFlags::empty();
         for c in result.capabilities {
             match c {
@@ -714,14 +714,14 @@ impl RdpServer {
                                 rdp::capability_sets::RemoteFxContainer::ClientContainer(c),
                             ) if self.opts.with_remote_fx => {
                                 for caps in c.caps_data.0 .0 {
-                                    rfxcodec = Some((caps.entropy_bits, codec.id));
+                                    update_codecs.set_remotefx(Some((caps.entropy_bits, codec.id)));
                                 }
                             }
                             rdp::capability_sets::CodecProperty::ImageRemoteFx(
                                 rdp::capability_sets::RemoteFxContainer::ClientContainer(c),
                             ) if self.opts.with_remote_fx => {
                                 for caps in c.caps_data.0 .0 {
-                                    rfxcodec = Some((caps.entropy_bits, codec.id));
+                                    update_codecs.set_remotefx(Some((caps.entropy_bits, codec.id)));
                                 }
                             }
                             rdp::capability_sets::CodecProperty::NsCodec(_) => (),
@@ -734,7 +734,7 @@ impl RdpServer {
         }
 
         let desktop_size = self.display.lock().await.size().await;
-        let encoder = UpdateEncoder::new(desktop_size, surface_flags, rfxcodec);
+        let encoder = UpdateEncoder::new(desktop_size, surface_flags, update_codecs);
 
         let state = self
             .client_loop(reader, writer, result.io_channel_id, result.user_channel_id, encoder)