From b2cbd90d266aedc83de435bc1e80fbf8fde00f12 Mon Sep 17 00:00:00 2001
From: Moritz Ruth <dev@moritzruth.de>
Date: Tue, 28 Feb 2023 23:23:46 +0100
Subject: [PATCH] And again

---
 .idea/runConfigurations/Run__dry_.xml |  2 +-
 src/config.rs                         | 26 +++++-----
 src/main.rs                           | 63 ++++++++++++++++++++----
 src/modules/mod.rs                    | 46 ++++++++++-------
 src/modules/power.rs                  | 30 ++++-------
 src/mqtt.rs                           | 71 +++++++++++++++++++++++----
 6 files changed, 165 insertions(+), 73 deletions(-)

diff --git a/.idea/runConfigurations/Run__dry_.xml b/.idea/runConfigurations/Run__dry_.xml
index 4fd01b9..4f00742 100644
--- a/.idea/runConfigurations/Run__dry_.xml
+++ b/.idea/runConfigurations/Run__dry_.xml
@@ -10,7 +10,7 @@
     <option name="buildTarget" value="REMOTE" />
     <option name="backtrace" value="SHORT" />
     <envs>
-      <env name="RUST_LOG" value="hassliebe=debug" />
+      <env name="RUST_LOG" value="hassliebe=trace" />
     </envs>
     <option name="isRedirectInput" value="false" />
     <option name="redirectInputPath" value="" />
diff --git a/src/config.rs b/src/config.rs
index f831bde..6b4ca81 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -2,8 +2,9 @@ use std::borrow::ToOwned;
 use std::fs;
 use std::fs::File;
 use std::io::{ErrorKind, Read, Write};
+use std::path::Path;
 
-use anyhow::{anyhow, bail, Context, Result};
+use anyhow::{bail, Context, Result};
 use rand::distributions::Alphanumeric;
 use rand::Rng;
 use regex::Regex;
@@ -48,10 +49,10 @@ pub struct Config {
 }
 
 fn validate_unique_id(value: &str) -> Result<(), ValidationError> {
-    if Regex::new(r"^[a-zA-Z0-9]+(_[a-zA-Z0-9])*$").unwrap().is_match(value) {
+    if Regex::new(r"^[a-zA-Z0-9]+(_[a-zA-Z0-9]+)*$").unwrap().is_match(value) {
         Ok(())
     } else {
-        Err(ValidationError::new("invalid_unique_id"))
+        Err(ValidationError::new("does not match regex"))
     }
 }
 
@@ -75,37 +76,32 @@ fn create_example_config() -> Config {
     }
 }
 
