osprey 0.1.0

data/ext/osprey/src/structs/mod.rs

@@ -0,0 +1,2 @@
+pub mod rack_request;
+pub mod rack_response;

data/ext/osprey/src/structs/rack_request.rs

@@ -0,0 +1,168 @@
+use crate::{RackResponse, ARRAY_PROC, ID_CALL, ID_TRANSFORM_VALUES, OSPREY};
+use dashmap::DashMap;
+use hyper::{
+    header::{HeaderName, HeaderValue},
+    HeaderMap,
+};
+use log::{debug, error};
+use magnus::{
+    block::Proc,
+    r_string::FString,
+    value::{Opaque, OpaqueId, ReprValue},
+    Class, Error, RHash, Ruby, Value,
+};
+use std::{
+    cell::RefCell,
+    collections::HashMap,
+    os::{fd::IntoRawFd, unix::net::UnixStream},
+    sync::{Arc, LazyLock},
+};
+use tokio::sync::oneshot;
+
+pub struct RackRequestInner {
+    pub path: String,
+    pub script_name: String,
+    pub query_string: String,
+    pub method: String,
+    pub version: String,
+    pub protocol: Vec<String>,
+    pub host: String,
+    pub scheme: String,
+    pub headers: HashMap<String, String>,
+    pub remote_addr: String,
+    pub sender: Arc<std::sync::Mutex<Option<oneshot::Sender<RackResponse>>>>,
+    pub port: i32,
+    pub stream_fd: i32,
+}
+
+#[magnus::wrap(class = "RackRequest", free_immediately, size)]
+pub struct RackRequest(pub RefCell<Option<RackRequestInner>>);
+
+unsafe impl Send for RackRequest {}
+
+pub fn define_request_info(ruby: &magnus::Ruby) -> Result<(), Error> {
+    ruby.define_class("RackRequest", ruby.class_object())?;
+    Ok(())
+}
+
+static INTERNED_RUBY_STRINGS: LazyLock<DashMap<String, Opaque<FString>>> =
+    LazyLock::new(|| DashMap::new());
+
+fn intern_ruby_str(rust_str: &str, ruby: &Ruby) -> Opaque<FString> {
+    if let Some(entry) = INTERNED_RUBY_STRINGS.get(rust_str) {
+        entry.clone()
+    } else {
+        let rstr = Opaque::from(ruby.str_new(rust_str).to_interned_str());
+        INTERNED_RUBY_STRINGS.insert(rust_str.to_owned(), rstr);
+        rstr
+    }
+}
+
+impl RackRequest {
+    pub fn process(&self, ruby: &Ruby) -> Result<(), Error> {
+        let req = self.0.take().unwrap();
+        let path_info = req.path.strip_prefix(&req.script_name).unwrap_or("");
+        let sender_clone = req.sender.clone();
+        let fd = req.stream_fd;
+        let hijack_proc = ruby.proc_from_fn(move |ruby, _args, _block| {
+            let (local_fd, remote_fd) = if fd == -1 {
+                let (local, remote) = UnixStream::pair()
+                    .map_err(|_| "Failed to create ephemeral socket pair")
+                    .unwrap();
+                (local.into_raw_fd() * -1, remote.into_raw_fd())
+            } else {
+                (fd * -1, fd)
+            };
+
+            let response_sender = sender_clone
+                .clone()
+                .lock()
+                .unwrap()
+                .take()
+                .expect("No response sender");
+            if let Err(err) = response_sender.send(RackResponse::from_parts(
+                local_fd as i16,
+                HeaderMap::new(),
+                ruby.qnil().as_value(),
+            )) {
+                debug!(
+                    "(PID={}) Failed to send response: {:?}. Connection dropped.",
+                    unsafe { libc::getpid() },
+                    err
+                );
+            }
+            let io_class: magnus::RClass = magnus::eval::<magnus::RClass>("::IO").unwrap();
+
+            io_class.new_instance((remote_fd, "r+b")).unwrap()
+        });
+
+        let response = ruby.get_inner(&OSPREY).funcall::<OpaqueId, (
+            Vec<&str>,
+            Opaque<FString>,
+            Opaque<FString>,
+            Opaque<FString>,
+            Opaque<FString>,
+            Opaque<FString>,
+            Vec<String>,
+            i32,
+            i32,
+            HashMap<String, String>,
+            Proc,
+        ), (i16, RHash, Value)>(
+            *ID_CALL,
+            (
+                vec![&req.query_string, &req.remote_addr, &path_info],
+                intern_ruby_str(&req.script_name, ruby),
+                intern_ruby_str(&req.method, ruby),
+                intern_ruby_str(&req.host, ruby),
+                intern_ruby_str(&req.scheme, ruby),
+                intern_ruby_str(&req.version, ruby),
+                req.protocol,
+                req.port,
+                req.stream_fd,
+                req.headers,
+                hijack_proc,
+            ),
+        );
+
+        if let Some(response_sender) = req.sender.lock().unwrap().take() {
+            let (status, headers, body): (i16, HeaderMap, Value) = match response {
+                Ok(res) => {
+                    let (status, headers, body) = res;
+                    let header_iter = headers
+                        .funcall_with_block::<OpaqueId, (), HashMap<String, Vec<String>>>(
+                            *ID_TRANSFORM_VALUES,
+                            (),
+                            ruby.get_inner(&ARRAY_PROC),
+                        )
+                        .unwrap()
+                        .into_iter()
+                        .flat_map(|(key, values)| {
+                            values.into_iter().map(move |value| {
+                                (
+                                    HeaderName::from_bytes(key.as_bytes()).unwrap(),
+                                    HeaderValue::from_bytes(value.as_bytes()).unwrap(),
+                                )
+                            })
+                        });
+
+                    (status, HeaderMap::from_iter(header_iter), body)
+                }
+                Err(err) => {
+                    error!("Error in Ruby response: {:?}", err);
+                    (500, HeaderMap::new(), ruby.qnil().as_value())
+                }
+            };
+
+            if let Err(err) = response_sender.send(RackResponse::from_parts(status, headers, body))
+            {
+                debug!(
+                    "(PID={}) Failed to send response: {:?}. Connection dropped.",
+                    unsafe { libc::getpid() },
+                    err
+                );
+            }
+        }
+        Ok(())
+    }
+}

