itsi-server 0.1.1 → 0.1.13

data/ext/itsi_server/src/ruby_types/itsi_http_response.rs

@@ -0,0 +1,388 @@
+use bytes::{Bytes, BytesMut};
+use derive_more::Debug;
+use futures::stream::{unfold, StreamExt};
+use http::{
+    header::TRANSFER_ENCODING, request::Parts, HeaderMap, HeaderName, HeaderValue, Request,
+    Response, StatusCode,
+};
+use http_body_util::{combinators::BoxBody, Empty, Full, StreamBody};
+use hyper::{body::Frame, upgrade::Upgraded};
+use hyper_util::rt::TokioIo;
+use itsi_error::Result;
+use itsi_tracing::error;
+use magnus::error::Result as MagnusResult;
+use parking_lot::RwLock;
+use std::{
+    collections::HashMap,
+    io,
+    os::{fd::FromRawFd, unix::net::UnixStream},
+    str::FromStr,
+    sync::Arc,
+};
+use tokio::{
+    io::AsyncReadExt,
+    net::UnixStream as TokioUnixStream,
+    sync::{
+        mpsc::{self},
+        watch,
+    },
+};
+use tokio_stream::wrappers::ReceiverStream;
+use tokio_util::io::ReaderStream;
+use tracing::warn;
+
+use crate::server::{
+    byte_frame::ByteFrame, serve_strategy::single_mode::RunningPhase, types::HttpResponse,
+};
+
+#[magnus::wrap(class = "Itsi::HttpResponse", free_immediately, size)]
+#[derive(Debug, Clone)]
+pub struct ItsiHttpResponse {
+    pub data: Arc<ResponseData>,
+}
+
+#[derive(Debug)]
+pub struct ResponseData {
+    pub response: RwLock<Option<HttpResponse>>,
+    pub response_writer: RwLock<Option<mpsc::Sender<ByteFrame>>>,
+    pub response_buffer: RwLock<BytesMut>,
+    pub hijacked_socket: RwLock<Option<UnixStream>>,
+    pub parts: Parts,
+}
+
+impl ItsiHttpResponse {
+    pub async fn build(
+        &self,
+        first_frame: ByteFrame,
+        receiver: mpsc::Receiver<ByteFrame>,
+        shutdown_rx: watch::Receiver<RunningPhase>,
+    ) -> HttpResponse {
+        if self.is_hijacked() {
+            return match self.process_hijacked_response().await {
+                Ok(result) => result,
+                Err(e) => {
+                    error!("Error processing hijacked response: {}", e);
+                    Response::new(BoxBody::new(Empty::new()))
+                }
+            };
+        }
+        let mut response = self.data.response.write().take().unwrap();
+        *response.body_mut() = if matches!(first_frame, ByteFrame::Empty) {
+            BoxBody::new(Empty::new())
+        } else if matches!(first_frame, ByteFrame::End(_)) {
+            BoxBody::new(Full::new(first_frame.into()))
+        } else {
+            let initial_frame = tokio_stream::once(Ok(Frame::data(Bytes::from(first_frame))));
+            let frame_stream = unfold(
+                (ReceiverStream::new(receiver), shutdown_rx),
+                |(mut receiver, mut shutdown_rx)| async move {
+                    if let RunningPhase::ShutdownPending = *shutdown_rx.borrow() {
+                        return None;
+                    }
+                    loop {
+                        tokio::select! {
+                            maybe_bytes = receiver.next() => {
+                              match maybe_bytes {
+                                Some(ByteFrame::Data(bytes)) | Some(ByteFrame::End(bytes)) => {
+                                  return Some((Ok(Frame::data(bytes)), (receiver, shutdown_rx)));
+                                }
+                                _ => {
+                                  return None;
+                                }
+                              }
+                            },
+                            _ = shutdown_rx.changed() => {
+                                match *shutdown_rx.borrow() {
+                                    RunningPhase::ShutdownPending => {
+                                        warn!("Disconnecting streaming client.");
+                                        return None;
+                                    },
+                                    _ => continue,
+                                }
+                            }
+                        }
+                    }
+                },
+            );
+
+            let combined_stream = initial_frame.chain(frame_stream);
+            BoxBody::new(StreamBody::new(combined_stream))
+        };
+        response
+    }
+
+    pub fn close(&self) {
+        self.data.response_writer.write().take();
+    }
+
+    async fn two_way_bridge(upgraded: Upgraded, local: TokioUnixStream) -> io::Result<()> {
+        let client_io = TokioIo::new(upgraded);
+
+        // Split each side
+        let (mut lr, mut lw) = tokio::io::split(local);
+        let (mut cr, mut cw) = tokio::io::split(client_io);
+
+        let to_ruby = tokio::spawn(async move {
+            if let Err(e) = tokio::io::copy(&mut cr, &mut lw).await {
+                eprintln!("Error copying upgraded->local: {:?}", e);
+            }
+        });
+        let from_ruby = tokio::spawn(async move {
+            if let Err(e) = tokio::io::copy(&mut lr, &mut cw).await {
+                eprintln!("Error copying upgraded->local: {:?}", e);
+            }
+        });
+
+        let _ = to_ruby.await;
+        let _ = from_ruby.await;
+        Ok(())
+    }
+
+    async fn read_response_headers(&self, reader: &mut TokioUnixStream) -> Result<Vec<u8>> {
+        let mut buf = [0u8; 1];
+        let mut collected = Vec::new();
+        loop {
+            let n = reader.read(&mut buf).await?;
+            if n == 0 {
+                // EOF reached unexpectedly
+                break;
+            }
+            collected.push(buf[0]);
+            if collected.ends_with(b"\r\n\r\n") {
+                break;
+            }
+        }
+
+        Ok(collected)
+    }
+
+    pub async fn read_hijacked_headers(
+        &self,
+    ) -> Result<(HeaderMap, StatusCode, bool, TokioUnixStream)> {
+        let hijacked_socket =
+            self.data
+                .hijacked_socket
+                .write()
+                .take()
+                .ok_or(itsi_error::ItsiError::InvalidInput(
+                    "Couldn't hijack stream".to_owned(),
+                ))?;
+        let mut reader = TokioUnixStream::from_std(hijacked_socket).unwrap();
+        let response_headers = self.read_response_headers(&mut reader).await?;
+        let mut headers = [httparse::EMPTY_HEADER; 64];
+        let mut resp = httparse::Response::new(&mut headers);
+        resp.parse(&response_headers)?;
+
+        let status = StatusCode::from_u16(resp.code.unwrap_or(200)).unwrap_or(StatusCode::OK);
+        let mut headers = HeaderMap::new();
+        for header in resp.headers.iter() {
+            headers.insert(
+                HeaderName::from_str(header.name).unwrap(),
+                HeaderValue::from_bytes(header.value).unwrap(),
+            );
+        }
+        let requires_upgrade = status == StatusCode::SWITCHING_PROTOCOLS;
+        Ok((headers, status, requires_upgrade, reader))
+    }
+
+    pub async fn process_hijacked_response(&self) -> Result<HttpResponse> {
+        let (headers, status, requires_upgrade, reader) = self.read_hijacked_headers().await?;
+        let mut response = if requires_upgrade {
+            let parts = self.data.parts.clone();
+            tokio::spawn(async move {
+                let mut req = Request::from_parts(parts, Empty::<Bytes>::new());
+                match hyper::upgrade::on(&mut req).await {
+                    Ok(upgraded) => {
+                        Self::two_way_bridge(upgraded, reader)
+                            .await
+                            .expect("Error in creating two way bridge");
+                    }
+                    Err(e) => eprintln!("upgrade error: {:?}", e),
+                }
+            });
+            Response::new(BoxBody::new(Empty::new()))
+        } else {
+            let stream = ReaderStream::new(reader);
+            let boxed_body = if headers
+                .get(TRANSFER_ENCODING)
+                .is_some_and(|h| h == "chunked")
+            {
+                BoxBody::new(StreamBody::new(unfold(
+                    (stream, Vec::new()),
+                    |(mut stream, mut buf)| async move {
+                        loop {
+                            if let Some(pos) = buf.iter().position(|&b| b == b'\n') {
+                                let line = buf.drain(..=pos).collect::<Vec<u8>>();
+                                let line = std::str::from_utf8(&line).ok()?.trim();
+                                let chunk_size = usize::from_str_radix(line, 16).ok()?;
+                                if chunk_size == 0 {
+                                    return None;
+                                }
+                                while buf.len() < chunk_size {
+                                    match stream.next().await {
+                                        Some(Ok(chunk)) => buf.extend_from_slice(&chunk),
+                                        _ => return None,
+                                    }
+                                }
+                                let data = buf.drain(..chunk_size).collect::<Vec<u8>>();
+                                if buf.starts_with(b"\r\n") {
+                                    buf.drain(..2);
+                                }
+                                return Some((Ok(Frame::data(Bytes::from(data))), (stream, buf)));
+                            }
+                            match stream.next().await {
+                                Some(Ok(chunk)) => buf.extend_from_slice(&chunk),
+                                _ => return None,
+                            }
+                        }
+                    },
+                )))
+            } else {
+                BoxBody::new(StreamBody::new(stream.map(
+                    |result: std::result::Result<Bytes, io::Error>| {
+                        result
+                            .map(Frame::data)
+                            .map_err(|e| unreachable!("unexpected io error: {:?}", e))
+                    },
+                )))
+            };
+            Response::new(boxed_body)
+        };
+
+        *response.status_mut() = status;
+        *response.headers_mut() = headers;
+        Ok(response)
+    }
+
+    pub fn internal_server_error(&self, message: String) {
+        error!(message);
+        self.data.response_writer.write().take();
+        if let Some(ref mut response) = *self.data.response.write() {
+            *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
+        }
+    }
+
+    pub fn send_frame(&self, frame: Bytes) -> MagnusResult<()> {
+        self.send_frame_into(ByteFrame::Data(frame), &self.data.response_writer)
+    }
+
+    pub fn recv_frame(&self) {
+        // not implemented
+    }
+
+    pub fn send_and_close(&self, frame: Bytes) -> MagnusResult<()> {
+        let result = self.send_frame_into(ByteFrame::End(frame), &self.data.response_writer);
+        self.data.response_writer.write().take();
+        result
+    }
+
+    pub fn send_frame_into(
+        &self,
+        frame: ByteFrame,
+        writer: &RwLock<Option<mpsc::Sender<ByteFrame>>>,
+    ) -> MagnusResult<()> {
+        if let Some(writer) = writer.write().as_ref() {
+            return Ok(writer
+                .blocking_send(frame)
+                .map_err(|_| itsi_error::ItsiError::ClientConnectionClosed)?);
+        }
+        Ok(())
+    }
+
+    pub fn is_hijacked(&self) -> bool {
+        self.data.hijacked_socket.read().is_some()
+    }
+
+    pub fn close_write(&self) -> MagnusResult<bool> {
+        self.data.response_writer.write().take();
+        Ok(true)
+    }
+    fn accept_str(&self) -> &str {
+        self.data
+            .parts
+            .headers
+            .get("Content-Type")
+            .and_then(|hv| hv.to_str().ok()) // handle invalid utf-8
+            .unwrap_or("application/x-www-form-urlencoded")
+    }
+
+    pub fn is_html(&self) -> bool {
+        self.accept_str().starts_with("text/html")
+    }
+
+    pub fn is_json(&self) -> bool {
+        self.accept_str().starts_with("application/json")
+    }
+
+    pub fn close_read(&self) -> MagnusResult<bool> {
+        Ok(true)
+    }
+
+    pub fn new(parts: Parts, response_writer: mpsc::Sender<ByteFrame>) -> Self {
+        Self {
+            data: Arc::new(ResponseData {
+                response: RwLock::new(Some(Response::new(BoxBody::new(Empty::new())))),
+                response_writer: RwLock::new(Some(response_writer)),
+                response_buffer: RwLock::new(BytesMut::new()),
+                hijacked_socket: RwLock::new(None),
+                parts,
+            }),
+        }
+    }
+
+    pub fn add_header(&self, name: Bytes, value: Bytes) -> MagnusResult<()> {
+        let header_name: HeaderName = HeaderName::from_bytes(&name).map_err(|e| {
+            itsi_error::ItsiError::InvalidInput(format!("Invalid header name {:?}: {:?}", name, e))
+        })?;
+        let header_value = unsafe { HeaderValue::from_maybe_shared_unchecked(value) };
+        if let Some(ref mut resp) = *self.data.response.write() {
+            resp.headers_mut().insert(header_name, header_value);
+        }
+        Ok(())
+    }
+
+    pub fn add_headers(&self, headers: HashMap<Bytes, Vec<Bytes>>) -> MagnusResult<()> {
+        if let Some(ref mut resp) = *self.data.response.write() {
+            let headers_mut = resp.headers_mut();
+            for (name, values) in headers {
+                let header_name = HeaderName::from_bytes(&name).map_err(|e| {
+                    itsi_error::ItsiError::InvalidInput(format!(
+                        "Invalid header name {:?}: {:?}",
+                        name, e
+                    ))
+                })?;
+                for value in values {
+                    let header_value = unsafe { HeaderValue::from_maybe_shared_unchecked(value) };
+                    headers_mut.insert(&header_name, header_value);
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn set_status(&self, status: u16) -> MagnusResult<()> {
+        if let Some(ref mut resp) = *self.data.response.write() {
+            *resp.status_mut() = StatusCode::from_u16(status).map_err(|e| {
+                itsi_error::ItsiError::InvalidInput(format!(
+                    "Invalid status code {:?}: {:?}",
+                    status, e
+                ))
+            })?;
+        }
+        Ok(())
+    }
+
+    pub fn hijack(&self, fd: i32) -> MagnusResult<()> {
+        let stream = unsafe { UnixStream::from_raw_fd(fd) };
+
+        *self.data.hijacked_socket.write() = Some(stream);
+        if let Some(writer) = self.data.response_writer.write().as_ref() {
+            writer
+                .blocking_send(ByteFrame::Empty)
+                .map_err(|_| itsi_error::ItsiError::ClientConnectionClosed)?
+        }
+        self.close();
+        Ok(())
+    }
+}

data/ext/itsi_server/src/ruby_types/itsi_server/file_watcher.rs

@@ -0,0 +1,225 @@
+use derive_more::Debug;
+use globset::{Glob, GlobSet, GlobSetBuilder};
+use magnus::error::Result;
+use nix::unistd::{close, fork, pipe, read};
+use notify::{event::ModifyKind, EventKind, RecommendedWatcher};
+use notify::{Event, RecursiveMode, Watcher};
+use std::path::Path;
+use std::sync::mpsc::Sender;
+use std::time::{Duration, Instant};
+use std::{collections::HashSet, fs};
+use std::{
+    os::fd::{AsRawFd, IntoRawFd, OwnedFd},
+    path::PathBuf,
+    process::Command,
+    sync::mpsc,
+    thread::{self},
+};
+
+/// Represents a set of patterns and commands.
+#[derive(Debug, Clone)]
+struct PatternGroup {
+    base_dir: PathBuf,
+    glob_set: GlobSet,
+    pattern: String,
+    commands: Vec<Vec<String>>,
+    last_triggered: Option<Instant>,
+}
+
+/// Extracts the base directory from a wildcard pattern by taking the portion up to the first
+/// component that contains a wildcard character.
+fn extract_and_canonicalize_base_dir(pattern: &str) -> PathBuf {
+    if !(pattern.contains("*") || pattern.contains("?") || pattern.contains('[')) {
+        if let Ok(metadata) = fs::metadata(pattern) {
+            if metadata.is_dir() {
+                return fs::canonicalize(pattern).unwrap();
+            }
+            if metadata.is_file() {
+                return fs::canonicalize(pattern)
+                    .unwrap()
+                    .parent()
+                    .unwrap()
+                    .to_path_buf();
+            }
+        }
+    }
+
+    let path = Path::new(pattern);
+    let mut base = PathBuf::new();
+    for comp in path.components() {
+        let comp_str = comp.as_os_str().to_string_lossy();
+        if comp_str.contains('*') || comp_str.contains('?') || comp_str.contains('[') {
+            break;
+        } else {
+            base.push(comp);
+        }
+    }
+    // If no base was built, default to "."
+    let base = if base.as_os_str().is_empty() || !base.exists() {
+        PathBuf::from(".")
+    } else {
+        base
+    };
+
+    // Canonicalize to get the absolute path.
+    fs::canonicalize(&base).unwrap_or(base)
+}
+
+/// Minimum time between triggering the same pattern group (debounce time)
+const DEBOUNCE_DURATION: Duration = Duration::from_millis(500);
+
+pub fn watch_groups(pattern_groups: Vec<(String, Vec<Vec<String>>)>) -> Result<Option<OwnedFd>> {
+    let (r_fd, w_fd): (OwnedFd, OwnedFd) = pipe().map_err(|e| {
+        magnus::Error::new(
+            magnus::exception::exception(),
+            format!("Failed to create watcher pipe: {}", e),
+        )
+    })?;
+
+    let fork_result = unsafe {
+        fork().map_err(|e| {
+            magnus::Error::new(
+                magnus::exception::exception(),
+                format!("Failed to fork file watcher: {}", e),
+            )
+        })
+    }?;
+
+    if fork_result.is_child() {
+        let _ = close(w_fd.into_raw_fd());
+        thread::spawn(move || {
+            let mut buf = [0u8; 1];
+            loop {
+                match read(r_fd.as_raw_fd(), &mut buf) {
+                    Ok(0) => {
+                        std::process::exit(0);
+                    }
+                    Ok(_) => {}
+                    Err(_) => {
+                        std::process::exit(0);
+                    }
+                }
+            }
+        });
+
+        let mut groups = Vec::new();
+        for (pattern, commands) in pattern_groups.into_iter() {
+            let base_dir = extract_and_canonicalize_base_dir(&pattern);
+            let glob = Glob::new(&pattern).map_err(|e| {
+                magnus::Error::new(
+                    magnus::exception::exception(),
+                    format!("Failed to create watch glob: {}", e),
+                )
+            })?;
+            let glob_set = GlobSetBuilder::new().add(glob).build().map_err(|e| {
+                magnus::Error::new(
+                    magnus::exception::exception(),
+                    format!("Failed to create watch glob set: {}", e),
+                )
+            })?;
+            groups.push(PatternGroup {
+                base_dir,
+                glob_set,
+                pattern,
+                commands,
+                last_triggered: None,
+            });
+        }
+
+        // Create a channel and a watcher.
+        let (tx, rx) = mpsc::channel::<notify::Result<Event>>();
+        let sender = tx.clone();
+        fn event_fn(sender: Sender<notify::Result<Event>>) -> impl Fn(notify::Result<Event>) {
+            move |res| match res {
+                Ok(event) => {
+                    sender.send(Ok(event)).unwrap();
+                }
+                Err(e) => println!("watch error: {:?}", e),
+            }
+        }
+
+        let mut watched_dirs = HashSet::new();
+        let mut watcher: RecommendedWatcher =
+            notify::recommended_watcher(event_fn(sender)).expect("Failed to create watcher");
+        for group in &groups {
+            if watched_dirs.insert(group.base_dir.clone()) {
+                watcher
+                    .watch(&group.base_dir, RecursiveMode::Recursive)
+                    .expect("Failed to add watch");
+            }
+        }
+
+        // Main event loop.
+        for res in rx {
+            match res {
+                Ok(event) => {
+                    if !matches!(event.kind, EventKind::Modify(ModifyKind::Metadata(_))) {
+                        continue;
+                    }
+                    let now = Instant::now();
+                    for group in &mut groups {
+                        for path in event.paths.iter() {
+                            if let Ok(rel_path) = path.strip_prefix(&group.base_dir) {
+                                if group.glob_set.is_match(rel_path)
+                                    || rel_path.to_str().is_some_and(|s| s == group.pattern)
+                                {
+                                    // Check if we should debounce this event
+                                    if let Some(last_triggered) = group.last_triggered {
+                                        if now.duration_since(last_triggered) < DEBOUNCE_DURATION {
+                                            // Skip this event as we've recently triggered for this pattern
+                                            continue;
+                                        }
+                                    }
+
+                                    // Update the last triggered time
+                                    group.last_triggered = Some(now);
+
+                                    // Execute the commands for this group.
+                                    for command in &group.commands {
+                                        if command.is_empty() {
+                                            continue;
+                                        }
+                                        let mut cmd = Command::new(&command[0]);
+                                        if command.len() > 1 {
+                                            cmd.args(&command[1..]);
+                                        }
+                                        match cmd.spawn() {
+                                            Ok(mut child) => {
+                                                if let Err(e) = child.wait() {
+                                                    eprintln!(
+                                                        "Command {:?} failed: {:?}",
+                                                        command, e
+                                                    );
+                                                }
+                                            }
+                                            Err(e) => {
+                                                eprintln!(
+                                                    "Failed to execute command {:?}: {:?}",
+                                                    command, e
+                                                );
+                                            }
+                                        }
+                                    }
+                                    break;
+                                }
+                            }
+                        }
+                    }
+                }
+                Err(e) => println!("Watch error: {:?}", e),
+            }
+        }
+
+        // Clean up the watches.
+        for group in &groups {
+            watcher
+                .unwatch(&group.base_dir)
+                .expect("Failed to remove watch");
+        }
+        drop(watcher);
+        std::process::exit(0);
+    } else {
+        let _ = close(r_fd.into_raw_fd());
+        Ok(Some(w_fd))
+    }
+}