From 3cd3c2b23f5946775810465be45ce868304ae378 Mon Sep 17 00:00:00 2001
From: fwcd <fwcd@fwcd.dev>
Date: Thu, 18 Jul 2024 01:04:05 +0200
Subject: [PATCH] Handle case where no new messages are available

---
 lighthouse-client/src/error.rs      | 2 ++
 lighthouse-client/src/lighthouse.rs | 6 +++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/lighthouse-client/src/error.rs b/lighthouse-client/src/error.rs
index 93c27ec..03d6fa9 100644
--- a/lighthouse-client/src/error.rs
+++ b/lighthouse-client/src/error.rs
@@ -17,6 +17,8 @@ pub enum Error {
     Value(#[from] ValueError),
     #[error("Server error: {} {} (warnings: {:?})", code, message.clone().unwrap_or_else(|| "(no message)".to_string()), warnings)]
     Server { code: i32, message: Option<String>, warnings: Vec<String> },
+    #[error("No next message available")]
+    NoNextMessage,
     #[error("Custom error")]
     Custom(String),
 }
diff --git a/lighthouse-client/src/lighthouse.rs b/lighthouse-client/src/lighthouse.rs
index 0ef6f7e..117e775 100644
--- a/lighthouse-client/src/lighthouse.rs
+++ b/lighthouse-client/src/lighthouse.rs
@@ -85,6 +85,10 @@ impl<S> Lighthouse<S>
                         warn!("Got message without request id from server: {:?}", msg);
                     }
                 },
+                Err(Error::NoNextMessage) => {
+                    info!("No next message available, closing receive loop");
+                    break
+                },
                 Err(e) => error!("Bad message: {:?}", e),
             }
         }
@@ -104,7 +108,7 @@ impl<S> Lighthouse<S>
     #[tracing::instrument(skip(ws_stream))]
     async fn receive_raw_from(ws_stream: &mut SplitStream<S>) -> Result<Vec<u8>> {
         loop {
-            let message = ws_stream.next().await.ok_or_else(|| Error::custom("Got no message"))??;
+            let message = ws_stream.next().await.ok_or_else(|| Error::NoNextMessage)??;
             match message {
                 Message::Binary(bytes) => break Ok(bytes),
                 // We ignore pings for now