ping fix

src/main.rs

@@ -29,6 +29,7 @@ type Writer = SplitSink<Sock, Message>;
 
 enum OutgoingMessage {
     Application(ClientMsg),
+    Ping,
     Pong(Vec<u8>),
 }
 
@@ -121,11 +122,38 @@ async fn process_messages(
     mut writer_closed: oneshot::Receiver<()>,
     config: &Arc<Config>,
 ) -> Result<(), String> {
+    let mut ping_interval = tokio::time::interval_at(
+        tokio::time::Instant::now() + config.ping_interval,
+        config.ping_interval,
+    );
+
+    let pong_timer = tokio::time::sleep(Duration::from_secs(0));
+    tokio::pin!(pong_timer);
+    let mut awaiting_pong = false;
+
     loop {
         tokio::select! {
             _ = &mut writer_closed => return Err("Write half closed".into()),
+            _ = ping_interval.tick() => {
+                if awaiting_pong {
+                    return Err("Server did not respond to ping".into());
+                }
+
+                send_outgoing(&outgoing_tx, OutgoingMessage::Ping).await?;
+                awaiting_pong = true;
+                pong_timer.as_mut().reset(
+                    tokio::time::Instant::now() + config.ping_timeout,
+                );
+            }
+            _ = &mut pong_timer, if awaiting_pong => {
+                return Err("Server did not respond to ping".into());
+            }
             next_message = stream.next() => {
-                handle_stream_message(next_message, &outgoing_tx, config.max_output_bytes).await?;
+                let action = handle_stream_message(next_message, &outgoing_tx, config.max_output_bytes).await?;
+
+                if action.clear_pong_deadline {
+                    awaiting_pong = false;
+                }
             }
         }
     }
@@ -160,6 +188,10 @@ async fn write_message(sink: &mut Writer, message: OutgoingMessage) -> Result<()
                 .await
                 .map_err(|e| e.to_string())
         }
+        OutgoingMessage::Ping => sink
+            .send(Message::Ping(Vec::new().into()))
+            .await
+            .map_err(|e| e.to_string()),
         OutgoingMessage::Pong(payload) => sink
             .send(Message::Pong(payload.into()))
             .await
@@ -167,6 +199,10 @@ async fn write_message(sink: &mut Writer, message: OutgoingMessage) -> Result<()
     }
 }
 
+struct StreamAction {
+    clear_pong_deadline: bool,
+}
+
 async fn send_outgoing(
     outgoing_tx: &mpsc::Sender<OutgoingMessage>,
     message: OutgoingMessage,
@@ -178,17 +214,29 @@ async fn handle_stream_message(
     message: Option<Result<Message, tokio_tungstenite::tungstenite::Error>>,
     outgoing_tx: &mpsc::Sender<OutgoingMessage>,
     max_output_bytes: u64,
-) -> Result<(), String> {
+) -> Result<StreamAction, String> {
     match message {
         Some(Ok(Message::Text(text))) => {
-            handle_server_message(parse_server_message(&text)?, outgoing_tx, max_output_bytes).await
+            handle_server_message(parse_server_message(&text)?, outgoing_tx, max_output_bytes)
+                .await?;
+            Ok(StreamAction {
+                clear_pong_deadline: false,
+            })
         }
         Some(Ok(Message::Ping(payload))) => {
-            send_outgoing(outgoing_tx, OutgoingMessage::Pong(payload.to_vec())).await
+            send_outgoing(outgoing_tx, OutgoingMessage::Pong(payload.to_vec())).await?;
+            Ok(StreamAction {
+                clear_pong_deadline: false,
+            })
         }
         Some(Ok(Message::Close(_))) | None => Err("Connection closed".into()),
         Some(Err(error)) => Err(error.to_string()),
-        Some(Ok(_)) => Ok(()),
+        Some(Ok(Message::Pong(_))) => Ok(StreamAction {
+            clear_pong_deadline: true,
+        }),
+        Some(Ok(_)) => Ok(StreamAction {
+            clear_pong_deadline: false,
+        }),
     }
 }