From 9c21ddd35b5000c9fcbd10a196cbe8691b8550a3 Mon Sep 17 00:00:00 2001
From: Erwin Boskma <erwin@datarift.nl>
Date: Tue, 19 Apr 2022 20:28:04 +0200
Subject: [PATCH] Ping and Pong

Send a ping message every 5 seconds, and handle incoming Ping and Pong messages.
---
 src/homeassistant.rs | 136 +++++++++++++++++++++++++++++--------------
 1 file changed, 93 insertions(+), 43 deletions(-)

diff --git a/src/homeassistant.rs b/src/homeassistant.rs
index ec230a3..0a72541 100644
--- a/src/homeassistant.rs
+++ b/src/homeassistant.rs
@@ -1,4 +1,4 @@
-use std::{fmt::Display, io::Write};
+use std::{fmt::Display, io::Write, time::Duration};
 
 use async_tungstenite::{tokio::connect_async, tungstenite};
 use color_eyre::{eyre::bail, Result};
@@ -60,13 +60,15 @@ impl HomeAssistantBuilder {
     }
 
     pub(crate) fn build(self) -> HomeAssistant {
-        HomeAssistant::new(
-            self.host.unwrap(),
-            self.entity.unwrap(),
-            self.insecure,
-            self.token.unwrap(),
-            self.format,
-        )
+        HomeAssistant {
+            host: self.host.unwrap(),
+            entity: self.entity.unwrap(),
+            insecure: self.insecure,
+            token: self.token.unwrap(),
+            format: self.format,
+            id: 1,
+            auth_complete: false,
+        }
     }
 }
 
@@ -77,6 +79,7 @@ pub(crate) struct HomeAssistant {
     host: String,
     id: u64,
     format: OutputFormat,
+    auth_complete: bool,
 }
 
 #[derive(Debug, Default, Clone, Serialize, Deserialize)]
@@ -209,6 +212,18 @@ impl Message {
         msg
     }
 
+    fn ping() -> Self {
+        let mut msg = Self::default();
+        msg.message_type = MessageType::Ping;
+        msg
+    }
+
+    fn pong() -> Self {
+        let mut msg = Self::default();
+        msg.message_type = MessageType::Pong;
+        msg
+    }
+
     fn to_json(self) -> String {
         serde_json::to_string(&self).unwrap()
     }
@@ -219,22 +234,23 @@ impl Message {
 }
 
 impl HomeAssistant {
-    pub(crate) fn new(
-        host: String,
-        entity: String,
-        insecure: bool,
-        token: String,
-        format: OutputFormat,
-    ) -> Self {
-        Self {
-            token,
-            entity,
-            insecure,
-            host,
-            format,
-            id: 1,
-        }
-    }
+    // pub(crate) fn new(
+    //     host: String,
+    //     entity: String,
+    //     insecure: bool,
+    //     token: String,
+    //     format: OutputFormat,
+    // ) -> Self {
+    //     Self {
+    //         token,
+    //         entity,
+    //         insecure,
+    //         host,
+    //         format,
+    //         id: 1,
+    //         auth_complete: false,
+    //     }
+    // }
 
     pub(crate) fn builder() -> HomeAssistantBuilder {
         HomeAssistantBuilder::new()
@@ -257,28 +273,48 @@ impl HomeAssistant {
 
         let (mut ws_stream, _) = connect_async(&api_url).await?;
 
-        while let Some(msg) = ws_stream.next().await {
-            let msg = msg?;
+        let mut ping_interval = tokio::time::interval(Duration::from_millis(5000));
 
-            match self.handle_message(msg).await {
-                Ok(response_messages) => {
-                    for response in response_messages {
-                        ws_stream.send(response).await?
-                    }
-                }
-                Err(e) => {
-                    if let Some(err) = e.downcast_ref::<HomeAssistantError>() {
-                        match err {
-                            HomeAssistantError::EmptyMessage => {
-                                debug!("Received empty message, ignoring...")
+        loop {
+            tokio::select! {
+                msg = ws_stream.next() => match msg {
+                    Some(msg) => {
+                        let msg = msg?;
+
+                        match self.handle_message(msg).await {
+                            Ok(response_messages) => {
+                                for response in response_messages {
+                                    ws_stream.send(response).await?
+                                }
                             }
-                            HomeAssistantError::UnhandledMessage(message_type) => debug!(
+                            Err(e) => {
+                                if let Some(err) = e.downcast_ref::<HomeAssistantError>() {
+                                    match err {
+                                        HomeAssistantError::EmptyMessage => {
+                                            debug!("Received empty message, ignoring...")
+                                        }
+                                        HomeAssistantError::UnhandledMessage(message_type) => {
+                                            debug!(
                                 "Received '{message_type}', we're currently not handling that."
-                            ),
-                            HomeAssistantError::JsonError(e) => error!("{e}"),
+                            )
+                                        }
+                                        HomeAssistantError::JsonError(e) => error!("{e}"),
+                                    }
+                                } else {
+                                    error!("{e}");
+                                }
+                            }
                         }
-                    } else {
-                        error!("{e}");
+                    }
+                    None => break,
+                },
+                _ = ping_interval.tick() => {
+                    if self.auth_complete {
+                        let mut ping_message = Message::ping();
+                        ping_message.id = Some(self.incrementing_id());
+                        debug!("[S] Ping!");
+
+                        ws_stream.send(ping_message.to_message()).await?;
                     }
                 }
             }
@@ -315,6 +351,8 @@ impl HomeAssistant {
                 let mut get_states = Message::get_states();
                 get_states.id = Some(self.incrementing_id());
 
+                self.auth_complete = true;
+
                 vec![subscribe.to_message(), get_states.to_message()]
             }
             MessageType::Event => {
@@ -351,6 +389,18 @@ impl HomeAssistant {
 
                 vec![]
             }
+            MessageType::Ping => {
+                debug!("[R] Ping!");
+                let mut pong = Message::pong();
+                pong.id = Some(self.incrementing_id());
+
+                vec![pong.to_message()]
+            }
+            MessageType::Pong => {
+                debug!("[R] Pong!");
+
+                vec![]
+            }
             _ => bail!(HomeAssistantError::UnhandledMessage(
                 message.message_type.to_string()
             )),
@@ -367,7 +417,7 @@ impl HomeAssistant {
             let is_muted = attributes["is_volume_muted"].as_bool().unwrap_or(false);
 
             let volume = if is_muted {
-                String::from("[  ﱝ ]")
+                String::from("[  \u{fc5d} ]")
             } else if volume_raw >= 0. {
                 format!("[{:3.0}%] ", volume_raw * 100.)
             } else {