Add loop server/proxy

Signed-off-by: Fabrice Desclaux <[email protected]>
cea-sec · Jan 19, 2023 · 3b85311 · 3b85311
1 parent 5f65881
commit 3b85311
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 2 deletions.
diff --git a/sanzu/src/bin/sanzu_proxy.rs b/sanzu/src/bin/sanzu_proxy.rs
@@ -1,7 +1,7 @@
 #[macro_use]
 extern crate log;
 
-use clap::{Arg, Command};
+use clap::{Arg, ArgAction, Command};
 use sanzu::{config::read_server_config, proxy, utils::ArgumentsProxy};
 use sanzu_common::proto::VERSION;
 use std::net::IpAddr;
@@ -106,6 +106,13 @@ Protocol version: {:?}
                 .long("shm_is_xwd")
                 .num_args(0),
         )
+        .arg(
+            Arg::new("loop")
+                .help("Endless server clients")
+                .short('d')
+                .long("loop")
+                .action(ArgAction::SetTrue),
+        )
         .get_matches();
 
     let server_addr = matches
@@ -129,6 +136,7 @@ Protocol version: {:?}
     let vsock = matches.get_flag("vsock");
     let import_video_shm = matches.get_one::<String>("import_video_shm").cloned();
     let shm_is_xwd = matches.get_flag("shm_is_xwd");
+    let endless_loop = matches.get_flag("loop");
 
     let arguments = ArgumentsProxy {
         vsock,
@@ -141,6 +149,7 @@ Protocol version: {:?}
         audio,
         video_shared_mem: import_video_shm,
         shm_is_xwd,
+        endless_loop,
     };
 
     if let Err(err) = proxy::run(&conf, &arguments) {

diff --git a/sanzu/src/bin/sanzu_server.rs b/sanzu/src/bin/sanzu_server.rs
@@ -162,6 +162,13 @@ Protocol version: {:?}
                 .long("rdonly")
                 .action(ArgAction::SetTrue),
         )
+        .arg(
+            Arg::new("loop")
+                .help("Endless server clients")
+                .short('d')
+                .long("loop")
+                .action(ArgAction::SetTrue),
+        )
         .get_matches();
 
     let address = matches.get_one::<String>("listen").unwrap();
@@ -188,6 +195,7 @@ Protocol version: {:?}
     let extern_img_source = matches.get_one::<String>("use_extern_img_source").cloned();
     let avoid_img_extraction = matches.get_flag("avoid_img_extraction");
     let rdonly = matches.get_flag("rdonly");
+    let endless_loop = matches.get_flag("loop");
 
     let arguments = ArgumentsSrv {
         vsock,
@@ -206,6 +214,7 @@ Protocol version: {:?}
         extern_img_source,
         avoid_img_extraction,
         rdonly,
+        endless_loop,
     };
 
     if let Err(err) = server::run(&conf, &arguments) {

diff --git a/sanzu/src/proxy.rs b/sanzu/src/proxy.rs
@@ -136,7 +136,22 @@ macro_rules! send_srv_msg_type {
     }};
 }
 
+/// Exec main loop
+///
 pub fn run(config: &ConfigServer, arguments: &ArgumentsProxy) -> Result<()> {
+    if arguments.endless_loop {
+        loop {
+            if let Err(err) = run_server(config, arguments) {
+                error!("Server error");
+                err.chain().for_each(|cause| error!(" - due to {}", cause));
+            }
+        }
+    } else {
+        run_server(config, arguments)
+    }
+}
+
+pub fn run_server(config: &ConfigServer, arguments: &ArgumentsProxy) -> Result<()> {
     let codec_name = get_encoder_category(&arguments.encoder_name)?;
 
     let mut sound_encoder = opus::Encoder::new(

diff --git a/sanzu/src/server.rs b/sanzu/src/server.rs
@@ -97,6 +97,21 @@ fn auth_client(
     Ok((tls_conn, username))
 }
 
+/// Exec main loop
+///
+pub fn run(config: &ConfigServer, arguments: &ArgumentsSrv) -> Result<()> {
+    if arguments.endless_loop {
+        loop {
+            if let Err(err) = run_server(config, arguments) {
+                error!("Server error");
+                err.chain().for_each(|cause| error!(" - due to {}", cause));
+            }
+        }
+    } else {
+        run_server(config, arguments)
+    }
+}
+
 /// Server main loop
 ///
 /// The loop is composed of the following actions:
@@ -106,7 +121,7 @@ fn auth_client(
 /// - encode image
 /// - serialize / send events to client
 /// - receive / handle client events
-pub fn run(config: &ConfigServer, arguments: &ArgumentsSrv) -> Result<()> {
+pub fn run_server(config: &ConfigServer, arguments: &ArgumentsSrv) -> Result<()> {
     info!("Start server");
 
     let mut sock: Box<dyn ReadWrite> = match (arguments.vsock, arguments.stdio, arguments.unixsock)

diff --git a/sanzu/src/utils.rs b/sanzu/src/utils.rs
@@ -36,6 +36,7 @@ pub struct ArgumentsSrv<'a> {
     pub extern_img_source: Option<String>,
     pub avoid_img_extraction: bool,
     pub rdonly: bool,
+    pub endless_loop: bool,
 }
 
 pub struct ArgumentsClient<'a> {
@@ -71,6 +72,7 @@ pub struct ArgumentsProxy<'a> {
     pub audio: bool,
     pub video_shared_mem: Option<String>,
     pub shm_is_xwd: bool,
+    pub endless_loop: bool,
 }
 
 const MAX_HEADER_SIZE: u32 = 0x100;