Locutus/locutus-server
2
1
Fork 0
Code Issues Pull Requests Projects Releases Activity

Fix WebSocket connection closes

Browse Source
This commit is contained in:
Lorenz Stechauner
2022-08-27 17:39:19 +02:00
parent f3c940a96c
commit d94daa2c93
3 changed files with 107 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/usimp/handler/mod.rs
View File

@ -5,10 +5,11 @@ mod subscribe;
use crate::usimp::*;
use tokio::sync::mpsc;
use crate::websocket::WebSocketEnvelope;
pub async fn endpoint(
input: &InputEnvelope,
tx: Option<mpsc::Sender<OutputEnvelope>>,
tx: Option<mpsc::Sender<WebSocketEnvelope>>,
) -> Result<OutputEnvelope, Error> {
if input.from_domain != None {
// TODO

7
src/usimp/handler/subscribe.rs
View File

@ -4,6 +4,7 @@ use crate::usimp::*;
use serde::{Deserialize, Serialize};
use serde_json::{from_value, to_value, Value};
use tokio::sync::mpsc;
use crate::websocket::WebSocketEnvelope;
#[derive(Serialize, Deserialize, Clone)]
struct Input {}
@ -16,7 +17,7 @@ struct Output {
pub async fn handle(
input: &InputEnvelope,
session: Option<Session>,
tx: Option<mpsc::Sender<OutputEnvelope>>,
tx: Option<mpsc::Sender<WebSocketEnvelope>>,
) -> Result<Value, Error> {
Ok(to_value(
subscribe(
@ -33,7 +34,7 @@ async fn subscribe(
_input: Input,
session: Option<Session>,
req_nr: Option<u64>,
tx: Option<mpsc::Sender<OutputEnvelope>>,
tx: Option<mpsc::Sender<WebSocketEnvelope>>,
) -> Result<Output, Error> {
let account = get_account(&session)?;
let mut rx = subscription::subscribe_account(account).await;
@ -46,7 +47,7 @@ async fn subscribe(
error: None,
request_nr: req_nr,
data: serde_json::json![{"events": [event]}],
})
}.into())
.await;
}
});

93
src/websocket.rs
View File

@ -11,11 +11,45 @@ use hyper_tungstenite::{tungstenite::protocol::Role, WebSocketStream};
use serde_json::{Map, Value};
use tokio::sync::mpsc;
pub enum WebSocketEnvelope {
Close,
Ping(Vec<u8>),
Pong(Vec<u8>),
Text(OutputEnvelope),
}
impl From<OutputEnvelope> for WebSocketEnvelope {
fn from(envelope: OutputEnvelope) -> Self {
WebSocketEnvelope::Text(envelope)
}
}
async fn sender(
mut sink: SplitSink<WebSocketStream<Upgraded>, Message>,
mut rx: mpsc::Receiver<OutputEnvelope>,
mut rx: mpsc::Receiver<WebSocketEnvelope>,
) {
while let Some(msg) = rx.recv().await {
match msg {
WebSocketEnvelope::Ping(data) => {
if let Err(error) = sink.send(Message::Ping(data)).await {
eprintln!("{:?}", error);
return;
}
}
WebSocketEnvelope::Pong(data) => {
if let Err(error) = sink.send(Message::Pong(data)).await {
eprintln!("{:?}", error);
return;
}
}
WebSocketEnvelope::Close => {
if let Err(error) = sink.send(Message::Close(None)).await {
eprintln!("{:?}", error);
return;
}
break;
}
WebSocketEnvelope::Text(msg) => {
let mut envelope = Value::Object(Map::new());
envelope["data"] = msg.data;
envelope["request_nr"] = match msg.request_nr {
@ -33,24 +67,58 @@ async fn sender(
}
if let Err(error) = sink.send(Message::Text(envelope.to_string())).await {
eprintln!("{:?}", error);
break;
return;
}
}
}
}
}
async fn receiver(
mut stream: SplitStream<WebSocketStream<Upgraded>>,
tx: mpsc::Sender<OutputEnvelope>,
tx: mpsc::Sender<WebSocketEnvelope>,
) {
while let Some(res) = stream.next().await {
match res {
Ok(msg) => {
let _res;
if msg.is_ping() {
_res = tx.send(WebSocketEnvelope::Pong(msg.into_data())).await;
} else if msg.is_pong() {
// Ignore
} else if msg.is_close() {
_res = tx.send(WebSocketEnvelope::Close).await;
break;
} else if msg.is_binary() {
_res = tx.send(WebSocketEnvelope::Text(OutputEnvelope {
error: Some(Error {
kind: ErrorKind::WebSocketError,
class: ErrorClass::ClientProtocolError,
msg: Some("Binary frames are not allowed".to_string()),
desc: None,
}),
request_nr: None,
data: Value::Null,
})).await;
} else if msg.is_text() {
let input: InputEnvelope = serde_json::from_slice(&msg.into_data()[..]).unwrap();
let output = match usimp::endpoint(&input, Some(tx.clone())).await {
Ok(output) => output,
Err(error) => input.error(error),
};
let _res = tx.send(output).await;
_res = tx.send(WebSocketEnvelope::Text(output)).await;
} else {
_res = tx.send(WebSocketEnvelope::Text(OutputEnvelope {
error: Some(Error {
kind: ErrorKind::WebSocketError,
class: ErrorClass::ClientProtocolError,
msg: Some("Unknown frame opcode".to_string()),
desc: None,
}),
request_nr: None,
data: Value::Null,
})).await;
}
}
Err(error) => println!("{:?}", error),
}
@ -71,7 +139,7 @@ pub async fn handler(
ErrorClass::ClientProtocolError,
None,
))),
)
);
}
}
@ -85,7 +153,7 @@ pub async fn handler(
ErrorClass::ClientProtocolError,
None,
))),
)
);
}
};
let key = handshake::derive_accept_key(key.as_bytes());
@ -100,19 +168,20 @@ pub async fn handler(
ErrorClass::ClientProtocolError,
None,
))),
)
);
}
}
tokio::spawn(async move {
match hyper::upgrade::on(req).await {
Ok(upgraded) => {
let ws_stream =
WebSocketStream::from_raw_socket(upgraded, Role::Server, None).await;
let (tx, rx) = mpsc::channel::<OutputEnvelope>(64);
let ws_stream = WebSocketStream::from_raw_socket(upgraded, Role::Server, None).await;
let (tx, rx) = mpsc::channel::<WebSocketEnvelope>(64);
let (sink, stream) = ws_stream.split();
tokio::spawn(async move { sender(sink, rx).await });
receiver(stream, tx).await
let other = tokio::spawn(async move { sender(sink, rx).await });
receiver(stream, tx).await;
other.await.unwrap();
}
Err(error) => eprintln!("Unable to upgrade: {}", error),
}