From fd110b89abd714f3663ea07a7335549455863816 Mon Sep 17 00:00:00 2001
From: Ulrik <ulrikm@spotify.com>
Date: Thu, 20 Oct 2022 14:03:30 +0200
Subject: [PATCH] Add test for behavior when limits is hit for fd:s or threads

Right now, when the limit on file-descriptors are hit, the server does not
typically recover. Incoming-requests returns EndOfStream, panic poisons a
Mutex, and recovery is unclear.
---
 Cargo.toml      |   1 +
 tests/limits.rs | 108 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+)
 create mode 100644 tests/limits.rs

diff --git a/Cargo.toml b/Cargo.toml
index 1bc68513..ce8f1a34 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -31,6 +31,7 @@ zeroize = { version = "1", optional = true }
 rustc-serialize = "0.3"
 sha1 = "0.6.0"
 fdlimit = "0.1"
+libc = "0.2"
 
 [package.metadata.docs.rs]
 # Enable just one SSL implementation
diff --git a/tests/limits.rs b/tests/limits.rs
new file mode 100644
index 00000000..d072ce38
--- /dev/null
+++ b/tests/limits.rs
@@ -0,0 +1,108 @@
+extern crate tiny_http;
+
+use std::io::{Read, Write};
+use std::net::{TcpListener, TcpStream};
+
+use libc::{__rlimit_resource_t, getrlimit, rlimit, setrlimit, RLIMIT_NOFILE, RLIMIT_NPROC};
+
+fn set_limit(limit: __rlimit_resource_t, value: u64) {
+    let mut current = rlimit {
+        rlim_cur: 0,
+        rlim_max: 0,
+    };
+    assert_eq!(0, unsafe { getrlimit(limit, &mut current) });
+    current.rlim_cur = value;
+    assert_eq!(0, unsafe { setrlimit(limit, &mut current) });
+}
+
+struct ServerProcess {
+    pid: libc::pid_t,
+}
+
+impl ServerProcess {
+    fn start(setup: impl FnOnce()) -> (std::net::SocketAddr, ServerProcess) {
+        let listener = TcpListener::bind("0.0.0.0:0").unwrap();
+        let pid = unsafe { libc::fork() };
+        if pid == 0 {
+            setup();
+            let server = tiny_http::Server::from_listener(listener, None).unwrap();
+            for req in server.incoming_requests() {
+                req.respond(tiny_http::Response::empty(204)).unwrap();
+            }
+            std::process::exit(0);
+        } else {
+            let addr = listener.local_addr().unwrap();
+            drop(listener);
+            (addr, Self { pid })
+        }
+    }
+}
+
+impl Drop for ServerProcess {
+    fn drop(&mut self) {
+        unsafe {
+            libc::kill(self.pid, libc::SIGKILL);
+            libc::waitpid(self.pid, std::ptr::null_mut(), 0);
+        }
+    }
+}
+
+fn make_request_with_keep_alive(addr: std::net::SocketAddr) -> std::io::Result<TcpStream> {
+    TcpStream::connect(addr).and_then(|mut s| {
+        let mut buf = [0; 1024];
+        write!(
+            s,
+            "GET / HTTP/1.1\r\nHost: localhost\r\nConnection: keep-alive\r\n\r\n"
+        )?;
+        s.read(&mut buf)?;
+        Ok(s)
+    })
+}
+
+#[test]
+fn survives_fd_limit_1() {
+    let (addr, _server) = ServerProcess::start(|| {
+        // One off from survives_fd_limit_2 to cover both fd-point of the fd-duplication
+         set_limit(RLIMIT_NOFILE, 8);
+    });
+
+    let clients: Vec<_> = (0..10)
+        .map(|_| make_request_with_keep_alive(addr))
+        .collect();
+    assert!(clients.iter().any(Result::is_err));
+    drop(clients);
+
+    assert!(TcpStream::connect(addr).is_ok())
+}
+
+#[test]
+fn survives_fd_limit_2() {
+    let (addr, _server) = ServerProcess::start(|| {
+        // One off from survives_fd_limit_1 to cover both fd-point of the fd-duplication
+        set_limit(RLIMIT_NOFILE, 7);
+    });
+
+    let clients: Vec<_> = (0..10)
+        .map(|_| make_request_with_keep_alive(addr))
+        .collect();
+    assert!(clients.iter().any(Result::is_err));
+
+    drop(clients);
+
+    assert!(TcpStream::connect(addr).is_ok())
+}
+
+#[test]
+fn survives_proc_limit() {
+    let (addr, _server) = ServerProcess::start(|| {
+        set_limit(RLIMIT_NPROC, 7);
+    });
+
+    let clients: Vec<_> = (0..10)
+        .map(|_| make_request_with_keep_alive(addr))
+        .collect();
+    assert!(clients.iter().any(Result::is_err));
+    drop(clients);
+
+    assert!(TcpStream::connect(addr).is_ok())
+}