data/ext/osprey/src/structs/rack_response.rs

@@ -0,0 +1,152 @@
+use bytes::{Bytes, BytesMut};
+use http_body_util::{combinators::BoxBody, Full};
+use hyper::{HeaderMap, Response, StatusCode};
+use magnus::{
+    function, method, try_convert,
+    value::{LazyId, OpaqueId, ReprValue},
+    Error, IntoValue, Module, Object, Value,
+};
+use std::{
+    convert::Infallible,
+    sync::{Arc, Mutex},
+};
+
+#[derive(Debug, Clone)]
+pub struct RackResponse {
+    pub status: i16,
+    pub headers: HeaderMap,
+    pub body: Bytes,
+}
+
+static ID_EACH: LazyId = LazyId::new("each");
+static ID_CALL: LazyId = LazyId::new("call");
+static ID_CLOSE: LazyId = LazyId::new("close");
+static ID_TO_ENUM: LazyId = LazyId::new("to_enum");
+static ID_INJECT: LazyId = LazyId::new("inject");
+static ID_TO_ARY: LazyId = LazyId::new("to_ary");
+static ID_CONCAT: LazyId = LazyId::new("<<");
+
+impl RackResponse {
+    pub fn from_parts(status: i16, headers: HeaderMap, body: Value) -> Self {
+        let body = unsafe { RackResponse::gather_body(&body).unwrap_or_default() };
+        Self {
+            status,
+            headers,
+            body,
+        }
+    }
+
+    pub fn into_hyper_response(self) -> Response<BoxBody<Bytes, Infallible>> {
+        let mut response = Response::new(BoxBody::new(Full::new(self.body)));
+        *response.status_mut() = StatusCode::from_u16(self.status as u16).unwrap();
+        *response.headers_mut() = self.headers;
+        response
+    }
+
+    unsafe fn gather_body(body_value: &Value) -> Result<Bytes, Error> {
+        if body_value.respond_to(*ID_EACH, false)? {
+            let bytes: Bytes = body_value
+                .funcall::<OpaqueId, (OpaqueId,), Value>(*ID_TO_ENUM, (*ID_EACH,))
+                .unwrap()
+                .funcall::<OpaqueId, (OpaqueId,), Value>(*ID_INJECT, (*ID_CONCAT,))
+                .unwrap()
+                .to_r_string()
+                .unwrap()
+                .to_bytes();
+
+            body_value.check_funcall::<OpaqueId, (), Value>(*ID_CLOSE, ());
+            return Ok(bytes);
+        }
+
+        // 2) If body responds to `to_ary`, it can be treated like an array of strings
+        if let Some(Ok(array)) = body_value.check_funcall::<_, (), Vec<Bytes>>(*ID_TO_ARY, ()) {
+            body_value.check_funcall::<OpaqueId, (), Value>(*ID_CLOSE, ());
+            return Ok(array.into_iter().flatten().collect());
+        }
+
+        // 3) If body responds to `call`, treat it like a Streaming Body
+        if body_value.respond_to(*ID_CALL, false)? {
+            let collector = make_collector()?;
+            let _: Option<Value> = body_value.funcall(*ID_CALL, (collector,))?;
+            let data = get_collector_bytes(&collector)?;
+            body_value.check_funcall::<OpaqueId, (), Value>(*ID_CLOSE, ());
+            return Ok(data);
+        }
+
+        // 4) If none of the above apply, stringify the body and return it
+        let bytes = body_value.to_r_string().unwrap().to_bytes();
+        body_value.check_funcall::<OpaqueId, (), Value>(*ID_CLOSE, ());
+        Ok(bytes)
+    }
+}
+
+#[magnus::wrap(class = "BodyCollector")]
+struct BodyCollector {
+    buffer: Arc<Mutex<BytesMut>>,
+}
+
+impl BodyCollector {
+    fn new() -> Self {
+        BodyCollector {
+            buffer: Arc::new(Mutex::new(BytesMut::new())),
+        }
+    }
+
+    // e.g. def write(str)
+    fn write(&self, s: Bytes) -> Result<(), Error> {
+        self.buffer
+            .lock()
+            .expect("Unlock body collector buffer")
+            .extend(&s);
+        Ok(())
+    }
+
+    // def <<(str)
+    fn append(&self, s: Bytes) -> Result<Self, Error> {
+        self.buffer
+            .lock()
+            .expect("Unlock body collector buffer")
+            .extend(&s);
+        Ok(Self {
+            buffer: Arc::clone(&self.buffer),
+        })
+    }
+
+    fn flush(&self) -> Result<(), Error> {
+        Ok(())
+    }
+
+    fn close(&self) -> Result<(), Error> {
+        Ok(())
+    }
+
+    fn get_buffer(&self) -> Bytes {
+        self.buffer
+            .lock()
+            .expect("Unlock body collector buffer")
+            .clone()
+            .freeze()
+    }
+}
+
+// Then to create a Ruby object of this class:
+fn make_collector() -> Result<Value, Error> {
+    let collector = BodyCollector::new();
+    Ok(collector.into_value())
+}
+
+fn get_collector_bytes(collector_val: &Value) -> Result<Bytes, Error> {
+    let collector: &BodyCollector = try_convert::TryConvert::try_convert(*collector_val)?;
+    Ok(collector.get_buffer())
+}
+
+pub fn define_collector(ruby: &magnus::Ruby) -> Result<(), Error> {
+    let class = ruby.define_class("BodyCollector", ruby.class_object())?;
+    class.define_singleton_method("new", function!(BodyCollector::new, 0))?;
+    class.define_method("write", method!(BodyCollector::write, 1))?;
+    class.define_method("append", method!(BodyCollector::append, 1))?;
+    class.define_method("flush", method!(BodyCollector::flush, 0))?;
+    class.define_method("close", method!(BodyCollector::close, 0))?;
+
+    Ok(())
+}