itsi-server 0.1.1 → 0.1.4

data/ext/itsi_server/src/server/thread_worker.rs

@@ -0,0 +1,368 @@
+use super::itsi_server::RequestJob;
+use crate::{request::itsi_request::ItsiRequest, ITSI_SERVER};
+use itsi_rb_helpers::{
+    call_with_gvl, call_without_gvl, create_ruby_thread, kill_threads, HeapValue,
+};
+use itsi_tracing::{debug, error, info, warn};
+use magnus::{
+    error::Result,
+    value::{InnerValue, Lazy, LazyId, Opaque, ReprValue},
+    Module, RClass, Ruby, Thread, Value,
+};
+use nix::unistd::Pid;
+use parking_lot::{Mutex, RwLock};
+use std::{
+    num::NonZeroU8,
+    ops::Deref,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    thread,
+    time::{Duration, Instant},
+};
+use tokio::{runtime::Builder as RuntimeBuilder, sync::watch};
+use tracing::instrument;
+pub struct ThreadWorker {
+    pub id: String,
+    pub app: Opaque<Value>,
+    pub receiver: Arc<async_channel::Receiver<RequestJob>>,
+    pub sender: async_channel::Sender<RequestJob>,
+    pub thread: RwLock<Option<HeapValue<Thread>>>,
+    pub terminated: Arc<AtomicBool>,
+    pub scheduler_class: Option<Opaque<Value>>,
+}
+
+static ID_CALL: LazyId = LazyId::new("call");
+static ID_ALIVE: LazyId = LazyId::new("alive?");
+static ID_SCHEDULER: LazyId = LazyId::new("scheduler");
+static ID_SCHEDULE: LazyId = LazyId::new("schedule");
+static ID_BLOCK: LazyId = LazyId::new("block");
+static ID_YIELD: LazyId = LazyId::new("yield");
+static ID_CONST_GET: LazyId = LazyId::new("const_get");
+static CLASS_FIBER: Lazy<RClass> = Lazy::new(|ruby| {
+    ruby.module_kernel()
+        .const_get::<_, RClass>("Fiber")
+        .unwrap()
+});
+
+pub struct TerminateWakerSignal(bool);
+
+#[instrument(name = "Boot", parent=None, skip(threads, app, pid, scheduler_class))]
+pub fn build_thread_workers(
+    pid: Pid,
+    threads: NonZeroU8,
+    app: Opaque<Value>,
+    scheduler_class: Option<String>,
+) -> Result<(Arc<Vec<ThreadWorker>>, async_channel::Sender<RequestJob>)> {
+    let (sender, receiver) = async_channel::bounded(20);
+    let receiver_ref = Arc::new(receiver);
+    let sender_ref = sender;
+    let (app, scheduler_class) = load_app(app, scheduler_class)?;
+    Ok((
+        Arc::new(
+            (1..=u8::from(threads))
+                .map(|id| {
+                    info!(pid = pid.as_raw(), id, "Thread");
+                    ThreadWorker::new(
+                        format!("{:?}#{:?}", pid, id),
+                        app,
+                        receiver_ref.clone(),
+                        sender_ref.clone(),
+                        scheduler_class,
+                    )
+                })
+                .collect::<Result<Vec<_>>>()?,
+        ),
+        sender_ref,
+    ))
+}
+
+pub fn load_app(
+    app: Opaque<Value>,
+    scheduler_class: Option<String>,
+) -> Result<(Opaque<Value>, Option<Opaque<Value>>)> {
+    call_with_gvl(|ruby| {
+        let app = app.get_inner_with(&ruby);
+        let app = Opaque::from(
+            app.funcall::<_, _, Value>(*ID_CALL, ())
+                .expect("Couldn't load app"),
+        );
+        let scheduler_class = if let Some(scheduler_class) = scheduler_class {
+            Some(Opaque::from(
+                ruby.module_kernel()
+                    .funcall::<_, _, Value>(*ID_CONST_GET, (scheduler_class,))?,
+            ))
+        } else {
+            None
+        };
+        Ok((app, scheduler_class))
+    })
+}
+impl ThreadWorker {
+    pub fn new(
+        id: String,
+        app: Opaque<Value>,
+        receiver: Arc<async_channel::Receiver<RequestJob>>,
+        sender: async_channel::Sender<RequestJob>,
+        scheduler_class: Option<Opaque<Value>>,
+    ) -> Result<Self> {
+        let mut worker = Self {
+            id,
+            app,
+            receiver,
+            sender,
+            thread: RwLock::new(None),
+            terminated: Arc::new(AtomicBool::new(false)),
+            scheduler_class,
+        };
+        worker.run()?;
+        Ok(worker)
+    }
+
+    #[instrument(skip(self), fields(id = self.id))]
+    pub async fn request_shutdown(&self) {
+        match self.sender.send(RequestJob::Shutdown).await {
+            Ok(_) => {}
+            Err(err) => error!("Failed to send shutdown request: {}", err),
+        };
+        info!("Requesting shutdown");
+    }
+
+    #[instrument(skip(self, deadline), fields(id = self.id))]
+    pub fn poll_shutdown(&self, deadline: Instant) -> bool {
+        call_with_gvl(|_ruby| {
+            if let Some(thread) = self.thread.read().deref() {
+                if Instant::now() > deadline {
+                    warn!("Worker shutdown timed out. Killing thread");
+                    self.terminated.store(true, Ordering::SeqCst);
+                    kill_threads(vec![thread.as_value()]);
+                }
+                if thread.funcall::<_, _, bool>(*ID_ALIVE, ()).unwrap_or(false) {
+                    return true;
+                }
+                info!("Thread has shut down");
+            }
+            self.thread.write().take();
+
+            false
+        })
+    }
+
+    pub fn run(&mut self) -> Result<()> {
+        let id = self.id.clone();
+        let app = self.app;
+        let receiver = self.receiver.clone();
+        let terminated = self.terminated.clone();
+        let scheduler_class = self.scheduler_class;
+        call_with_gvl(|_| {
+            *self.thread.write() = Some(
+                create_ruby_thread(move || {
+                    if let Some(scheduler_class) = scheduler_class {
+                        if let Err(err) =
+                            Self::fiber_accept_loop(id, app, receiver, scheduler_class, terminated)
+                        {
+                            error!("Error in fiber_accept_loop: {:?}", err);
+                        }
+                    } else {
+                        Self::accept_loop(id, app, receiver, terminated);
+                    }
+                })
+                .into(),
+            );
+            Ok::<(), magnus::Error>(())
+        })?;
+        Ok(())
+    }
+
+    pub fn build_scheduler_proc(
+        app: Opaque<Value>,
+        leader: &Arc<Mutex<Option<RequestJob>>>,
+        receiver: &Arc<async_channel::Receiver<RequestJob>>,
+        terminated: &Arc<AtomicBool>,
+        waker_sender: &watch::Sender<TerminateWakerSignal>,
+    ) -> magnus::block::Proc {
+        let leader = leader.clone();
+        let receiver = receiver.clone();
+        let terminated = terminated.clone();
+        let waker_sender = waker_sender.clone();
+        Ruby::get().unwrap().proc_from_fn(move |ruby, _args, _blk| {
+            let scheduler = ruby
+                .get_inner(&CLASS_FIBER)
+                .funcall::<_, _, Value>(*ID_SCHEDULER, ())
+                .unwrap();
+            let server = ruby.get_inner(&ITSI_SERVER);
+            let thread_current = ruby.thread_current();
+            let leader_clone = leader.clone();
+            let receiver = receiver.clone();
+            let terminated = terminated.clone();
+            let waker_sender = waker_sender.clone();
+            let mut batch = Vec::with_capacity(MAX_BATCH_SIZE as usize);
+
+            static MAX_BATCH_SIZE: i32 = 25;
+            call_without_gvl(move || loop {
+                let mut idle_counter = 0;
+                if let Some(v) = leader_clone.lock().take() {
+                    match v {
+                        RequestJob::ProcessRequest(itsi_request) => {
+                            batch.push(RequestJob::ProcessRequest(itsi_request))
+                        }
+                        RequestJob::Shutdown => {
+                            waker_sender.send(TerminateWakerSignal(true)).unwrap();
+                            break;
+                        }
+                    }
+                }
+                for _ in 0..MAX_BATCH_SIZE {
+                    if let Ok(req) = receiver.try_recv() {
+                        batch.push(req);
+                    } else {
+                        break;
+                    }
+                }
+
+                let shutdown_requested = call_with_gvl(|_| {
+                    for req in batch.drain(..) {
+                        match req {
+                            RequestJob::ProcessRequest(request) => {
+                                let response = request.response.clone();
+                                if let Err(err) =
+                                    server.funcall::<_, _, Value>(*ID_SCHEDULE, (app, request))
+                                {
+                                    ItsiRequest::internal_error(ruby, response, err)
+                                }
+                            }
+                            RequestJob::Shutdown => return true,
+                        }
+                    }
+                    false
+                });
+
+                if shutdown_requested || terminated.load(Ordering::Relaxed) {
+                    waker_sender.send(TerminateWakerSignal(true)).unwrap();
+                    break;
+                }
+
+                let yield_result = if receiver.is_empty() {
+                    waker_sender.send(TerminateWakerSignal(false)).unwrap();
+                    idle_counter = (idle_counter + 1) % 100;
+                    call_with_gvl(|ruby| {
+                        if idle_counter == 0 {
+                            ruby.gc_start();
+                        }
+                        scheduler.funcall::<_, _, Value>(*ID_BLOCK, (thread_current, None::<u8>))
+                    })
+                } else {
+                    call_with_gvl(|_| scheduler.funcall::<_, _, Value>(*ID_YIELD, ()))
+                };
+
+                if yield_result.is_err() {
+                    break;
+                }
+            })
+        })
+    }
+
+    #[instrument(skip_all, fields(thread_worker=id))]
+    pub fn fiber_accept_loop(
+        id: String,
+        app: Opaque<Value>,
+        receiver: Arc<async_channel::Receiver<RequestJob>>,
+        scheduler_class: Opaque<Value>,
+        terminated: Arc<AtomicBool>,
+    ) -> Result<()> {
+        let ruby = Ruby::get().unwrap();
+        let (waker_sender, waker_receiver) = watch::channel(TerminateWakerSignal(false));
+        let leader: Arc<Mutex<Option<RequestJob>>> = Arc::new(Mutex::new(None));
+        let server = ruby.get_inner(&ITSI_SERVER);
+        let scheduler_proc =
+            Self::build_scheduler_proc(app, &leader, &receiver, &terminated, &waker_sender);
+        let (scheduler, scheduler_fiber) = server.funcall::<_, _, (Value, Value)>(
+            "start_scheduler_loop",
+            (scheduler_class, scheduler_proc),
+        )?;
+        Self::start_waker_thread(
+            scheduler.into(),
+            scheduler_fiber.into(),
+            leader,
+            receiver,
+            waker_receiver,
+        );
+        Ok(())
+    }
+
+    #[allow(clippy::await_holding_lock)]
+    pub fn start_waker_thread(
+        scheduler: Opaque<Value>,
+        scheduler_fiber: Opaque<Value>,
+        leader: Arc<Mutex<Option<RequestJob>>>,
+        receiver: Arc<async_channel::Receiver<RequestJob>>,
+        mut waker_receiver: watch::Receiver<TerminateWakerSignal>,
+    ) {
+        create_ruby_thread(move || {
+            let scheduler = scheduler.get_inner_with(&Ruby::get().unwrap());
+            let leader = leader.clone();
+            call_without_gvl(|| {
+                RuntimeBuilder::new_current_thread()
+                    .build()
+                    .expect("Failed to build Tokio runtime")
+                    .block_on(async {
+                        loop {
+                            waker_receiver.changed().await.ok();
+                            if waker_receiver.borrow().0 {
+                                break;
+                            }
+                            tokio::select! {
+                                _ = waker_receiver.changed() => {
+                                  if waker_receiver.borrow().0 {
+                                      break;
+                                  }
+                                },
+                                next_msg = receiver.recv() => {
+                                  *leader.lock() = next_msg.ok();
+                                  call_with_gvl(|_| {
+                                      scheduler
+                                          .funcall::<_, _, Value>(
+                                              "unblock",
+                                              (None::<u8>, scheduler_fiber),
+                                          )
+                                          .ok();
+                                  });
+                                }
+                            }
+                        }
+                    })
+            });
+        });
+    }
+
+    #[instrument(skip_all, fields(thread_worker=id))]
+    pub fn accept_loop(
+        id: String,
+        app: Opaque<Value>,
+        receiver: Arc<async_channel::Receiver<RequestJob>>,
+        terminated: Arc<AtomicBool>,
+    ) {
+        let ruby = Ruby::get().unwrap();
+        let server = ruby.get_inner(&ITSI_SERVER);
+        call_without_gvl(|| loop {
+            match receiver.recv_blocking() {
+                Ok(RequestJob::ProcessRequest(request)) => {
+                    if terminated.load(Ordering::Relaxed) {
+                        break;
+                    }
+                    call_with_gvl(|_ruby| {
+                        request.process(&ruby, server, app).ok();
+                    })
+                }
+                Ok(RequestJob::Shutdown) => {
+                    debug!("Shutting down thread worker");
+                    break;
+                }
+                Err(_) => {
+                    thread::sleep(Duration::from_micros(1));
+                }
+            }
+        });
+    }
+}

