103 lines
3.5 KiB
Rust
103 lines
3.5 KiB
Rust
use std::sync::Arc;
|
|
|
|
use futures::StreamExt;
|
|
use hyper::upgrade::OnUpgrade;
|
|
use hyper_util::rt::TokioIo;
|
|
use tokio_tungstenite::WebSocketStream;
|
|
use tungstenite::{Message, Utf8Bytes};
|
|
|
|
use crate::data::communication::{CommunicationType, CommunicationValue};
|
|
use crate::log;
|
|
use crate::server::omikron_connection::OmikronConnection;
|
|
|
|
pub fn handle(path: String, upgrades: OnUpgrade) {
|
|
tokio::spawn(async move {
|
|
match upgrades.await {
|
|
Ok(upgraded_stream) => {
|
|
log!("Valid WebSocket upgrade");
|
|
let raw_stream = TokioIo::new(upgraded_stream);
|
|
|
|
let ws_stream = WebSocketStream::from_raw_socket(
|
|
raw_stream,
|
|
tungstenite::protocol::Role::Server,
|
|
None,
|
|
)
|
|
.await;
|
|
log!(
|
|
"WebSocket handshake successful, handling connection for {}",
|
|
path
|
|
);
|
|
|
|
let (writer, reader) = ws_stream.split();
|
|
if path == "/ws/omikron" {
|
|
let connection = OmikronConnection::new(writer, reader);
|
|
start_connecteable_handler(connection).await;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
log!("WebSocket upgrade failed after response: {:?}", e);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
|
|
use futures::SinkExt;
|
|
use tokio::time::Duration;
|
|
|
|
const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
|
|
|
|
loop {
|
|
let mut receiver = connection.receiver.write().await;
|
|
|
|
match tokio::time::timeout(IDLE_TIMEOUT, receiver.next()).await {
|
|
Ok(Some(Ok(msg))) => {
|
|
// Drop the lock so other tasks can use the receiver if needed,
|
|
// and so we can handle the message without holding the lock.
|
|
drop(receiver);
|
|
|
|
match msg {
|
|
Message::Text(text) => {
|
|
let conn_clone = connection.clone();
|
|
tokio::spawn(async move {
|
|
conn_clone.handle_message(text.to_string()).await;
|
|
});
|
|
}
|
|
Message::Close(_) => {
|
|
break;
|
|
}
|
|
Message::Pong(_) => {
|
|
// Received a pong, connection is alive.
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
Ok(Some(Err(e))) => {
|
|
log!("WS Error: {}", e);
|
|
break;
|
|
}
|
|
Ok(None) => {
|
|
// Stream is closed
|
|
break;
|
|
}
|
|
Err(_) => {
|
|
// Timeout, we need to send a ping.
|
|
// Drop receiver lock before acquiring sender lock to avoid deadlock.
|
|
drop(receiver);
|
|
log!("WebSocket connection is idle. Sending a ping.");
|
|
let mut sender = connection.sender.write().await;
|
|
if let Err(e) = sender
|
|
.send(Message::Text(Utf8Bytes::from(
|
|
CommunicationValue::new(CommunicationType::ping)
|
|
.to_json()
|
|
.to_string(),
|
|
)))
|
|
.await
|
|
{
|
|
log!("Failed to send ping: {}. Closing connection.", e);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
connection.handle_close().await;
|
|
}
|