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 .
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) {
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 for Error {
Error::Var(err)
}
}
+
+impl From 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 {
+async fn main() -> Result<(), Error> {
pretty_env_logger::init();
let matches = command!()
@@ -60,9 +60,12 @@ async fn main() -> Result {
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 {
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(())
}