From 768a259f02f5743f7812904657a99f1f58b8e835 Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Sun, 12 Mar 2023 16:31:56 +0100
Subject: [PATCH] Add timeout while opening connection

This should hopefully fix an issue where instances would get stuck in a
"Connecting" state after suspending the host machine or maybe just
connecting it to a different wifi. I've not managed to reliably
reproduce the bug, so this may not fix it at all.
---
 CHANGELOG.md        |  3 +++
 src/bot/instance.rs |  8 ++++----
 src/conn.rs         | 10 ++++++++--
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e8118d9..6aee1ea 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,15 +18,18 @@ Procedure when bumping the version number:
 - `bot::command::parse_prefix_initiated`
 - `bot::commands::Commands::fallthrough`
 - `bot::commands::Commands::set_fallthrough`
+- `conn::Error::ConnectionTimedOut`
 
 ### Changed
 - `bot::command::ClapCommand::execute` now returns a `Result<bool, E>` instead of a `Result<(), E>`
 - `bot::command::Command::execute` now returns a `Result<bool, E>` instead of a `Result<(), E>`
 - `bot::commands::Commands::handle_packet` now returns a `Result<bool, E>` instead of a `Result<(), E>`
 - `bot::instance::Instance` now implements `Clone`
+- `conn::Conn::connect` now returns `conn::Result`
 
 ### Fixed
 - `phone` and `mobile` emoji
+- Instances getting stuck in "Connecting" state
 
 ## v0.3.1 - 2023-02-26
 
diff --git a/src/bot/instance.rs b/src/bot/instance.rs
index 83a1117..e6d2741 100644
--- a/src/bot/instance.rs
+++ b/src/bot/instance.rs
@@ -251,7 +251,7 @@ enum Request {
 enum RunError {
     StoppedManually,
     InstanceDropped,
-    CouldNotConnect(tungstenite::Error),
+    CouldNotConnect(conn::Error),
     Conn(conn::Error),
 }
 
@@ -394,9 +394,9 @@ impl Instance {
                     idebug!(config, "Instance dropped");
                     break;
                 }
-                Err(RunError::CouldNotConnect(tungstenite::Error::Http(response)))
-                    if response.status() == StatusCode::NOT_FOUND =>
-                {
+                Err(RunError::CouldNotConnect(conn::Error::Tungstenite(
+                    tungstenite::Error::Http(response),
+                ))) if response.status() == StatusCode::NOT_FOUND => {
                     iwarn!(config, "Failed to connect: room does not exist");
                     break;
                 }
diff --git a/src/conn.rs b/src/conn.rs
index dcfc948..2a5325b 100644
--- a/src/conn.rs
+++ b/src/conn.rs
@@ -30,6 +30,8 @@ pub type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
 pub enum Error {
     /// The connection is now closed.
     ConnectionClosed,
+    /// The connection was not opened in time.
+    ConnectionTimedOut,
     /// The server didn't reply to one of our commands in time.
     CommandTimedOut,
     /// The server did something that violated the api specification.
@@ -45,6 +47,7 @@ impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             Self::ConnectionClosed => write!(f, "connection closed"),
+            Self::ConnectionTimedOut => write!(f, "connection did not open in time"),
             Self::CommandTimedOut => write!(f, "server did not reply to command in time"),
             Self::ProtocolViolation(msg) => write!(f, "{msg}"),
             Self::Euph(msg) => write!(f, "{msg}"),
@@ -613,7 +616,7 @@ impl Conn {
         human: bool,
         cookies: Option<HeaderValue>,
         timeout: Duration,
-    ) -> tungstenite::Result<(Self, Vec<HeaderValue>)> {
+    ) -> Result<(Self, Vec<HeaderValue>)> {
         let human = if human { "?h=1" } else { "" };
         let uri = format!("wss://{domain}/room/{room}/ws{human}");
         debug!("Connecting to {uri} with cookies: {cookies:?}");
@@ -622,7 +625,10 @@ impl Conn {
             request.headers_mut().append(header::COOKIE, cookies);
         }
 
-        let (ws, response) = tokio_tungstenite::connect_async(request).await?;
+        let (ws, response) =
+            tokio::time::timeout(timeout, tokio_tungstenite::connect_async(request))
+                .await
+                .map_err(|_| Error::ConnectionTimedOut)??;
         let (mut parts, _) = response.into_parts();
         let cookies_set = match parts.headers.entry(header::SET_COOKIE) {
             header::Entry::Occupied(entry) => entry.remove_entry_mult().1.collect(),