From f42b201791ef38ff796d1d05b02100edd20d1203 Mon Sep 17 00:00:00 2001
From: Erwin Boskma <erwin@datarift.nl>
Date: Thu, 28 Sep 2023 20:37:45 +0200
Subject: [PATCH] This should reconnect when the connection is lost

---
 flake.nix            |  2 ++
 src/homeassistant.rs | 12 ++++++++----
 src/main.rs          | 33 ++++++++++++++++++++++++---------
 3 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/flake.nix b/flake.nix
index b3e0bb4..e7ae030 100644
--- a/flake.nix
+++ b/flake.nix
@@ -113,6 +113,8 @@
             rust-analyzer
             taplo
             gitflow
+
+            dsniff # has the tcpkill command to test dropped connections
           ];
         };
 
diff --git a/src/homeassistant.rs b/src/homeassistant.rs
index be079ec..aa178c9 100644
--- a/src/homeassistant.rs
+++ b/src/homeassistant.rs
@@ -284,7 +284,7 @@ impl HomeAssistant {
                                                     "Received '{message_type}', we're currently not handling that."
                                                 )
                                             }
-                                            HomeAssistantError::JsonError(e) => error!("{e}"),
+                                            HomeAssistantError::JsonError(e) => error!("Error deserializing JSON: {e}"),
                                         }
                                     } else {
                                         error!("{e}");
@@ -294,10 +294,14 @@ impl HomeAssistant {
                         },
                         Err(e) => {
                             error!("Error in receiving message: {e}");
-                            break;
+                            // break;
                         }
                     }
-                    None => break,
+                    None => {
+                        // I recon we should reconnect here. Probably.
+                        error!("Received 'None' from handle_message");
+                        break;
+                    },
                 },
                 _ = ping_interval.tick() => {
                     if self.auth_complete {
@@ -307,7 +311,7 @@ impl HomeAssistant {
 
                         ws_stream.send(ping_message.to_message()).await?;
                     }
-                }
+                },
             }
         }
 
diff --git a/src/main.rs b/src/main.rs
index 330c3cb..3d8ae3d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,7 +5,8 @@ use clap::Parser;
 use color_eyre::{eyre::WrapErr, Result};
 use serde_json::json;
 use std::fmt::Display;
-use tracing::debug;
+use tokio::signal::unix::SignalKind;
+use tracing::{debug, error};
 use tracing_error::ErrorLayer;
 use tracing_subscriber::{prelude::*, EnvFilter, Registry};
 use tracing_tree::HierarchicalLayer;
@@ -127,19 +128,33 @@ async fn main() -> Result<()> {
                 .wrap_err_with(|| format!("Unable to execute command {cmd}"))?;
         } else {
             debug!("Retrieving state of '{entity}' on '{host}'");
-            let ha = homeassistant::HomeAssistant::builder()
-                .host(host)
-                .entity(entity)
-                .token(token)
-                .insecure(insecure)
-                .format(format)
-                .build();
-            ha.open_connection().await.unwrap_or(());
+            loop {
+                tokio::select! {
+                    _ = homeassistant::HomeAssistant::builder()
+                        .host(host.clone())
+                        .entity(entity.clone())
+                        .token(token.clone())
+                        .insecure(insecure)
+                        .format(format)
+                        .build().open_connection() => {
+                            error!("Connection closed? Reconnecting...");
+                        },
+                    _ = tokio::signal::ctrl_c() => break,
+                    _ = sig_term() => break,
+                };
+            }
         }
     }
     Ok(())
 }
 
+async fn sig_term() -> Result<()> {
+    tokio::signal::unix::signal(SignalKind::terminate())?
+        .recv()
+        .await;
+    Ok(())
+}
+
 fn setup(debug: bool) -> Result<()> {
     if debug {
         std::env::set_var("RUST_LOG", "debug");