From afecbd9cd893bef4bbf1d91c2bc8145172a97a01 Mon Sep 17 00:00:00 2001
From: Jake Hillion <jake@hillion.co.uk>
Date: Tue, 24 May 2022 17:22:47 +0100
Subject: [PATCH] tcp listener ctrl-c handler

---
 Cargo.lock               |  1 +
 Cargo.toml               |  1 +
 examples/tls/listener.rs | 55 ++++++++++++++++++++++++++++++++++++++++
 examples/tls/main.rs     | 24 ++----------------
 4 files changed, 59 insertions(+), 22 deletions(-)
 create mode 100644 examples/tls/listener.rs

diff --git a/Cargo.lock b/Cargo.lock
index 86a3f61..e513080 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -132,6 +132,7 @@ dependencies = [
  "exitcode",
  "httparse",
  "ipnetwork",
+ "lazy_static",
  "libc",
  "log",
  "nix",
diff --git a/Cargo.toml b/Cargo.toml
index 7d53d83..9507c40 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -30,6 +30,7 @@ anyhow = "1"
 httparse = "1"
 rustls = "0.20"
 rustls-pemfile = "1"
+lazy_static = "1"
 
 [[bench]]
 name = "clone3"
diff --git a/examples/tls/listener.rs b/examples/tls/listener.rs
new file mode 100644
index 0000000..06f5d0a
--- /dev/null
+++ b/examples/tls/listener.rs
@@ -0,0 +1,55 @@
+use std::fs::File;
+use std::io::ErrorKind;
+use std::net::TcpListener;
+use std::os::unix::io::AsRawFd;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use nix::poll::{poll, PollFd, PollFlags};
+use nix::sys::signal::{signal, SigHandler, Signal};
+use nix::Error as NixError;
+
+use lazy_static::lazy_static;
+
+lazy_static! {
+    static ref RUNNING: AtomicBool = AtomicBool::new(true);
+}
+
+pub(crate) fn handler(tls_handler_trigger: File, listener: TcpListener) -> i32 {
+    println!("connection_listener entered");
+
+    // SAFETY: only unsafe if you use the result
+    unsafe { signal(Signal::SIGINT, SigHandler::Handler(handle_sigint)) }.unwrap();
+
+    listener.set_nonblocking(true).unwrap();
+
+    let mut to_poll = [PollFd::new(listener.as_raw_fd(), PollFlags::POLLIN)];
+    while RUNNING.load(Ordering::Relaxed) {
+        if let Err(e) = poll(&mut to_poll, 1000) {
+            if e == NixError::EINTR {
+                continue; // timed out
+            }
+            Err(e).unwrap()
+        }
+
+        let stream = match listener.accept() {
+            Ok(s) => s,
+            Err(e) => {
+                if e.kind() != ErrorKind::WouldBlock {
+                    Err(e).unwrap()
+                } else {
+                    continue;
+                }
+            }
+        };
+
+        println!("received a new connection");
+        super::tls_handler(&tls_handler_trigger, stream.0);
+    }
+
+    exitcode::OK
+}
+
+extern "C" fn handle_sigint(signal: libc::c_int) {
+    let signal = Signal::try_from(signal).unwrap();
+    RUNNING.store(signal != Signal::SIGINT, Ordering::Relaxed);
+}
diff --git a/examples/tls/main.rs b/examples/tls/main.rs
index 94b75f0..d2f95e6 100644
--- a/examples/tls/main.rs
+++ b/examples/tls/main.rs
@@ -1,4 +1,5 @@
 mod http;
+mod listener;
 mod tls;
 
 use std::fs::File;
@@ -41,29 +42,8 @@ fn connection_listener_entrypoint() {
         .expect("tcp listener should be a file descriptor");
     let tcp_listener = unsafe { TcpListener::from_raw_fd(tcp_listener) };
 
-    // actual function body
-    fn connection_listener(tls_handler_trigger: File, tcp_listener: TcpListener) -> i32 {
-        println!("connection_listener entered");
-
-        // handle incoming connections
-        for stream in tcp_listener.incoming() {
-            let stream = match stream {
-                Ok(s) => s,
-                Err(e) => {
-                    println!("connection listener: error: {}", e);
-                    return 1;
-                }
-            };
-
-            println!("received a new connection");
-            tls_handler(&tls_handler_trigger, stream);
-        }
-
-        exitcode::OK
-    }
-
     // run function
-    std::process::exit(connection_listener(tls_handler_trigger, tcp_listener));
+    std::process::exit(listener::handler(tls_handler_trigger, tcp_listener));
 }
 
 fn tls_handler(trigger_socket: &File, stream: TcpStream) {
-- 
2.46.0