FEAT: WebTransport datagrams

Author: BryanFRD

src/datas/docker.rs

@@ -0,0 +1,26 @@
+use bollard::secret::ContainerSummary;
+use serde_json::json;
+
+use super::EventDTO;
+
+impl EventDTO for ContainerSummary {
+  fn to_json(&self) -> String {
+    json!({
+      "id": self.id,
+      "names": self.names,
+      "image": self.image,
+      "command": self.command,
+      "created": self.created,
+      "state": self.state,
+      "status": self.status,
+      "ports": self.ports,
+      "labels": self.labels,
+    }).to_string()
+  }
+}
+
+impl EventDTO for Vec<ContainerSummary> {
+  fn to_json(&self) -> String {
+    json!(self.iter().map(|f| f.to_json()).collect::<Vec<String>>()).to_string()
+  }
+}

src/datas/mod.rs

@@ -0,0 +1,5 @@
+pub mod docker;
+
+pub trait EventDTO {
+  fn to_json(&self) -> String;
+}

src/events/docker.rs

@@ -0,0 +1,14 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+#[serde(tag = "type")]
+pub enum DockerEvent {
+  DockerStatus,
+  DockerContainersRestart { data: DockerContainersRestartData }
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct DockerContainersRestartData {
+  #[serde(rename = "containerId")]
+  pub container_id: String
+}

src/events/mod.rs

@@ -0,0 +1,13 @@
+use serde::{Deserialize, Serialize};
+use system::SystemEvent;
+use docker::DockerEvent;
+
+pub mod system;
+pub mod docker;
+
+#[derive(Serialize, Deserialize, Debug)]
+#[serde(untagged)]
+pub enum Event {
+  System(SystemEvent),
+  Docker(DockerEvent)
+}

src/events/system.rs

@@ -0,0 +1,7 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+#[serde(tag = "type")]
+pub enum SystemEvent {
+  SystemStatus
+}

src/main.rs

@@ -1,3 +1,5 @@
+mod events;
+mod datas;
 mod errors;
 mod webtransport;
 mod services;

src/webtransport/mod.rs

@@ -1,8 +1,12 @@
-use std::{error::Error, sync::Arc};
+use std::error::Error;
 use axum::Json;
-use tokio::sync::{broadcast, Mutex};
+use tokio::sync::broadcast;
 use wtransport::{endpoint::IncomingSession, Connection, Endpoint, Identity, ServerConfig};
+use crate::datas::EventDTO;
+use crate::events::docker::DockerEvent;
+use crate::events::system::SystemEvent;
 use crate::services::docker;
+use crate::events::Event;
 
 pub async fn start_webtransport() -> Result<(), Box<dyn Error + Send + Sync>> {
     let identity = match Identity::load_pemfiles("localhost.pem", "localhost-key.pem").await {
@@ -62,17 +66,18 @@ async fn handle_connection(incoming_session: IncomingSession, tx: broadcast::Sen
         }
     };
 
-    log::info!("Accepted connection from {:?}", connection.remote_address());
-    
-    let datagram_handle = tokio::spawn(handle_datagram(connection.clone()));
+    let datagram_handle = tokio::spawn(handle_datagram(connection.clone(), tx.clone()));
 
     let bidirectional_handle = tokio::spawn(handle_bidirectionnal(connection, tx));
 
     let _ = tokio::join!(datagram_handle, bidirectional_handle);
     Ok(())
 }
 
-async fn handle_datagram(connection: Connection) -> Result<(), Box<dyn Error + Send + Sync>> {
+async fn handle_datagram(connection: Connection, tx: broadcast::Sender<String>) -> Result<(), Box<dyn Error + Send + Sync>> {
+    log::info!("Accepted datagram connection from {:?}", connection.remote_address());
+    let mut rx = tx.subscribe();
+    
     loop {
         let datagram = match connection.receive_datagram().await {
             Ok(datagram) => datagram,
@@ -82,59 +87,63 @@ async fn handle_datagram(connection: Connection) -> Result<(), Box<dyn Error + S
             }
         };
         let received_message = String::from_utf8_lossy(&datagram);
-        log::info!("Received message: {:?}", received_message);
+        log::info!("Received datagram message: {:?}", received_message);
 
-        let response = b"Hello from server via datagram!";
-        match connection.send_datagram(response) {
-            Ok(_) => {},
+        let event: Event = match serde_json::from_str(&received_message) {
+            Ok(event) => event,
             Err(e) => {
-                log::error!("Failed to send datagram: {:?}", e);
-                return Err(Box::new(e));
+                log::error!("Failed to parse event: {:?}", e);
+                continue;
             }
         };
+        
+        match &event {
+            Event::Docker(docker_event) => {
+                match docker_event {
+                    DockerEvent::DockerStatus => {
+                        let event = "DockerStatus".to_string();
+                        connection.send_datagram(event.as_bytes())?;
+                    },
+                    DockerEvent::DockerContainersRestart { data } => {
+                        let containers = docker::get_containers().await?;
+                        log::info!("Restarting container: {:?}", containers.to_json().as_bytes().len());
+                        if let Err(error) = connection.send_datagram(containers.to_json()) {
+                            log::error!("Failed to send event: {:?}", error);
+                        }
+                    }
+                }
+            },
+            Event::System(system_event) => {
+                match system_event {
+                    SystemEvent::SystemStatus => {
+                        let event = "SystemStatus".to_string();
+                        connection.send_datagram(event.as_bytes())?;
+                    }
+                }
+            },
+        }
     }
 }
 
 async fn handle_bidirectionnal(connection: Connection, tx: broadcast::Sender<String>) -> Result<(), Box<dyn Error + Send + Sync>> {
-    log::info!("Accepted connection from {:?}", connection.remote_address());
+    log::info!("Accepted bidirectional connection from {:?}", connection.remote_address());
 
-    while let Ok((send_stream, mut recv_stream)) = connection.accept_bi().await {
+    while let Ok((mut send_stream, mut recv_stream)) = connection.accept_bi().await {
         log::trace!("Accepted bidirectional stream");
 
         let mut rx = tx.subscribe();
-        let send_stream = Arc::new(Mutex::new(send_stream));
 
-        {
-            let send_stream = Arc::clone(&send_stream);
-            tokio::spawn(async move {
-                let mut buffer = vec![0; 1024];
-                
-                while let Ok(Some(bytes_read)) = recv_stream.read(&mut buffer).await {
-                    let received_message = String::from_utf8_lossy(&buffer[..bytes_read]);
-                    log::info!("Received message: {:?}", received_message);
-                    
-                    let mut send_stream = send_stream.lock().await;
-                    
-                    let _ = send_stream.write_all(b"Hello from server!").await;
-                }
-            });
-        }
-        
-        {
-            let send_stream = Arc::clone(&send_stream);
-            tokio::spawn(async move {
-                while let Ok(event) = rx.recv().await {
-                    let mut send_stream = send_stream.lock().await;
-                    let _ = send_stream.write_all(event.as_bytes()).await;
-                }
-            });
-        }
+        tokio::spawn(async move {
+            while let Ok(event) = rx.recv().await {
+                let _ = send_stream.write_all(event.as_bytes()).await;
+            }
+        });
     }
 
     Ok(())
 }
 
-pub fn handle_event(event: String, data: Json<String>) {
+fn handle_event(event: String, data: Json<String>) {
     log::info!("Received event: {:?}", event);
     log::info!("Received data: {:?}", data);