data/ext/itsi_server/src/server/tls.rs

@@ -1,20 +1,69 @@
 use base64::{engine::general_purpose, Engine as _};
 use itsi_error::Result;
-use itsi_tracing::{info, warn};
+use itsi_tracing::info;
 use rcgen::{CertificateParams, DnType, KeyPair, SanType};
+use rustls::pki_types::{CertificateDer, PrivateKeyDer};
 use rustls_pemfile::{certs, pkcs8_private_keys};
-use std::{collections::HashMap, fs, io::BufReader};
-use tokio_rustls::rustls::{Certificate, PrivateKey, ServerConfig};
+use std::{
+    collections::HashMap,
+    env, fs,
+    io::{BufReader, Error},
+    sync::Arc,
+};
+use tokio::sync::Mutex;
+use tokio_rustls::{rustls::ServerConfig, TlsAcceptor};
+use tokio_rustls_acme::{caches::DirCache, AcmeAcceptor, AcmeConfig, AcmeState};
 
 const ITS_CA_CERT: &str = include_str!("./itsi_ca/itsi_ca.crt");
 const ITS_CA_KEY: &str = include_str!("./itsi_ca/itsi_ca.key");
 
+#[derive(Clone)]
+pub enum ItsiTlsAcceptor {
+    Manual(TlsAcceptor),
+    Automatic(
+        AcmeAcceptor,
+        Arc<Mutex<AcmeState<Error>>>,
+        Arc<ServerConfig>,
+    ),
+}
+
 // Generates a TLS configuration based on either :
 // * Input "cert" and "key" options (either paths or Base64-encoded strings) or
 // * Performs automatic certificate generation/retrieval. Generated certs use an internal self-signed Isti CA.
 // If a non-local host or optional domain parameter is provided,
 // an automated certificate will attempt to be fetched using let's encrypt.
