Cleanly shutdown on ctrl_c

Signed-off-by: pep <pep@bouah.net>

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?)
+	}
 }

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)
+	}
+}

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 {
+		tokio::select! {
-			let io = TokioIo::new(tcp);
+			accept = tcp_bind.accept() => {
-			tokio::task::spawn(async move {
+				if let Ok((tcp, _)) = accept {
-				if let Err(err) = http1::Builder::new()
+					let io = TokioIo::new(tcp);
-					.timer(TokioTimer::new())
+					tokio::task::spawn(async move {
-					.serve_connection(
+						if let Err(err) = http1::Builder::new()
-						io,
+							.timer(TokioTimer::new())
-						service_fn(|request| {
+							.serve_connection(
-							let value_tx = value_tx.clone();
+								io,
-							let secret = secret.clone();
+								service_fn(|request| {
-							async move { hooks(request, &secret, value_tx).await }
+									let value_tx = value_tx.clone();
-						}),
+									let secret = secret.clone();
-					)
+									async move { hooks(request, &secret, value_tx).await }
-					.await
+								}),
-				{
+							)
-					println!("Error serving connection: {:?}", err);
+							.await
+						{
+							println!("Error serving connection: {:?}", err);
+						}
+					});
 				}
-			});
+			},
+			_ = ctrl_c() => {
+				break;
+			},
 		}
 	}
+
+	let _ = xmpp_handle.await;
+
+	Ok(())
 }