itsi 0.1.6 → 0.1.8

data/gems/scheduler/ext/itsi_server/Cargo.toml

@@ -19,6 +19,7 @@ parking_lot = "0.12.3"
 rustls-pemfile = "2.2.0"
 tokio-rustls = "0.26.2"
 bytes = "1.3"
+tokio-rustls-acme = "0.6.0"
 rcgen = { version = "0.13.2", features = ["x509-parser", "pem"] }
 base64 = "0.22.1"
 http-body-util = "0.1.2"
@@ -39,8 +40,8 @@ httparse = "1.10.1"
 async-channel = "2.3.1"
 tempfile = "3.18.0"
 sysinfo = "0.33.1"
-tokio-rustls-acme = "0.6.0"
 rustls = "0.23.23"
 fs2 = "0.4.3"
 ring = "0.17.14"
 async-trait = "0.1.87"
+dirs = "6.0.0"

data/gems/scheduler/ext/itsi_server/src/env.rs

@@ -0,0 +1,43 @@
+use std::{
+    env::{var, VarError},
+    path::PathBuf,
+    sync::LazyLock,
+};
+
+type StringVar = LazyLock<String>;
+type MaybeStringVar = LazyLock<Result<String, VarError>>;
+type PathVar = LazyLock<PathBuf>;
+
+/// ACME Configuration for auto-generating production certificates
+/// *ITSI_ACME_CACHE_DIR* - Directory to store cached certificates
+/// so that these are not regenerated every time the server starts
+pub static ITSI_ACME_CACHE_DIR: StringVar = LazyLock::new(|| {
+    var("ITSI_ACME_CACHE_DIR").unwrap_or_else(|_| "./.rustls_acme_cache".to_string())
+});
+
+/// *ITSI_ACME_CONTACT_EMAIL* - Contact Email address to provide to ACME server during certificate renewal
+pub static ITSI_ACME_CONTACT_EMAIL: MaybeStringVar =
+    LazyLock::new(|| var("ITSI_ACME_CONTACT_EMAIL"));
+
+/// *ITSI_ACME_CA_PEM_PATH* - Optional CA Pem path, used for testing with non-trusted CAs for certifcate generation.
+pub static ITSI_ACME_CA_PEM_PATH: MaybeStringVar = LazyLock::new(|| var("ITSI_ACME_CA_PEM_PATH"));
+
+/// *ITSI_ACME_DIRECTORY_URL* - Directory URL to use for ACME certificate generation.
+pub static ITSI_ACME_DIRECTORY_URL: StringVar = LazyLock::new(|| {
+    var("ITSI_ACME_DIRECTORY_URL")
+        .unwrap_or_else(|_| "https://acme-v02.api.letsencrypt.org/directory".to_string())
+});
+
+/// *ITSI_ACME_LOCK_FILE_NAME* - Name of the lock file used to prevent concurrent certificate generation.
+pub static ITSI_ACME_LOCK_FILE_NAME: StringVar =
+    LazyLock::new(|| var("ITSI_ACME_LOCK_FILE_NAME").unwrap_or(".acme.lock".to_string()));
+
+pub static ITSI_LOCAL_CA_DIR: PathVar = LazyLock::new(|| {
+    var("ITSI_LOCAL_CA_DIR")
+        .map(PathBuf::from)
+        .unwrap_or_else(|_| {
+            dirs::home_dir()
+                .expect("Failed to find HOME directory when initializing ITSI_LOCAL_CA_DIR")
+                .join(".itsi")
+        })
+});

data/gems/scheduler/ext/itsi_server/src/lib.rs

@@ -6,6 +6,7 @@ use server::{itsi_server::Server, signal::reset_signal_handlers};
 use tracing::*;
 
 pub mod body_proxy;
+pub mod env;
 pub mod request;
 pub mod response;
 pub mod server;

data/gems/scheduler/ext/itsi_server/src/request/itsi_request.rs

@@ -6,7 +6,7 @@ use crate::{
     response::itsi_response::ItsiResponse,
     server::{
         itsi_server::{RequestJob, Server},
-        listener::{SockAddr, TokioListener},
+        listener::{ListenerInfo, SockAddr},
         serve_strategy::single_mode::RunningPhase,
     },
 };