-pub fn load() -> Result<Option<Config>> {
-    let dirs = directories::ProjectDirs::from("", "", "Hassliebe")
-        .ok_or_else(|| anyhow!("Could not determine a valid home directory path. Please specify a custom config file path via HASS_CONFIG"))?;
-
-    let mut path = dirs.config_dir().to_owned();
-    path.push("config.toml");
-
-    match File::open(&path) {
+pub(crate) fn load(config_file_path: &Path) -> Result<Option<Config>> {
+    match File::open(config_file_path) {
         Ok(mut file) => {
-            log::info!("Reading config file: {}", path.to_string_lossy());
+            log::info!("Reading config file: {}", config_file_path.to_string_lossy());
 
             let mut string_content = String::new();
             file.read_to_string(&mut string_content).context("while reading the config file")?;
 
             let parsed = toml::from_str::<Config>(string_content.as_str()).context("while parsing the config file")?;
+            parsed.validate().context("while validating the config file")?;
 
             Ok(Some(parsed))
         }
 
         Err(error) if error.kind() == ErrorKind::NotFound => {
-            if let Some(parent) = path.parent() {
+            if let Some(parent) = config_file_path.parent() {
                 fs::create_dir_all(parent)?;
             }
 
-            let mut file = File::create(&path).context("while creating the default config file")?;
+            let mut file = File::create(config_file_path).context("while creating the default config file")?;
 
             let default = toml::to_string::<Config>(&create_example_config()).expect("create_example_config() should be valid");
 
             file.write_all(default.as_bytes())?;
 
-            log::warn!("Created an example config file at {}", path.to_string_lossy());
+            log::warn!("Created an example config file at {}", config_file_path.to_string_lossy());
             log::warn!("Make sure to edit the file before running again.");
 
             Ok(None)
diff --git a/src/main.rs b/src/main.rs
index 8dd67af..17ed216 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,19 +1,54 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
+use std::path::{Path, PathBuf};
 use std::process::exit;
 
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 
-use crate::modules::ModuleContext;
+use crate::modules::{ModuleContext, ModuleContextMqtt};
+use crate::mqtt::OwnedTopicsService;
 
-mod modules;
 mod config;
+mod modules;
 mod mqtt;
 
+struct Paths {
+    data_directory: Box<Path>,
+    config: Box<Path>,
+}
+
+async fn get_paths_and_create_directories() -> Result<Paths> {
+    let project_dirs = directories::ProjectDirs::from("", "", "Hassliebe");
+
+    let paths = Paths {
+        config: std::env::var_os("HASS_CONFIG")
+            .map(PathBuf::from)
+            // I don’t like this clone
+            .or(project_dirs.clone().map(|d| d.config_dir().to_path_buf().join("config.toml")))
+            .ok_or_else(|| anyhow!("Please specify a config file via HASS_CONFIG"))?
+            .into_boxed_path(),
+        data_directory: std::env::var_os("HASS_DATA_DIR")
+            .map(PathBuf::from)
+            .or(project_dirs.map(|d| d.data_dir().to_path_buf()))
+            .ok_or_else(|| anyhow!("Please specify a data directory via HASS_DATA_DIR"))?
+            .into_boxed_path(),
+    };
+
+    if let Some(p) = paths.config.parent() {
+        tokio::fs::create_dir_all(p).await?;
+    }
+
+    tokio::fs::create_dir_all(&paths.data_directory).await?;
+
+    Ok(paths)
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     env_logger::init();
 
-    let Some(config) = config::load()? else {
+    let paths = get_paths_and_create_directories().await?;
+
+    let Some(config) = config::load(&paths.config)? else {
         exit(exitcode::CONFIG);
     };
 
@@ -21,16 +56,24 @@ async fn main() -> Result<()> {
     let (mqtt_client, event_loop) = mqtt::create_client(&config, &availability_topic).await?;
     let discovery_device_object = mqtt::create_discovery_device_object(&config);
 
+    let owned_topics_service = OwnedTopicsService::new(&paths.data_directory).await?;
+
     let mut module_context = ModuleContext {
         config: &config,
-        mqtt_client: &mqtt_client,
-        mqtt_availability_topic: availability_topic.as_str(),
-        mqtt_discovery_device_object: &discovery_device_object,
-        mqtt_message_handler_by_topic: HashMap::new(),
+        mqtt: ModuleContextMqtt {
+            client: &mqtt_client,
+            availability_topic: availability_topic.as_str(),
+            discovery_device_object: &discovery_device_object,
+            message_handler_by_topic: HashMap::new(),
+            owned_topics: HashSet::new(),
+        },
     };
 
     modules::init_all(&mut module_context).await?;
 
+    owned_topics_service
+        .clear_old_and_save_new(module_context.mqtt.client, &module_context.mqtt.owned_topics)
+        .await?;
+
     mqtt::start_communication(&module_context, event_loop).await
 }
-
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 0cf24ef..8b8f980 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -1,37 +1,49 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use anyhow::Result;
 use json::JsonValue;
-use rumqttc::AsyncClient as MqttClient;
+use rumqttc::{AsyncClient as MqttClient, ClientError, QoS};
 
 pub mod power;
 
 type MqttMessageHandler<'a> = dyn Fn(&str) -> Result<()> + 'a;
 
+pub struct ModuleContextMqtt<'a> {
+    pub discovery_device_object: &'a JsonValue,
+    pub availability_topic: &'a str,
+    pub client: &'a MqttClient,
+    pub message_handler_by_topic: HashMap<String, Box<MqttMessageHandler<'a>>>,
+
+    // Owned topics are topics which a retained message was sent into.
+    pub owned_topics: HashSet<String>,
+}
+
 pub struct ModuleContext<'a> {
     pub config: &'a super::config::Config,
-    pub mqtt_discovery_device_object: &'a JsonValue,
-    pub mqtt_availability_topic: &'a str,
-    pub mqtt_client: &'a MqttClient,
-    pub mqtt_message_handler_by_topic: HashMap<String, Box<MqttMessageHandler<'a>>>,
+    pub mqtt: ModuleContextMqtt<'a>,
 }
 
 impl<'a> ModuleContext<'a> {
-    fn get_mqtt_topic(&self, component_type: &str, entity_id: &str, suffix: &str) -> String {
-        format!("homeassistant/{}/{}/{}", component_type, entity_id, suffix)
-    }
-
     fn get_entity_id(&self, module_id: &str, sub_id: &str) -> String {
         format!("{}_{}_{}", self.config.unique_id, module_id, sub_id)
     }
+}
 
-    fn subscribe_mqtt_topic<F: Fn(&str) -> Result<()> + 'a>(
-        &mut self,
-        topic: impl Into<String>,
-        handler: F,
-    ) {
-        self.mqtt_message_handler_by_topic
-            .insert(topic.into(), Box::new(handler));
+impl<'a> ModuleContextMqtt<'a> {
+    fn get_topic(&self, component_type: &str, entity_id: &str, suffix: &str) -> String {
+        format!("homeassistant/{}/{}/{}", component_type, entity_id, suffix)
+    }
+
+    fn subscribe<F: Fn(&str) -> Result<()> + 'a>(&mut self, topic: impl Into<String>, handler: F) {
+        self.message_handler_by_topic.insert(topic.into(), Box::new(handler));
+    }
+
+    async fn send_retained_message(&mut self, topic: impl Into<String>, message: impl Into<String>) -> std::result::Result<(), ClientError> {
+        let topic = topic.into();
+        let message = message.into();
+
+        self.owned_topics.insert(topic.to_owned());
+        self.client.publish(topic, QoS::AtLeastOnce, true, message).await
     }
 }
 
diff --git a/src/modules/power.rs b/src/modules/power.rs
index 9342e55..cb5b3b0 100644
--- a/src/modules/power.rs
+++ b/src/modules/power.rs
@@ -1,5 +1,4 @@
 use anyhow::Result;
-use rumqttc::QoS;
 use tokio::process::Command;
 
 use super::ModuleContext;
@@ -16,27 +15,20 @@ pub(crate) async fn init(context: &mut ModuleContext<'_>) -> Result<()> {
     Ok(())
 }
 
-async fn init_command_button(
-    context: &mut ModuleContext<'_>,
-    sub_id: &str,
-    name: &str,
-    command: impl Into<String>,
-) -> Result<()> {
+async fn init_command_button(context: &mut ModuleContext<'_>, sub_id: &str, name: &str, command: impl Into<String>) -> Result<()> {
     let command = command.into();
 
     let entity_id = context.get_entity_id(MODULE_ID, sub_id);
-    let command_topic = context.get_mqtt_topic("button", &entity_id, "trigger");
+    let command_topic = context.mqtt.get_topic("button", &entity_id, "trigger");
 
     context
-        .mqtt_client
-        .publish(
-            context.get_mqtt_topic("button", &entity_id, "config"),
-            QoS::AtLeastOnce,
-            true,
+        .mqtt
+        .send_retained_message(
+            context.mqtt.get_topic("button", &entity_id, "config"),
             json::stringify(json::object! {
-                "availability_topic": context.mqtt_availability_topic,
+                "availability_topic": context.mqtt.availability_topic,
                 "command_topic": command_topic.as_str(),
-                "device": context.mqtt_discovery_device_object.clone(),
+                "device": context.mqtt.discovery_device_object.clone(),
                 "icon": "mdi:power",
                 "name": name,
                 "payload_press": BUTTON_TRIGGER_TEXT,
@@ -46,7 +38,7 @@ async fn init_command_button(
         )
         .await?;
 
-    context.subscribe_mqtt_topic(command_topic, move |text| {
+    context.mqtt.subscribe(command_topic, move |text| {
         if text == BUTTON_TRIGGER_TEXT {
             run_command(command.clone());
         }
@@ -62,11 +54,7 @@ fn run_command(command: String) {
         let is_dry_run = cfg!(feature = "dry_run");
 
         let status = {
-            log::info!(
-                "Executing command{}: {}",
-                if is_dry_run { "(dry run)" } else { "" },
-                command
-            );
+            log::info!("Executing command{}: {}", if is_dry_run { " (dry run)" } else { "" }, command);
 
             let mut command_parts = command.split(' ').collect::<Vec<_>>();
             let mut actual_command = Command::new(command_parts[0]);
diff --git a/src/mqtt.rs b/src/mqtt.rs
index ba7259f..8e02c84 100644
--- a/src/mqtt.rs
+++ b/src/mqtt.rs
@@ -1,17 +1,22 @@
+use std::collections::HashSet;
+use std::io::SeekFrom;
+use std::path::Path;
 use std::time::Duration;
 
 use anyhow::Result;
 use json::JsonValue;
-use rumqttc::{AsyncClient as MqttClient, AsyncClient, EventLoop, LastWill, MqttOptions, Packet, SubscribeFilter};
+use rumqttc::{AsyncClient as MqttClient, EventLoop, LastWill, MqttOptions, Packet, SubscribeFilter};
 use rumqttc::Event::Incoming;
 use rumqttc::QoS;
+use tokio::fs::{File, OpenOptions};
+use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
 use tokio::time::{Instant, sleep};
 
 use crate::modules::ModuleContext;
 
 use super::config;
 
-pub async fn create_client(config: &config::Config, availability_topic: &str) -> Result<(AsyncClient, EventLoop)> {
+pub async fn create_client(config: &config::Config, availability_topic: &str) -> Result<(MqttClient, EventLoop)> {
     let mut options = MqttOptions::new(&config.internal.stable_id, config.mqtt.host.to_owned(), config.mqtt.port);
     options.set_clean_session(true);
     options.set_keep_alive(Duration::from_secs(5));
@@ -34,6 +39,55 @@ pub fn create_discovery_device_object(config: &config::Config) -> JsonValue {
     }
 }
 
+pub struct OwnedTopicsService {
+    file: File,
+    old_topics: HashSet<String>,
+}
+
+impl OwnedTopicsService {
+    pub async fn new(data_directory_path: &Path) -> Result<OwnedTopicsService> {
+        let path = data_directory_path.join("owned_topics");
+        let mut file = OpenOptions::new().write(true).read(true).create(true).open(path).await?;
+
+        let mut content = String::new();
+        file.read_to_string(&mut content).await?;
+
+        let old_topics = content.split_terminator('\n').skip(1).map(|s| s.to_owned()).collect::<HashSet<_>>();
+
+        Ok(OwnedTopicsService { file, old_topics })
+    }
+
+    pub async fn clear_old_and_save_new(mut self, mqtt_client: &MqttClient, new_topics: &HashSet<String>) -> Result<()> {
+        let unused_topics = self.old_topics.difference(new_topics).map(|s| s.to_owned()).collect::<Vec<_>>();
+
+        log::info!(
+            "{} unused owned topics will be cleared. Now using {} owned topics.",
+            unused_topics.len(),
+            new_topics.len()
+        );
+
+        for topic in unused_topics.iter() {
+            log::trace!("Clearing owned topic: {}", topic);
+
+            // Deletes the retained message
+            mqtt_client.publish(topic, QoS::AtLeastOnce, true, Vec::new()).await?;
+        }
+
+        let mut new_content = String::new();
+        new_content.push_str("# DO NOT EDIT THIS FILE. It is automatically generated at each run of Hassliebe.\n");
+
+        for topic in new_topics {
+            new_content.push_str(topic);
+            new_content.push('\n')
+        }
+
+        self.file.set_len(0).await?;
+        self.file.seek(SeekFrom::Start(0)).await?;
+        self.file.write_all(new_content.as_bytes()).await?;
+        Ok(())
+    }
+}
+
 #[derive(Debug, PartialEq, Eq)]
 enum ConnectionState {
     NotConnected,
@@ -50,10 +104,12 @@ pub async fn start_communication(context: &ModuleContext<'_>, mut event_loop: Ev
     log::info!("Connecting to MQTT broker at {}:{}", context.config.mqtt.host, context.config.mqtt.port);
 
     context
-        .mqtt_client
+        .mqtt
+        .client
         .subscribe_many(
             context
-                .mqtt_message_handler_by_topic
+                .mqtt
+                .message_handler_by_topic
                 .keys()
                 .map(|k| SubscribeFilter::new(k.to_owned(), QoS::AtLeastOnce)),
         )
@@ -103,6 +159,7 @@ pub async fn start_communication(context: &ModuleContext<'_>, mut event_loop: Ev
                     _ => {}
                 }
             }
+
             Ok(event) => match event {
                 Incoming(Packet::ConnAck(_)) => {
                     if connection_state == ConnectionState::NotConnected {
@@ -117,15 +174,11 @@ pub async fn start_communication(context: &ModuleContext<'_>, mut event_loop: Ev
                 Incoming(Packet::Publish(message)) => {
                     let text = std::str::from_utf8(&message.payload)?;
 
-                    if let Some(handler) = context.mqtt_message_handler_by_topic.get(message.topic.as_str()) {
+                    if let Some(handler) = context.mqtt.message_handler_by_topic.get(message.topic.as_str()) {
                         handler(text)?;
                     }
                 }
 
-                Incoming(packet) => {
-                    log::trace!("Unhandled packet received: {:?}", packet)
-                }
-
                 _ => {}
             },
         }