Skip to content

Commit

Permalink
fix(backend): quit socket_handler if there's a websocket error (#312)
Browse files Browse the repository at this point in the history 
* fix(backend): quit if there's a websocket error

* fix(backend): fix 'outer loop label
  • Loading branch information
@ravenclaw900
ravenclaw900 authored Aug 2, 2022
1 parent f956963 commit 5bb4142
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 36 deletions.
49 changes: 27 additions & 22 deletions src/page_handlers.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,12 @@ fn main_handler_getter(
})
}

// Return true if error was related to websocket, false otherwise
#[instrument(skip_all)]
pub async fn main_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) {
pub async fn main_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) -> bool {
let mut cpu_collector = handle_error!(
psutil::cpu::CpuPercentCollector::new().context("Couldn't init cpu collector"),
return
return false
);

let mut net_collector = psutil::network::NetIoCountersCollector::default();
Expand All @@ -53,14 +54,14 @@ pub async fn main_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChan
loop {
tokio::select! {
biased;
data = data_recv.recv() => if let Some(Some(_)) = data {} else { break },
data = data_recv.recv() => if let Some(Some(_)) = data {} else { return false },
res = socket_send
.send(json_msg!(&handle_error!(main_handler_getter(&mut cpu_collector, &mut net_collector, &mut prev_data), shared::SysData::default()), continue))
=> {
sleep(Duration::from_secs(1)).await;
if res.is_err() {
tracing::debug!("Socket send failed, returning");
break
return true;
}
},
}
Expand Down Expand Up @@ -92,14 +93,14 @@ fn process_handler_helper(data: &shared::Request) -> anyhow::Result<()> {
}

#[instrument(skip_all)]
pub async fn process_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) {
pub async fn process_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) -> bool {
loop {
tokio::select! {
biased;
data = data_recv.recv() => if let Some(Some(data)) = data {
handle_error!(process_handler_helper(&data));
} else {
break;
return false;
},
res = socket_send
.send(json_msg!(
Expand All @@ -110,7 +111,7 @@ pub async fn process_handler(socket_send: &mut SocketSend, data_recv: &mut RecvC
sleep(Duration::from_secs(1)).await;
if res.is_err() {
tracing::debug!("Socket send failed, returning");
break;
return true;
}
},
}
Expand Down Expand Up @@ -150,7 +151,7 @@ pub async fn software_handler_helper(
}

#[instrument(skip_all)]
pub async fn software_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) {
pub async fn software_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) -> bool {
let software = handle_error!(systemdata::dpsoftware().await, (Vec::new(), Vec::new()));
if socket_send
.send(json_msg!(
Expand All @@ -159,13 +160,13 @@ pub async fn software_handler(socket_send: &mut SocketSend, data_recv: &mut Recv
installed: software.1,
response: String::new(),
},
return
return false
))
.await
.is_err()
{
tracing::debug!("Socket send failed, returning");
return;
return true;
}
while let Some(Some(data)) = data_recv.recv().await {
let out = handle_error!(
Expand All @@ -174,23 +175,24 @@ pub async fn software_handler(socket_send: &mut SocketSend, data_recv: &mut Recv
);
if socket_send.send(json_msg!(&out, continue)).await.is_err() {
tracing::debug!("Socket send failed, returning");
break;
return true;
}
}
false
}

#[instrument(skip_all)]
pub async fn management_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) {
pub async fn management_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) -> bool {
if socket_send
.send(json_msg!(
&handle_error!(systemdata::host().await, shared::HostData::default()),
return
return false
))
.await
.is_err()
{
tracing::debug!("Socket send failed, returning");
return;
return true;
}
while let Some(Some(data)) = data_recv.recv().await {
tracing::info!("Running command {}", &data.cmd);
Expand All @@ -200,22 +202,23 @@ pub async fn management_handler(socket_send: &mut SocketSend, data_recv: &mut Re
.map(|_| ())
.with_context(|| format!("Couldn't spawn command {}", &data.cmd)));
}
false
}

#[instrument(skip_all)]
pub async fn service_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) {
pub async fn service_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) -> bool {
if socket_send
.send(json_msg!(
&shared::ServiceList {
services: handle_error!(systemdata::services().await, Vec::new()),
},
return
return false
))
.await
.is_err()
{
tracing::debug!("Socket send failed, returning");
return;
return true;
}
while let Some(Some(data)) = data_recv.recv().await {
handle_error!(Command::new("systemctl")
Expand All @@ -234,9 +237,10 @@ pub async fn service_handler(socket_send: &mut SocketSend, data_recv: &mut RecvC
.is_err()
{
tracing::debug!("Socket send failed, returning");
break;
return true;
}
}
false
}

async fn browser_refresh(path: &std::path::Path) -> anyhow::Result<shared::BrowserList> {
Expand Down Expand Up @@ -309,7 +313,7 @@ async fn browser_handler_helper(data: &shared::Request) -> anyhow::Result<shared
}

#[instrument(skip_all)]
pub async fn browser_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) {
pub async fn browser_handler(socket_send: &mut SocketSend, data_recv: &mut RecvChannel) -> bool {
// Get initial listing of $HOME
if socket_send
.send(json_msg!(
Expand All @@ -322,13 +326,13 @@ pub async fn browser_handler(socket_send: &mut SocketSend, data_recv: &mut RecvC
Vec::new()
),
},
return
return false
))
.await
.is_err()
{
tracing::debug!("Socket send failed, returning");
return;
return true;
}

'outer: while let Some(Some(mut data)) = data_recv.recv().await {
Expand All @@ -338,7 +342,7 @@ pub async fn browser_handler(socket_send: &mut SocketSend, data_recv: &mut RecvC
let list = handle_error!(res, shared::BrowserList::default());
if socket_send.send(json_msg!(&list, continue)).await.is_err() {
tracing::debug!("Socket send failed, returning");
break 'outer;
return true;
}
break;
},
Expand All @@ -349,4 +353,5 @@ pub async fn browser_handler(socket_send: &mut SocketSend, data_recv: &mut RecvC
}
}
}
false
}
24 changes: 10 additions & 14 deletions src/socket_handlers.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,23 +98,15 @@ pub async fn socket_handler(socket: warp::ws::WebSocket) {
return;
}
while let Some(Some(message)) = data_recv.recv().await {
match message.page.as_str() {
if match message.page.as_str() {
"/" => page_handlers::main_handler(&mut socket_send, &mut data_recv).await,
"/process" => {
page_handlers::process_handler(&mut socket_send, &mut data_recv).await;
}
"/software" => {
page_handlers::software_handler(&mut socket_send, &mut data_recv).await;
}
"/process" => page_handlers::process_handler(&mut socket_send, &mut data_recv).await,
"/software" => page_handlers::software_handler(&mut socket_send, &mut data_recv).await,
"/management" => {
page_handlers::management_handler(&mut socket_send, &mut data_recv).await;
}
"/service" => {
page_handlers::service_handler(&mut socket_send, &mut data_recv).await;
}
"/browser" => {
page_handlers::browser_handler(&mut socket_send, &mut data_recv).await;
page_handlers::management_handler(&mut socket_send, &mut data_recv).await
}
"/service" => page_handlers::service_handler(&mut socket_send, &mut data_recv).await,
"/browser" => page_handlers::browser_handler(&mut socket_send, &mut data_recv).await,
"/login" => {
tracing::debug!("Sending login message");
// Internal poll, see other thread
Expand All @@ -128,10 +120,14 @@ pub async fn socket_handler(socket: warp::ws::WebSocket) {
{
break;
}
false
}
_ => {
tracing::debug!("Got page {}, not handling", message.page);
false
}
} {
break;
}
}
}
Expand Down

0 comments on commit 5bb4142

Please sign in to comment.