@@ -44,7 +44,7 @@ pub struct ItsiRequest {
     pub remote_addr: String,
     pub version: String,
     #[debug(skip)]
-    pub(crate) listener: Arc<TokioListener>,
+    pub(crate) listener: Arc<ListenerInfo>,
     #[debug(skip)]
     pub server: Arc<Server>,
     pub response: ItsiResponse,
@@ -128,7 +128,7 @@ impl ItsiRequest {
         hyper_request: Request<Incoming>,
         sender: async_channel::Sender<RequestJob>,
         server: Arc<Server>,
-        listener: Arc<TokioListener>,
+        listener: Arc<ListenerInfo>,
         addr: SockAddr,
         shutdown_rx: watch::Receiver<RunningPhase>,
     ) -> itsi_error::Result<Response<BoxBody<Bytes, Infallible>>> {
@@ -153,7 +153,7 @@ impl ItsiRequest {
         request: Request<Incoming>,
         sock_addr: SockAddr,
         server: Arc<Server>,
-        listener: Arc<TokioListener>,
+        listener: Arc<ListenerInfo>,
     ) -> (ItsiRequest, mpsc::Receiver<Option<Bytes>>) {
         let (parts, body) = request.into_parts();
         let body = if server.stream_body.is_some_and(|f| f) {
@@ -231,7 +231,7 @@ impl ItsiRequest {
             .uri
             .host()
             .map(|host| host.to_string())
-            .unwrap_or_else(|| self.listener.host()))
+            .unwrap_or_else(|| self.listener.host.clone()))
     }
 
     pub(crate) fn scheme(&self) -> MagnusResult<String> {
@@ -240,7 +240,7 @@ impl ItsiRequest {
             .uri
             .scheme()
             .map(|scheme| scheme.to_string())
-            .unwrap_or_else(|| self.listener.scheme()))
+            .unwrap_or_else(|| self.listener.scheme.clone()))
     }
 
     pub(crate) fn headers(&self) -> MagnusResult<Vec<(String, &str)>> {
@@ -264,7 +264,7 @@ impl ItsiRequest {
     }
 
     pub(crate) fn port(&self) -> MagnusResult<u16> {
-        Ok(self.parts.uri.port_u16().unwrap_or(self.listener.port()))
+        Ok(self.parts.uri.port_u16().unwrap_or(self.listener.port))
     }
 
     pub(crate) fn body(&self) -> MagnusResult<Value> {

data/gems/scheduler/ext/itsi_server/src/server/bind.rs

@@ -9,6 +9,7 @@ use std::{
     path::PathBuf,
     str::FromStr,
 };
+
 #[derive(Debug, Clone)]
 pub enum BindAddress {
     Ip(IpAddr),
@@ -129,13 +130,13 @@ impl FromStr for Bind {
             BindProtocol::Unix => None,
             BindProtocol::Unixs => Some(configure_tls(host, &options)?),
         };
-
-        Ok(Self {
+        let bind = Self {
             address,
             port,
             protocol,
             tls_config,
-        })
+        };
+        Ok(bind)
     }
 }
 

data/gems/scheduler/ext/itsi_server/src/server/itsi_server.rs

@@ -19,7 +19,7 @@ use parking_lot::Mutex;
 use std::{cmp::max, ops::Deref, sync::Arc};
 use tracing::{info, instrument};
 
-static DEFAULT_BIND: &str = "localhost:3000";
+static DEFAULT_BIND: &str = "http://localhost:3000";
 
 #[magnus::wrap(class = "Itsi::Server", free_immediately, size)]
 #[derive(Clone)]
@@ -231,13 +231,6 @@ impl Server {
             drop(strategy);
             Ok(())
         })?;
-        if let Ok(listeners) = Arc::try_unwrap(listeners) {
-            listeners.into_iter().for_each(|listener| {
-                if let Ok(listener) = Arc::try_unwrap(listener) {
-                    listener.unbind()
-                };
-            });
-        }
         clear_signal_handlers();
         Ok(())
     }

data/gems/scheduler/ext/itsi_server/src/server/listener.rs

@@ -1,6 +1,7 @@
 use super::bind::{Bind, BindAddress};
 use super::bind_protocol::BindProtocol;
 use super::io_stream::IoStream;
+use super::serve_strategy::single_mode::RunningPhase;
 use super::tls::ItsiTlsAcceptor;
 use itsi_error::{ItsiError, Result};
 use itsi_tracing::info;
@@ -11,6 +12,7 @@ use std::{os::unix::net::UnixListener, path::PathBuf};
 use tokio::net::TcpListener as TokioTcpListener;
 use tokio::net::UnixListener as TokioUnixListener;
 use tokio::net::{unix, TcpStream, UnixStream};
+use tokio::sync::watch::Receiver;
 use tokio_rustls::TlsAcceptor;
 use tokio_stream::StreamExt;
 use tracing::error;
