Handle Ping, Binary and Close

crates/cdk-axum/src/ws/mod.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use axum::extract::ws::{Message, WebSocket};
+use axum::extract::ws::{CloseFrame, Message, WebSocket};
 use cdk::mint::QuoteId;
 use cdk::nuts::nut17::NotificationPayload;
 use cdk::pub_sub::SubId;
@@ -80,14 +80,42 @@ pub async fn main_websocket(mut socket: WebSocket, state: MintState) {
                 }
             }
 
-            from_ws = socket.next() => {
-                let text = if let Some(Ok(Message::Text(message))) = from_ws {
-                    message
-                } else {
-                    // Client disconnected or sent something that is not a text
-                    return;
+            Some(from_ws) = socket.next() => {
+                let text = match from_ws {
+                    Ok(Message::Text(text)) => text.to_string(),
+                    Ok(Message::Binary(bin)) => String::from_utf8_lossy(&bin).to_string(),
+                    Ok(Message::Ping(payload)) => {
+                        // Reply with Pong with same payload
+                        if let Err(e) = socket.send(Message::Pong(payload)).await {
+                            tracing::error!("failed to send pong: {e}");
+                            break;
+                        }
+                        continue;
+                    },
+                    Ok(Message::Pong(_payload)) => {
+                        tracing::error!("Unexpected pong");
+                        continue;
+                    },
+                    Ok(Message::Close(frame)) => {
+                        if let Some(CloseFrame { code, reason }) = frame {
+                            tracing::info!("ws-close: code={code:?} reason='{reason}'");
+                        } else {
+                            tracing::info!("ws-close: no frame");
+                        }
+
+                        let _ = socket.send(Message::Close(Some(CloseFrame {
+                            code: axum::extract::ws::close_code::NORMAL,
+                            reason: "bye!".into(),
+                        }))).await;
+                        break;
+                    }
+                    Err(err) => {
+                        tracing::error!("ws-error: {err}");
+                        break;
+                    }
                 };
 
+
                 let request = match serde_json::from_str::<WsRequest>(&text) {
                     Ok(request) => request,
                     Err(err) => {
@@ -113,7 +141,9 @@ pub async fn main_websocket(mut socket: WebSocket, state: MintState) {
                 }
             }
             else =>  {
-
+                // Unexpected, we should exit the loop
+                tracing::warn!("Unexpected event, closing ws");
+                break;
             }
         }
     }