MyFavMirrors/vaultwarden
1
0
Fork 0
mirror of https://github.com/dani-garcia/vaultwarden.git synced 2024-12-27 15:45:57 -05: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(); let users2 = users.clone();
tokio::spawn(async move { tokio::spawn(async move {
let addr = (CONFIG.websocket_address(), CONFIG.websocket_port()); 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 listener = TcpListener::bind(addr).await.expect("Can't listen on websocket port");
let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>(); let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>();
@ -317,9 +318,11 @@ pub fn start_notification_server() -> WebSocketUsers {
users 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; 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 // Accept connection, do initial handshake, validate auth token and get the user ID
use handshake::server::{Request, Response}; use handshake::server::{Request, Response};
let mut stream = accept_hdr_async(stream, |req: &Request, res: 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() => { res = stream.next() => {
match res { match res {
Some(Ok(message)) => { 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 // We should receive an initial message with the protocol and version, and we will reply to it
if let Message::Text(ref message) = message { 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) { if serde_json::from_str(msg).ok() == Some(INITIAL_MESSAGE) {
stream.send(Message::binary(INITIAL_RESPONSE)).await?; stream.send(Message::binary(INITIAL_RESPONSE)).await?;
@ -377,13 +393,15 @@ async fn handle_connection(stream: TcpStream, users: WebSocketUsers, _addr: Sock
} }
_= interval.tick() => { _= interval.tick() => {
if stream.send(Message::Binary(create_ping())).await.is_err() { if stream.send(Message::Ping(create_ping())).await.is_err() {
break; break;
} }
} }
} }
} }
info!("Closing WS connection from {addr}");
// Delete from map // Delete from map
users.map.entry(user_uuid).or_default().retain(|(uuid, _)| uuid != &entry_uuid); users.map.entry(user_uuid).or_default().retain(|(uuid, _)| uuid != &entry_uuid);
Ok(()) Ok(())