@@ -23,45 +25,79 @@ pub(crate) enum Listener {
 }
 
 pub(crate) enum TokioListener {
-    Tcp {
-        listener: TokioTcpListener,
-        host: String,
-        port: u16,
-    },
-    TcpTls {
-        listener: TokioTcpListener,
-        acceptor: ItsiTlsAcceptor,
-        host: String,
-        port: u16,
-    },
-    Unix {
-        listener: TokioUnixListener,
-    },
-    UnixTls {
-        listener: TokioUnixListener,
-        acceptor: ItsiTlsAcceptor,
-    },
+    Tcp(TokioTcpListener),
+    TcpTls(TokioTcpListener, ItsiTlsAcceptor),
+    Unix(TokioUnixListener),
+    UnixTls(TokioUnixListener, ItsiTlsAcceptor),
+}
+
+#[derive(Debug, Clone)]
+pub struct ListenerInfo {
+    pub host: String,
+    pub port: u16,
+    pub scheme: String,
 }
 
 impl TokioListener {
-    pub fn unbind(self) {
+    pub fn listener_info(&self) -> ListenerInfo {
         match self {
-            TokioListener::Tcp { listener, .. } => drop(listener.into_std().unwrap()),
-            TokioListener::TcpTls { listener, .. } => drop(listener.into_std().unwrap()),
-            TokioListener::Unix { listener } => drop(listener.into_std().unwrap()),
-            TokioListener::UnixTls { listener, .. } => drop(listener.into_std().unwrap()),
-        };
+            TokioListener::Tcp(listener) => ListenerInfo {
+                host: listener
+                    .local_addr()
+                    .unwrap()
+                    .ip()
+                    .to_canonical()
+                    .to_string(),
+                port: listener.local_addr().unwrap().port(),
+                scheme: "http".to_string(),
+            },
+            TokioListener::TcpTls(listener, _) => ListenerInfo {
+                host: listener
+                    .local_addr()
+                    .unwrap()
+                    .ip()
+                    .to_canonical()
+                    .to_string(),
+                port: listener.local_addr().unwrap().port(),
+                scheme: "https".to_string(),
+            },
+            TokioListener::Unix(listener) => ListenerInfo {
+                host: listener
+                    .local_addr()
+                    .unwrap()
+                    .as_pathname()
+                    .unwrap()
+                    .to_str()
+                    .unwrap()
+                    .to_owned(),
+                port: 0,
+                scheme: "unix".to_string(),
+            },
+            TokioListener::UnixTls(listener, _) => ListenerInfo {
+                host: listener
+                    .local_addr()
+                    .unwrap()
+                    .as_pathname()
+                    .unwrap()
+                    .to_str()
+                    .unwrap()
+                    .to_owned(),
+                port: 0,
+                scheme: "ssl".to_string(),
+            },
+        }
     }
+
     pub(crate) async fn accept(&self) -> Result<IoStream> {
         match self {
-            TokioListener::Tcp { listener, .. } => TokioListener::accept_tcp(listener).await,
-            TokioListener::TcpTls {
-                listener, acceptor, ..
-            } => TokioListener::accept_tls(listener, acceptor).await,
-            TokioListener::Unix { listener, .. } => TokioListener::accept_unix(listener).await,
-            TokioListener::UnixTls {
-                listener, acceptor, ..
-            } => TokioListener::accept_unix_tls(listener, acceptor).await,
+            TokioListener::Tcp(listener) => TokioListener::accept_tcp(listener).await,
+            TokioListener::TcpTls(listener, acceptor) => {
+                TokioListener::accept_tls(listener, acceptor).await
+            }
+            TokioListener::Unix(listener) => TokioListener::accept_unix(listener).await,
+            TokioListener::UnixTls(listener, acceptor) => {
+                TokioListener::accept_unix_tls(listener, acceptor).await
+            }
         }
     }
 
@@ -70,17 +106,24 @@ impl TokioListener {
         Self::to_tokio_io(Stream::TcpStream(tcp_stream), None).await
     }
 
-    pub async fn spawn_state_task(&self) {
-        if let TokioListener::TcpTls {
-            acceptor: ItsiTlsAcceptor::Automatic(_acme_acceptor, state, _server_config),
-            ..
-        } = self
+    pub async fn spawn_state_task(&self, mut shutdown_receiver: Receiver<RunningPhase>) {
+        if let TokioListener::TcpTls(
+            _,
+            ItsiTlsAcceptor::Automatic(_acme_acceptor, state, _server_config),
+        ) = self
         {
             let mut state = state.lock().await;
             loop {
-                match StreamExt::next(&mut *state).await {
-                    Some(event) => info!("Received acme event: {:?}", event),
-                    None => error!("Received no acme event"),
+                tokio::select! {
+                  stream_event = StreamExt::next(&mut *state) => {
+                      match stream_event {
+                        Some(event) => info!("Received acme event: {:?}", event),
+                        None => error!("Received no acme event"),
+                      }
+                  },
+                  _ = shutdown_receiver.changed() => {
+                      break;
+                  }
                 }
             }
         }
@@ -175,33 +218,6 @@ impl TokioListener {
             },
         }
     }
-
-    pub(crate) fn scheme(&self) -> String {
-        match self {
-            TokioListener::Tcp { .. } => "http".to_string(),
-            TokioListener::TcpTls { .. } => "https".to_string(),
-            TokioListener::Unix { .. } => "http".to_string(),
-            TokioListener::UnixTls { .. } => "https".to_string(),
-        }
-    }
-
-    pub(crate) fn port(&self) -> u16 {
-        match self {
-            TokioListener::Tcp { port, .. } => *port,
-            TokioListener::TcpTls { port, .. } => *port,
-            TokioListener::Unix { .. } => 0,
-            TokioListener::UnixTls { .. } => 0,
-        }
-    }
-
-    pub(crate) fn host(&self) -> String {
-        match self {
-            TokioListener::Tcp { host, .. } => host.to_string(),
-            TokioListener::TcpTls { host, .. } => host.to_string(),
-            TokioListener::Unix { .. } => "unix".to_string(),
-            TokioListener::UnixTls { .. } => "unix".to_string(),
-        }
-    }
 }
 
 enum Stream {
@@ -227,48 +243,22 @@ impl std::fmt::Display for SockAddr {
 }
 
 impl Listener {
-    pub fn unbind(self) {
-        match self {
-            Listener::Tcp(listener) => drop(listener),
-            Listener::TcpTls((listener, _)) => drop(listener),
-            Listener::Unix(listener) => drop(listener),
-            Listener::UnixTls((listener, _)) => drop(listener),
-        };
-    }
     pub fn to_tokio_listener(&self) -> TokioListener {
         match self {
-            Listener::Tcp(listener) => TokioListener::Tcp {
-                listener: TokioTcpListener::from_std(TcpListener::try_clone(listener).unwrap())
-                    .unwrap(),
-                host: listener
-                    .local_addr()
-                    .unwrap()
-                    .ip()
-                    .to_canonical()
-                    .to_string(),
-                port: listener.local_addr().unwrap().port(),
-            },
-            Listener::TcpTls((listener, acceptor)) => TokioListener::TcpTls {
-                listener: TokioTcpListener::from_std(TcpListener::try_clone(listener).unwrap())
-                    .unwrap(),
-                acceptor: acceptor.clone(),
-                host: listener
-                    .local_addr()
-                    .unwrap()
-                    .ip()
-                    .to_canonical()
-                    .to_string(),
-                port: listener.local_addr().unwrap().port(),
-            },
-            Listener::Unix(listener) => TokioListener::Unix {
-                listener: TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap())
-                    .unwrap(),
-            },
-            Listener::UnixTls((listener, acceptor)) => TokioListener::UnixTls {
-                listener: TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap())
-                    .unwrap(),
-                acceptor: acceptor.clone(),
-            },
+            Listener::Tcp(listener) => TokioListener::Tcp(
+                TokioTcpListener::from_std(TcpListener::try_clone(listener).unwrap()).unwrap(),
+            ),
+            Listener::TcpTls((listener, acceptor)) => TokioListener::TcpTls(
+                TokioTcpListener::from_std(TcpListener::try_clone(listener).unwrap()).unwrap(),
+                acceptor.clone(),
+            ),
+            Listener::Unix(listener) => TokioListener::Unix(
+                TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap()).unwrap(),
+            ),
+            Listener::UnixTls((listener, acceptor)) => TokioListener::UnixTls(
+                TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap()).unwrap(),
+                acceptor.clone(),
+            ),
         }
     }
 }