-pub fn configure_tls(host: &str, query_params: &HashMap<String, String>) -> Result<ServerConfig> {
+pub fn configure_tls(
+    host: &str,
+    query_params: &HashMap<String, String>,
+) -> Result<ItsiTlsAcceptor> {
+    let domains = query_params
+        .get("domains")
+        .map(|v| v.split(',').map(String::from).collect::<Vec<_>>());
+
+    if query_params.get("cert").is_none() || query_params.get("key").is_none() {
+        if let Some(domains) = domains {
+            let directory_url = env::var("ACME_DIRECTORY_URL")
+                .unwrap_or_else(|_| "https://acme-v02.api.letsencrypt.org/directory".to_string());
+            info!(
+                domains = format!("{:?}", domains),
+                directory_url, "Requesting acme cert"
+            );
+            let acme_state = AcmeConfig::new(domains)
+                .contact(["mailto:wc@pico.net.nz"])
+                .cache(DirCache::new("./rustls_acme_cache"))
+                .directory(directory_url)
+                .state();
+            let rustls_config = ServerConfig::builder()
+                .with_no_client_auth()
+                .with_cert_resolver(acme_state.resolver());
+            let acceptor = acme_state.acceptor();
+            return Ok(ItsiTlsAcceptor::Automatic(
+                acceptor,
+                Arc::new(Mutex::new(acme_state)),
+                Arc::new(rustls_config),
+            ));
+        }
+    }
     let (certs, key) = if let (Some(cert_path), Some(key_path)) =
         (query_params.get("cert"), query_params.get("key"))
     {
@@ -22,41 +71,20 @@ pub fn configure_tls(host: &str, query_params: &HashMap<String, String>) -> Resu
         let certs = load_certs(cert_path);
         let key = load_private_key(key_path);
         (certs, key)
-    } else if query_params
-        .get("cert")
-        .map(|v| v == "auto")
-        .unwrap_or(false)
-    {
-        let domain_param = query_params.get("domain");
-        let host_string = host.to_string();
-        let domain = domain_param.or_else(|| {
-            if host_string != "localhost" {
-                Some(&host_string)
-            } else {
-                None
-            }
-        });
-
-        if let Some(domain) = domain {
-            retrieve_acme_cert(domain)?
-        } else {
-            generate_ca_signed_cert(host)?
-        }
     } else {
-        generate_ca_signed_cert(host)?
+        generate_ca_signed_cert(vec![host.to_owned()])?
     };
 
     let mut config = ServerConfig::builder()
-        .with_safe_defaults()
         .with_no_client_auth()
         .with_single_cert(certs, key)
         .expect("Failed to build TLS config");
 
     config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
-    Ok(config)
+    Ok(ItsiTlsAcceptor::Manual(TlsAcceptor::from(Arc::new(config))))
 }
 
-pub fn load_certs(path: &str) -> Vec<Certificate> {
+pub fn load_certs(path: &str) -> Vec<CertificateDer<'static>> {
     let data = if let Some(stripped) = path.strip_prefix("base64:") {
         general_purpose::STANDARD
             .decode(stripped)
@@ -74,14 +102,20 @@ pub fn load_certs(path: &str) -> Vec<Certificate> {
             })
             .collect::<Result<_>>()
             .expect("Failed to parse certificate file");
-        certs_der.into_iter().map(Certificate).collect()
+        certs_der
+            .into_iter()
+            .map(|vec| {
+                // Convert the owned Vec<u8> into a CertificateDer and force 'static.
+                unsafe { std::mem::transmute(CertificateDer::from(vec)) }
+            })
+            .collect()
     } else {
-        vec![Certificate(data)]
+        vec![CertificateDer::from(data)]
     }
 }
 
 /// Loads a private key from a file or Base64.
