octoleo/vaultwarden
1
0
Fork 0
mirror of https://github.com/dani-garcia/vaultwarden.git synced 2024-06-08 19:32:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix incorrect pings sent, and respond to pings from the client

Browse Source
This commit is contained in:
Daniel García 2022-05-21 19:12:38 +02:00
parent 699777be9e
commit 0de5919a16
No known key found for this signature in database
GPG Key ID: FC8A7D14C3CD543A
1 changed files with 21 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
src/api/notifications.rs
View File

@ -293,6 +293,7 @@ pub fn start_notification_server() -> WebSocketUsers {
let users2 = users.clone();
tokio::spawn(async move {
let addr = (CONFIG.websocket_address(), CONFIG.websocket_port());
info!("Starting WebSockets server on {}:{}", addr.0, addr.1);
let listener = TcpListener::bind(addr).await.expect("Can't listen on websocket port");
let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>();
@ -317,9 +318,11 @@ pub fn start_notification_server() -> WebSocketUsers {
users
}
async fn handle_connection(stream: TcpStream, users: WebSocketUsers, _addr: SocketAddr) -> Result<(), Error> {
async fn handle_connection(stream: TcpStream, users: WebSocketUsers, addr: SocketAddr) -> Result<(), Error> {
let mut user_uuid: Option<String> = None;
info!("Accepting WS connection from {addr}");
// Accept connection, do initial handshake, validate auth token and get the user ID
use handshake::server::{Request, Response};
let mut stream = accept_hdr_async(stream, |req: &Request, res: Response| {
@ -346,9 +349,22 @@ async fn handle_connection(stream: TcpStream, users: WebSocketUsers, _addr: Sock
res = stream.next() => {
match res {
Some(Ok(message)) => {
//info!("RECEIVED {message:?}");
// Respond to any pings
if let Message::Ping(ping) = message {
if stream.send(Message::Pong(ping)).await.is_err() {
break;
}
continue;
} else if let Message::Pong(_) = message {
/* Ignored */
continue;
}
// We should receive an initial message with the protocol and version, and we will reply to it
if let Message::Text(ref message) = message {
let msg = message.strip_suffix('\u{1e}').unwrap_or(message);
let msg = message.strip_suffix(RECORD_SEPARATOR as char).unwrap_or(message);
if serde_json::from_str(msg).ok() == Some(INITIAL_MESSAGE) {
stream.send(Message::binary(INITIAL_RESPONSE)).await?;
@ -377,13 +393,15 @@ async fn handle_connection(stream: TcpStream, users: WebSocketUsers, _addr: Sock
}
_= interval.tick() => {
if stream.send(Message::Binary(create_ping())).await.is_err() {
if stream.send(Message::Ping(create_ping())).await.is_err() {
break;
}
}
}
}
info!("Closing WS connection from {addr}");
// Delete from map
users.map.entry(user_uuid).or_default().retain(|(uuid, _)| uuid != &entry_uuid);
Ok(())