@@ -307,6 +297,7 @@ fn connect_tcp_socket(addr: IpAddr, port: u16) -> Result<TcpListener> {
     socket.set_nonblocking(true).ok();
     socket.set_nodelay(true).ok();
     socket.set_recv_buffer_size(262_144).ok();
+    info!("Binding to {:?}", socket_address);
     socket.bind(&socket_address.into())?;
     socket.listen(1024)?;
     Ok(socket.into())

data/gems/scheduler/ext/itsi_server/src/server/serve_strategy/single_mode.rs

@@ -4,7 +4,7 @@ use crate::{
         io_stream::IoStream,
         itsi_server::{RequestJob, Server},
         lifecycle_event::LifecycleEvent,
-        listener::{Listener, TokioListener},
+        listener::{Listener, ListenerInfo},
         thread_worker::{build_thread_workers, ThreadWorker},
     },
 };
@@ -87,18 +87,25 @@ impl SingleMode {
     pub fn run(self: Arc<Self>) -> Result<()> {
         let mut listener_task_set = JoinSet::new();
         let self_ref = Arc::new(self);
-        self_ref.build_runtime().block_on(async {
+        let runtime = self_ref.build_runtime();
 
-          for listener in self_ref.listeners.clone().iter() {
-              let listener = Arc::new(listener.to_tokio_listener());
+        runtime.block_on(async {
+          let tokio_listeners = self_ref
+              .listeners
+              .iter()
+              .map(|list| Arc::new(list.to_tokio_listener()))
+              .collect::<Vec<_>>();
+          for listener in tokio_listeners.iter() {
               let mut lifecycle_rx = self_ref.lifecycle_channel.subscribe();
+              let listener_info = Arc::new(listener.listener_info());
               let self_ref = self_ref.clone();
               let listener = listener.clone();
               let (shutdown_sender, mut shutdown_receiver) = tokio::sync::watch::channel::<RunningPhase>(RunningPhase::Running);
               let listener_clone = listener.clone();
 
-              tokio::spawn(async move {
-                listener_clone.spawn_state_task().await;
+              let shutdown_receiver_clone = shutdown_receiver.clone();
+              listener_task_set.spawn(async move {
+                listener_clone.spawn_state_task(shutdown_receiver_clone).await;
               });
 
               listener_task_set.spawn(async move {
@@ -107,8 +114,11 @@ impl SingleMode {
                     tokio::select! {
                         accept_result = listener.accept() => match accept_result {
                           Ok(accept_result) => {
-                            if let Err(e) = strategy.serve_connection(accept_result, listener.clone(), shutdown_receiver.clone()).await {
-                              error!("Error in serve_connection {:?}", e)
+                            match strategy.serve_connection(accept_result, listener_info.clone(), shutdown_receiver.clone()).await {
+                              Ok(_) => {
+                                debug!("Connection accepted and served");
+                              },
+                              Err(e) => error!("Error in serve_connection {:?}", e)
                             }
                           },
                           Err(e) => debug!("Listener.accept failed {:?}", e),
@@ -130,23 +140,23 @@ impl SingleMode {
                       }
                     }
                 }
-                if let Ok(listener) = Arc::try_unwrap(listener){
-                  listener.unbind();
-                }
             });
 
           }
 
           while let Some(_res) = listener_task_set.join_next().await {}
+
         });
+        runtime.shutdown_timeout(Duration::from_millis(100));
 
+        info!("Runtime has shut down");
         Ok(())
     }
 
     pub(crate) async fn serve_connection(
         &self,
         stream: IoStream,
-        listener: Arc<TokioListener>,
+        listener: Arc<ListenerInfo>,
         shutdown_channel: tokio::sync::watch::Receiver<RunningPhase>,
     ) -> Result<()> {
         let sender_clone = self.sender.clone();

data/gems/scheduler/ext/itsi_server/src/server/tls/locked_dir_cache.rs

@@ -1,13 +1,14 @@
 use async_trait::async_trait;
 use fs2::FileExt;
 use parking_lot::Mutex;
-use std::env;
 use std::fs::{self, OpenOptions};
 use std::io::Error as IoError;
 use std::path::{Path, PathBuf};
 use tokio_rustls_acme::caches::DirCache;
 use tokio_rustls_acme::{AccountCache, CertCache};
 
+use crate::env::ITSI_ACME_LOCK_FILE_NAME;
+
 /// A wrapper around DirCache that locks a file before writing cert/account data.
 pub struct LockedDirCache<P: AsRef<Path> + Send + Sync> {
     inner: DirCache<P>,
@@ -19,8 +20,7 @@ impl<P: AsRef<Path> + Send + Sync> LockedDirCache<P> {
     pub fn new(dir: P) -> Self {
         let dir_path = dir.as_ref().to_path_buf();
         std::fs::create_dir_all(&dir_path).unwrap();
-        let lock_path =
-            dir_path.join(env::var("ITSI_ACME_LOCK_FILE_NAME").unwrap_or(".acme.lock".to_string()));
+        let lock_path = dir_path.join(&*ITSI_ACME_LOCK_FILE_NAME);
         Self::touch_file(&lock_path).expect("Failed to create lock file");
 
         Self {