-pub fn load_private_key(path: &str) -> PrivateKey {
+pub fn load_private_key(path: &str) -> PrivateKeyDer<'static> {
     let key_data = if let Some(stripped) = path.strip_prefix("base64:") {
         general_purpose::STANDARD
             .decode(stripped)
@@ -100,39 +134,55 @@ pub fn load_private_key(path: &str) -> PrivateKey {
             .collect::<Result<_>>()
             .expect("Failed to parse private key");
         if !keys.is_empty() {
-            return PrivateKey(keys[0].clone());
+            return PrivateKeyDer::try_from(keys[0].clone()).unwrap();
         }
     }
-    PrivateKey(key_data)
+    PrivateKeyDer::try_from(key_data).unwrap()
 }
 
-pub fn generate_ca_signed_cert(domain: &str) -> Result<(Vec<Certificate>, PrivateKey)> {
+pub fn generate_ca_signed_cert(
+    domains: Vec<String>,
+) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>)> {
     info!("Generating New Itsi CA - Self signed Certificate. Use `itsi ca export` to export the CA certificate for import into your local trust store.");
 
-    let ca_kp = KeyPair::from_pem(ITS_CA_KEY).unwrap();
-    let params = CertificateParams::from_ca_cert_pem(ITS_CA_CERT).unwrap();
+    let ca_kp = KeyPair::from_pem(ITS_CA_KEY).expect("Failed to load embedded CA key");
+    let ca_cert = CertificateParams::from_ca_cert_pem(ITS_CA_CERT)
+        .expect("Failed to parse embedded CA certificate")
+        .self_signed(&ca_kp)
+        .expect("Failed to self-sign embedded CA cert");
 
-    let ca_cert = params.self_signed(&ca_kp).unwrap();
-    let ee_key = KeyPair::generate().unwrap();
+    let ee_key = KeyPair::generate_for(&rcgen::PKCS_ECDSA_P256_SHA256).unwrap();
     let mut ee_params = CertificateParams::default();
 
-    // Set the domain in the subject alternative names (SAN)
-    ee_params.subject_alt_names = vec![SanType::DnsName(domain.try_into()?)];
-    // Optionally, set the Common Name (CN) in the distinguished name:
+    info!(
+        "Generated certificate will be valid for domains {:?}",
+        domains
+    );
+    use std::net::IpAddr;
+
+    ee_params.subject_alt_names = domains
+        .iter()
+        .map(|domain| {
+            if let Ok(ip) = domain.parse::<IpAddr>() {
+                SanType::IpAddress(ip)
+            } else {
+                SanType::DnsName(domain.clone().try_into().unwrap())
+            }
+        })
+        .collect();
+
     ee_params
         .distinguished_name
-        .push(DnType::CommonName, domain);
+        .push(DnType::CommonName, domains[0].clone());
 
     ee_params.use_authority_key_identifier_extension = true;
 
-    let ee_cert = ee_params.signed_by(&ee_key, &ca_cert, &ee_key).unwrap();
+    let ee_cert = ee_params.signed_by(&ee_key, &ca_cert, &ca_kp).unwrap();
     let ee_cert_der = ee_cert.der().to_vec();
-    let ee_cert = Certificate(ee_cert_der);
-    Ok((vec![ee_cert], PrivateKey(ee_key.serialize_der())))
-}
-
-/// Retrieves an ACME certificate for a given domain.
-pub fn retrieve_acme_cert(domain: &str) -> Result<(Vec<Certificate>, PrivateKey)> {
-    warn!("Retrieving ACME cert for {}", domain);
-    generate_ca_signed_cert(domain)
+    let ee_cert = CertificateDer::from(ee_cert_der);
+    let ca_cert = CertificateDer::from(ca_cert.der().to_vec());
+    Ok((
+        vec![ee_cert, ca_cert],
+        PrivateKeyDer::try_from(ee_key.serialize_der()).unwrap(),
+    ))
 }

data/ext/itsi_tracing/Cargo.toml

@@ -9,4 +9,8 @@ tracing-subscriber = { version = "0.3.19", features = [
   "env-filter",
   "std",
   "fmt",
+  "json",
+  "ansi",
 ] }
+tracing-attributes = "0.1"
+atty = "0.2.14"

data/ext/itsi_tracing/src/lib.rs

@@ -1,11 +1,41 @@
+use std::env;
+
+use atty::{Stream, is};
 pub use tracing::{debug, error, info, trace, warn};
-use tracing_subscriber::{EnvFilter, fmt};
+pub use tracing_attributes::instrument; // Explicitly export from tracing-attributes
+use tracing_subscriber::{
+    EnvFilter,
+    fmt::{self, format},
+};
 
+#[instrument]
 pub fn init() {
-    let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
-    let format = fmt::format().with_level(true).with_target(false).compact();
-    tracing_subscriber::fmt()
-        .with_env_filter(env_filter)
+    let env_filter = EnvFilter::builder()
+        .with_env_var("ITSI_LOG")
+        .try_from_env()
+        .unwrap_or_else(|_| EnvFilter::new("info"));
+
+    let format = fmt::format()
+        .compact()
+        .with_file(false)
+        .with_level(true)
+        .with_line_number(false)
+        .with_source_location(false)
+        .with_target(false)
+        .with_thread_ids(false);
+
+    let is_tty = is(Stream::Stdout);
+
+    let subscriber = tracing_subscriber::fmt()
         .event_format(format)
-        .init();
+        .with_env_filter(env_filter);
+
+    if (is_tty && env::var("ITSI_LOG_PLAIN").is_err()) || env::var("ITSI_LOG_ANSI").is_ok() {
+        subscriber.with_ansi(true).init();
+    } else {
+        subscriber
+            .fmt_fields(format::JsonFields::default())
+            .event_format(fmt::format().json())
+            .init();
+    }
 }