From dd9e22a78b9a7a16933c71738f03a7560d8eacef Mon Sep 17 00:00:00 2001
From: Orbital <orbitalturtle@protonmail.com>
Date: Fri, 16 Aug 2024 22:53:53 -0500
Subject: [PATCH] main: handle sigterm/sigint signals

---
 Cargo.lock  | 20 ++++++++++++++++++++
 Cargo.toml  |  1 +
 src/main.rs | 21 +++++++++++++++++++--
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index a85d392a..36c2c8c5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1231,6 +1231,7 @@ dependencies = [
  "rand_chacha",
  "rand_core 0.6.4",
  "rcgen",
+ "signal-hook",
  "tempfile",
  "tokio",
  "tonic 0.11.0",
@@ -2051,6 +2052,25 @@ dependencies = [
  "unsafe-libyaml",
 ]
 
+[[package]]
+name = "signal-hook"
+version = "0.3.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8621587d4798caf8eb44879d42e56b9a93ea5dcd315a6487c357130095b62801"
+dependencies = [
+ "libc",
+ "signal-hook-registry",
+]
+
+[[package]]
+name = "signal-hook-registry"
+version = "1.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "slab"
 version = "0.4.9"
diff --git a/Cargo.toml b/Cargo.toml
index 5fcf9f45..47011ce5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -32,6 +32,7 @@ configure_me = "0.4.0"
 bytes = "1.4.0"
 triggered = "0.1.2"
 prost = "0.12"
+signal-hook = "0.3.17"
 
 [dev-dependencies]
 bitcoincore-rpc = { package="core-rpc", version = "0.17.0" }
diff --git a/src/main.rs b/src/main.rs
index 2c76353c..daae8b19 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -19,10 +19,13 @@ use lndk::{
 };
 use lndkrpc::offers_server::OffersServer;
 use log::{error, info};
+use signal_hook::consts::{SIGINT, SIGTERM};
+use signal_hook::iterator::Signals;
 use std::fs::create_dir_all;
 use std::path::PathBuf;
 use std::process::exit;
 use std::sync::Arc;
+use std::thread;
 use tokio::select;
 use tonic::transport::{Server, ServerTlsConfig};
 use tonic_lnd::lnrpc::GetInfoRequest;
@@ -53,13 +56,27 @@ async fn main() -> Result<(), ()> {
     let lnd_args = LndCfg::new(config.address, creds.clone());
 
     let (shutdown, listener) = triggered::trigger();
-    let signals = LifecycleSignals { shutdown, listener };
+    let signals = LifecycleSignals {
+        shutdown: shutdown.clone(),
+        listener: listener.clone(),
+    };
     let args = Cfg {
         lnd: lnd_args,
         signals,
         skip_version_check: config.skip_version_check,
     };
 
+    let mut shutdown_signals = Signals::new([SIGINT, SIGTERM]).map_err(|e| {
+        error!("Error initializing shutdown signals: {e}.");
+    })?;
+
+    thread::spawn(move || {
+        if let Some(sig) = shutdown_signals.forever().next() {
+            info!("Received shutdown signal {:?}. Shutting down.", sig);
+            shutdown.trigger();
+        }
+    });
+
     let response_invoice_timeout = config.response_invoice_timeout;
     if let Some(timeout) = response_invoice_timeout {
         if timeout == 0 {
@@ -113,7 +130,7 @@ async fn main() -> Result<(), ()> {
         .tls_config(ServerTlsConfig::new().identity(identity))
         .expect("couldn't configure tls")
         .add_service(OffersServer::new(server))
-        .serve(addr);
+        .serve_with_shutdown(addr, listener);
 
     info!("Starting lndk's grpc server at address {grpc_host}:{grpc_port}");