rework Service accept

airbus-seclab · Feb 7, 2025 · 6feaf9d · 6feaf9d
1 parent 382f46d
commit 6feaf9d
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 29 deletions.
diff --git a/README.md b/README.md
@@ -347,8 +347,10 @@ and the `backend`.
 ### Steps
 
 1. Update `common/src/api.rs`:
-   - Define a unique name: `const SERVICE_PING: &str = "ping";`
-   - Add a `Ping` variant to the `Service` enum;
+   - Define a unique name: `const SERVICE_PING: &str = "ping";`,
+   - Add a `Ping` variant to the `Service` enum,
+   - Complete the functions `value`, `accept` and `try_from` of the `Service`
+     enum;
 2. Declare the new module in `common/src/lib.rs`: `pub mod ping`;
 3. Create a new module structure:
 
@@ -360,7 +362,9 @@ common/src/ping/
 ```
 
 4. Implement `service::Frontend` in `common/src/ping/frontend.rs` and
-   `service::Backend` in `common/src/ping/backend.rs`.
+   `service::Backend` in `common/src/ping/backend.rs`;
+5. Add proper initialization of the frontend part of the service in the
+   `init` function in `frontend/src/lib.rs`.
 
 Refer to `common/src/clipboard/` for examples.
 

diff --git a/common/src/api.rs b/common/src/api.rs
@@ -1,3 +1,8 @@
+use crate::{
+    clipboard, command, ftp,
+    service::{self, Backend},
+    socks5, stage0,
+};
 use std::{borrow, fmt, io, sync};
 
 pub const CHUNK_LENGTH: usize = 1600; // this is the max value
@@ -31,6 +36,16 @@ impl Service {
     const fn as_bytes(self) -> &'static [u8] {
         self.value().as_bytes()
     }
+
+    pub fn accept(&self, stream: service::RdpStream) -> Result<(), io::Error> {
+        match self {
+            Self::Clipboard => clipboard::backend::Server::accept(stream),
+            Self::Command => command::backend::Server::accept(stream),
+            Self::Ftp => ftp::backend::Server::accept(stream),
+            Self::Socks5 => socks5::backend::Server::accept(stream),
+            Self::Stage0 => stage0::backend::Server::accept(stream),
+        }
+    }
 }
 
 impl<'a> TryFrom<&'a [u8]> for Service {

diff --git a/common/src/service.rs b/common/src/service.rs
@@ -1,4 +1,4 @@
-use crate::{api, clipboard, command, ftp, socks5, stage0};
+use crate::api;
 use std::{
     collections::{self, hash_map},
     io::{self, Write},
@@ -98,31 +98,9 @@ impl Channel {
 
                     thread::Builder::new()
                         .name(format!("{service_kind} {service} {client_id:x}"))
-                        .spawn_scoped(scope, move || match service {
-                            api::Service::Clipboard => {
-                                if let Err(e) = clipboard::backend::Server::accept(stream) {
-                                    crate::debug!("error: {e}");
-                                }
-                            }
-                            api::Service::Command => {
-                                if let Err(e) = command::backend::Server::accept(stream) {
-                                    crate::debug!("error: {e}");
-                                }
-                            }
-                            api::Service::Ftp => {
-                                if let Err(e) = ftp::backend::Server::accept(stream) {
-                                    crate::error!("error: {e}");
-                                }
-                            }
-                            api::Service::Socks5 => {
-                                if let Err(e) = socks5::backend::Server::accept(stream) {
-                                    crate::error!("error: {e}");
-                                }
-                            }
-                            api::Service::Stage0 => {
-                                if let Err(e) = stage0::backend::Server::accept(stream) {
-                                    crate::error!("error: {e}");
-                                }
+                        .spawn_scoped(scope, move || {
+                            if let Err(e) = service.accept(stream) {
+                                crate::debug!("error: {e}");
                             }
                         })
                         .unwrap();