From e064263f9671e5fb9166478f9d0c5ae4a3014c91 Mon Sep 17 00:00:00 2001
From: pep <pep@bouah.net>
Date: Sat, 3 May 2025 13:29:43 +0200
Subject: [PATCH] Cleanly shutdown on ctrl_c

Signed-off-by: pep <pep@bouah.net>
---
 src/bot.rs   | 11 ++++++++++-
 src/error.rs |  9 +++++++++
 src/main.rs  | 56 ++++++++++++++++++++++++++++++++--------------------
 3 files changed, 54 insertions(+), 22 deletions(-)

diff --git a/src/bot.rs b/src/bot.rs
index 081dc27..968979e 100644
--- a/src/bot.rs
+++ b/src/bot.rs
@@ -14,9 +14,10 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 use crate::hooks::{format_hook, Hook};
+use crate::Error;
 
 use log::debug;
-use tokio::sync::mpsc;
+use tokio::{signal::ctrl_c, sync::mpsc};
 use xmpp::jid::{BareJid, Jid, ResourcePart};
 use xmpp::parsers::message::MessageType;
 use xmpp::{
@@ -85,6 +86,9 @@ impl XmppClient {
 	pub async fn receive(&mut self, mut rx: mpsc::UnboundedReceiver<Hook>) {
 		loop {
 			tokio::select! {
+				_ = ctrl_c() => {
+						return;  // Disconnecting
+				},
 				_ = self.next() => (),
 				wh = rx.recv() => {
 					if let Some(hook) = wh {
@@ -114,4 +118,9 @@ impl XmppClient {
 			debug!("XMPP Bot Processed Hook");
 		}
 	}
+
+	pub async fn disconnect(self) -> Result<(), Error> {
+		log::info!("Disconnecting...");
+		Ok(self.agent.disconnect().await?)
+	}
 }
diff --git a/src/error.rs b/src/error.rs
index c00d02c..8b55931 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -21,6 +21,7 @@ use std::str::Utf8Error;
 use hex::FromHexError;
 use hmac::digest::InvalidLength as HmacInvalidLength;
 use hyper::StatusCode;
+use xmpp::tokio_xmpp::Error as TokioXmppError;
 
 #[derive(Debug)]
 pub enum Error {
@@ -38,6 +39,7 @@ pub enum Error {
 	Toml(toml::de::Error),
 	Utf8(Utf8Error),
 	Var(VarError),
+	Xmpp(TokioXmppError),
 }
 
 impl Error {
@@ -74,6 +76,7 @@ impl std::fmt::Display for Error {
 			Error::Toml(e) => write!(fmt, "toml deserialization error: {}", e),
 			Error::Utf8(e) => write!(fmt, "Utf8 error: {}", e),
 			Error::Var(e) => write!(fmt, "Var error: {}", e),
+			Error::Xmpp(e) => write!(fmt, "Xmpp error: {}", e),
 		}
 	}
 }
@@ -125,3 +128,9 @@ impl From<VarError> for Error {
 		Error::Var(err)
 	}
 }
+
+impl From<TokioXmppError> for Error {
+	fn from(err: TokioXmppError) -> Error {
+		Error::Xmpp(err)
+	}
+}
diff --git a/src/main.rs b/src/main.rs
index 643d14b..9d002b4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -33,10 +33,10 @@ use clap::{command, value_parser, Arg};
 use hyper::{server::conn::http1, service::service_fn};
 use hyper_util::rt::tokio::{TokioIo, TokioTimer};
 use log::error;
-use tokio::{net::TcpListener, sync::mpsc};
+use tokio::{net::TcpListener, signal::ctrl_c, sync::mpsc};
 
 #[tokio::main]
-async fn main() -> Result<!, Error> {
+async fn main() -> Result<(), Error> {
 	pretty_env_logger::init();
 
 	let matches = command!()
@@ -60,9 +60,12 @@ async fn main() -> Result<!, Error> {
 		config.nickname,
 	);
 
-	tokio::task::spawn(async move {
+	let xmpp_handle = tokio::task::spawn(async move {
 		bot.receive(value_rx).await;
-		error!("XMPP client exited early");
+		match bot.disconnect().await {
+			Err(err) => error!("XMPP disconnect error: {err}"),
+			_ => (),
+		}
 	});
 
 	let tcp_bind = TcpListener::bind(config.addr).await?;
@@ -71,24 +74,35 @@ async fn main() -> Result<!, Error> {
 		let value_tx = value_tx.clone();
 		let secret = config.secret.clone();
 
-		if let Ok((tcp, _)) = tcp_bind.accept().await {
-			let io = TokioIo::new(tcp);
-			tokio::task::spawn(async move {
-				if let Err(err) = http1::Builder::new()
-					.timer(TokioTimer::new())
-					.serve_connection(
-						io,
-						service_fn(|request| {
-							let value_tx = value_tx.clone();
-							let secret = secret.clone();
-							async move { hooks(request, &secret, value_tx).await }
-						}),
-					)
-					.await
-				{
-					println!("Error serving connection: {:?}", err);
+		tokio::select! {
+			accept = tcp_bind.accept() => {
+				if let Ok((tcp, _)) = accept {
+					let io = TokioIo::new(tcp);
+					tokio::task::spawn(async move {
+						if let Err(err) = http1::Builder::new()
+							.timer(TokioTimer::new())
+							.serve_connection(
+								io,
+								service_fn(|request| {
+									let value_tx = value_tx.clone();
+									let secret = secret.clone();
+									async move { hooks(request, &secret, value_tx).await }
+								}),
+							)
+							.await
+						{
+							println!("Error serving connection: {:?}", err);
+						}
+					});
 				}
-			});
+			},
+			_ = ctrl_c() => {
+				break;
+			},
 		}
 	}
+
+	let _ = xmpp_handle.await;
+
+	Ok(())
 }