tonic-rails 0.1.0

data/ext/tonic_rails_native/src/body.rs

@@ -0,0 +1,289 @@
+//! gRPC response bodies and request accounting.
+//!
+//! gRPC carries `grpc-status` in trailing HEADERS, so responses use
+//! custom `http_body::Body` implementations that yield data frames
+//! followed by a trailers frame.
+
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::task::{Context, Poll};
+use std::time::Instant;
+
+use bytes::Bytes;
+use http::HeaderMap;
+use http_body::{Body, Frame};
+use tokio::sync::mpsc;
+use tonic::{Code, Status};
+
+use crate::framing::{COMPRESSION_THRESHOLD, frame, sanitize_for_log, status_trailers};
+use crate::gvl::StreamMessage;
+
+/// Request counters, exposed to Ruby via `Server#stats`.
+#[derive(Default)]
+pub struct ServerStats {
+    pub total: AtomicUsize,
+    pub successful: AtomicUsize,
+    pub errors: AtomicUsize,
+    pub cancelled: AtomicUsize,
+    pub deadline_exceeded: AtomicUsize,
+    pub resource_exhausted: AtomicUsize,
+    pub unimplemented: AtomicUsize,
+    /// Log every request (including OK) to stderr.  Errors are always logged.
+    pub log_requests: AtomicBool,
+}
+
+impl ServerStats {
+    pub fn record(&self, path: &str, code: u32, message: &str, start: Instant) {
+        let counter = match Code::from_i32(code as i32) {
+            Code::Ok => &self.successful,
+            Code::Cancelled => &self.cancelled,
+            _ => &self.errors,
+        };
+        counter.fetch_add(1, Ordering::Relaxed);
+        match Code::from_i32(code as i32) {
+            Code::DeadlineExceeded => {
+                self.deadline_exceeded.fetch_add(1, Ordering::Relaxed);
+            }
+            Code::ResourceExhausted => {
+                self.resource_exhausted.fetch_add(1, Ordering::Relaxed);
+            }
+            Code::Unimplemented => {
+                self.unimplemented.fetch_add(1, Ordering::Relaxed);
+            }
+            _ => {}
+        }
+
+        let is_error = !matches!(Code::from_i32(code as i32), Code::Ok | Code::Cancelled);
+        if is_error || self.log_requests.load(Ordering::Relaxed) {
+            let path = sanitize_for_log(path);
+            let dur_us = start.elapsed().as_micros();
+            if message.is_empty() {
+                eprintln!("tonic-rails: path={path} code={code} duration_us={dur_us}");
+            } else {
+                let message = sanitize_for_log(message);
+                eprintln!("tonic-rails: path={path} code={code} duration_us={dur_us} message={message}");
+            }
+        }
+    }
+}
+
+/// Trailers-only error response for requests rejected before dispatch.
+pub struct ErrorBody {
+    trailers: Option<HeaderMap>,
+}
+
+impl ErrorBody {
+    pub fn new(code: Code, message: &str) -> Self {
+        Self {
+            trailers: Some(status_trailers(code as u32, message)),
+        }
+    }
+}
+
+impl Body for ErrorBody {
+    type Data = Bytes;
+    type Error = Status;
+
+    fn poll_frame(
+        mut self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
+        Poll::Ready(self.trailers.take().map(|t| Ok(Frame::trailers(t))))
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.trailers.is_none()
+    }
+}
+
+/// Streaming response body fed by the Ruby handler through an mpsc
+/// channel.  Each `Data` message becomes a gRPC frame; `End` (or the
+/// deadline) produces the trailers.
+pub struct ResponseBody {
+    pub receiver: Option<mpsc::Receiver<StreamMessage>>,
+    pub stats: Arc<ServerStats>,
+    pub path: String,
+    pub start: Instant,
+    pub compress: bool,
+    pub deadline: Option<Pin<Box<tokio::time::Sleep>>>,
+    pub max_message_size: usize,
+    pub finished: bool,
+}
+
+impl ResponseBody {
+    fn finish(&mut self, code: Code, message: &str) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
+        self.receiver = None;
+        self.finished = true;
+        self.stats.record(&self.path, code as u32, message, self.start);
+        Poll::Ready(Some(Ok(Frame::trailers(status_trailers(code as u32, message)))))
+    }
+}
+
+impl Body for ResponseBody {
+    type Data = Bytes;
+    type Error = Status;
+
+    fn poll_frame(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
+        let Some(receiver) = self.receiver.as_mut() else {
+            return Poll::Ready(None);
+        };
+
+        // Drain ready messages before checking the deadline so a response
+        // that made it in time is never discarded.
+        match receiver.poll_recv(cx) {
+            Poll::Ready(Some(StreamMessage::Data(bytes))) => {
+                if bytes.len() > self.max_message_size {
+                    return self.finish(
+                        Code::ResourceExhausted,
+                        "response message exceeds max_message_size",
+                    );
+                }
+                let compress = self.compress && bytes.len() > COMPRESSION_THRESHOLD;
+                return Poll::Ready(Some(Ok(Frame::data(frame(&bytes, compress)))));
+            }
+            Poll::Ready(Some(StreamMessage::End(None))) => return self.finish(Code::Ok, ""),
+            Poll::Ready(Some(StreamMessage::End(Some((code, message))))) => {
+                return self.finish(Code::from_i32(code as i32), &message);
+            }
+            Poll::Ready(None) => return self.finish(Code::Internal, "handler ended without a response"),
+            Poll::Pending => {}
+        }
+
+        if let Some(deadline) = self.deadline.as_mut()
+            && deadline.as_mut().poll(cx).is_ready()
+        {
+            return self.finish(Code::DeadlineExceeded, "deadline exceeded");
+        }
+        Poll::Pending
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.receiver.is_none()
+    }
+}
+
+impl Drop for ResponseBody {
+    fn drop(&mut self) {
+        // Dropped before trailers were produced: the client went away
+        // (RST_STREAM / disconnect) — count it as cancelled.
+        if !self.finished {
+            self.stats.cancelled.fetch_add(1, Ordering::Relaxed);
+            if self.stats.log_requests.load(Ordering::Relaxed) {
+                eprintln!(
+                    "tonic-rails: path={} code=1 message=cancelled by client",
+                    sanitize_for_log(&self.path)
+                );
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::task::Waker;
+
+    type PolledFrame = Poll<Option<Result<Frame<Bytes>, Status>>>;
+
+    fn poll<B: Body<Data = Bytes, Error = Status> + Unpin>(body: &mut B) -> PolledFrame {
+        let mut cx = Context::from_waker(Waker::noop());
+        Pin::new(body).poll_frame(&mut cx)
+    }
+
+    fn response_body(rx: mpsc::Receiver<StreamMessage>, stats: Arc<ServerStats>) -> ResponseBody {
+        ResponseBody {
+            receiver: Some(rx),
+            stats,
+            path: "/test.Svc/Method".into(),
+            start: Instant::now(),
+            compress: false,
+            deadline: None,
+            max_message_size: 16,
+            finished: false,
+        }
+    }
+
+    fn trailers(p: PolledFrame) -> HeaderMap {
+        match p {
+            Poll::Ready(Some(Ok(frame))) => frame.into_trailers().expect("trailers frame"),
+            _ => panic!("expected trailers"),
+        }
+    }
+
+    #[test]
+    fn error_body_yields_trailers_once() {
+        let mut body = ErrorBody::new(Code::NotFound, "missing");
+        let t = trailers(poll(&mut body));
+        assert_eq!(t["grpc-status"], "5");
+        assert_eq!(t["grpc-message"], "missing");
+        assert!(matches!(poll(&mut body), Poll::Ready(None)));
+        assert!(body.is_end_stream());
+    }
+
+    #[test]
+    fn response_body_streams_data_then_ok() {
+        let (tx, rx) = mpsc::channel(4);
+        let stats = Arc::new(ServerStats::default());
+        let mut body = response_body(rx, stats.clone());
+        tx.try_send(StreamMessage::Data(b"abc".to_vec())).ok();
+        tx.try_send(StreamMessage::End(None)).ok();
+
+        match poll(&mut body) {
+            Poll::Ready(Some(Ok(frame))) => {
+                assert_eq!(
+                    frame.into_data().unwrap().as_ref(),
+                    &[0, 0, 0, 0, 3, b'a', b'b', b'c']
+                );
+            }
+            _ => panic!("expected data"),
+        }
+        assert_eq!(trailers(poll(&mut body))["grpc-status"], "0");
+        assert!(matches!(poll(&mut body), Poll::Ready(None)));
+        drop(body);
+        assert_eq!(stats.successful.load(Ordering::Relaxed), 1);
+        assert_eq!(stats.cancelled.load(Ordering::Relaxed), 0);
+    }
+
+    #[test]
+    fn response_body_rejects_oversized_message() {
+        let (tx, rx) = mpsc::channel(4);
+        let stats = Arc::new(ServerStats::default());
+        let mut body = response_body(rx, stats.clone());
+        tx.try_send(StreamMessage::Data(vec![0; 17])).ok();
+        assert_eq!(trailers(poll(&mut body))["grpc-status"], "8");
+        assert_eq!(stats.resource_exhausted.load(Ordering::Relaxed), 1);
+    }
+
+    #[test]
+    fn response_body_reports_handler_error() {
+        let (tx, rx) = mpsc::channel(4);
+        let mut body = response_body(rx, Arc::new(ServerStats::default()));
+        tx.try_send(StreamMessage::End(Some((7, "denied".into())))).ok();
+        let t = trailers(poll(&mut body));
+        assert_eq!(t["grpc-status"], "7");
+        assert_eq!(t["grpc-message"], "denied");
+    }
+
+    #[test]
+    fn response_body_closed_channel_is_internal() {
+        let (tx, rx) = mpsc::channel(4);
+        let mut body = response_body(rx, Arc::new(ServerStats::default()));
+        drop(tx);
+        assert_eq!(trailers(poll(&mut body))["grpc-status"], "13");
+    }
+
+    #[test]
+    fn dropped_body_counts_as_cancelled() {
+        let (_tx, rx) = mpsc::channel(4);
+        let stats = Arc::new(ServerStats::default());
+        let mut body = response_body(rx, stats.clone());
+        assert!(poll(&mut body).is_pending());
+        drop(body);
+        assert_eq!(stats.cancelled.load(Ordering::Relaxed), 1);
+    }
+}

data/ext/tonic_rails_native/src/dispatch.rs

@@ -0,0 +1,250 @@
+//! Raw-byte gRPC dispatch service.
+//!
+//! Every request that is not for the built-in health or reflection
+//! services lands here.  Requests are validated, the handler is queued
+//! for a Ruby dispatch thread *before* the body is read (so bidi RPCs can
+//! interleave), and a spawned task feeds request messages to the handler
+//! through a bounded channel — giving HTTP/2 flow-control backpressure.
+
+use std::collections::HashSet;
+use std::convert::Infallible;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::task::{Context, Poll};
+use std::time::Instant;
+
+use http::{HeaderValue, Request, Response};
+use http_body_util::BodyExt;
+use magnus::{Value, value::Opaque};
+use tokio::sync::mpsc;
+use tonic::Code;
+use tonic::body::Body as TonicBody;
+
+use crate::body::{ErrorBody, ResponseBody, ServerStats};
+use crate::framing::{MessageReader, parse_grpc_timeout};
+use crate::gvl::{DispatchContext, RequestError, StreamMessage};
+
+/// Messages buffered per direction before backpressure applies.
+const STREAM_BUFFER: usize = 32;
+
+/// Request metadata larger than this (sum of names + values) is rejected.
+const MAX_METADATA_BYTES: usize = 8 * 1024;
+
+/// Headers consumed by the transport and not exposed to handlers.
+fn is_reserved_header(name: &str) -> bool {
+    matches!(
+        name,
+        "content-type" | "te" | "grpc-timeout" | "grpc-encoding" | "grpc-accept-encoding"
+    )
+}
+
+pub struct DispatchInner {
+    pub callback: Opaque<Value>,
+    pub routes: HashSet<String>,
+    pub request_tx: crossbeam_channel::Sender<Box<DispatchContext>>,
+    pub stats: Arc<ServerStats>,
+    pub max_message_size: usize,
+    pub compression: Arc<AtomicBool>,
+}
+
+#[derive(Clone)]
+pub struct DispatchService(pub Arc<DispatchInner>);
+
+impl tower::Service<Request<axum::body::Body>> for DispatchService {
+    type Response = Response<TonicBody>;
+    type Error = Infallible;
+    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Infallible>> + Send>>;
+
+    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(&mut self, req: Request<axum::body::Body>) -> Self::Future {
+        let response = self.0.handle(req);
+        Box::pin(async move { Ok(response) })
+    }
+}
+
+fn grpc_response(body: TonicBody, compressed: bool) -> Response<TonicBody> {
+    let mut resp = Response::new(body);
+    let headers = resp.headers_mut();
+    headers.insert(
+        http::header::CONTENT_TYPE,
+        HeaderValue::from_static("application/grpc"),
+    );
+    headers.insert("grpc-accept-encoding", HeaderValue::from_static("gzip,identity"));
+    if compressed {
+        headers.insert("grpc-encoding", HeaderValue::from_static("gzip"));
+    }
+    resp
+}
+
+impl DispatchInner {
+    fn reject(&self, path: &str, start: Instant, code: Code, message: &str) -> Response<TonicBody> {
+        self.stats.record(path, code as u32, message, start);
+        grpc_response(TonicBody::new(ErrorBody::new(code, message)), false)
+    }
+
+    fn handle(&self, req: Request<axum::body::Body>) -> Response<TonicBody> {
+        let start = Instant::now();
+        self.stats.total.fetch_add(1, Ordering::Relaxed);
+        let path = req.uri().path().to_owned();
+        let headers = req.headers();
+
+        if req.method() != http::Method::POST {
+            return self.reject(&path, start, Code::Unimplemented, "gRPC requires HTTP POST");
+        }
+
+        // application/grpc, optionally with a +proto / ;params suffix.
+        // gRPC-Web and other framings are not supported.
+        let content_type = headers
+            .get(http::header::CONTENT_TYPE)
+            .and_then(|v| v.to_str().ok())
+            .unwrap_or("");
+        let base = content_type.split(';').next().unwrap_or("").trim();
+        if base != "application/grpc" && base != "application/grpc+proto" {
+            return self.reject(&path, start, Code::Unimplemented, "unsupported content-type");
+        }
+
+        if !self.routes.contains(&path) {
+            return self.reject(&path, start, Code::Unimplemented, "unknown method");
+        }
+
+        let gzip_request = match headers.get("grpc-encoding").map(|v| v.to_str().unwrap_or("")) {
+            None | Some("identity") => false,
+            Some("gzip") => true,
+            Some(_) => return self.reject(&path, start, Code::Unimplemented, "unsupported grpc-encoding"),
+        };
+
+        let accepts_gzip = headers
+            .get("grpc-accept-encoding")
+            .and_then(|v| v.to_str().ok())
+            .is_some_and(|v| v.split(',').any(|e| e.trim() == "gzip"));
+        let compress = accepts_gzip && self.compression.load(Ordering::Relaxed);
+
+        let timeout = headers
+            .get("grpc-timeout")
+            .and_then(|v| v.to_str().ok())
+            .and_then(parse_grpc_timeout);
+
+        let mut metadata_bytes = 0;
+        let mut metadata = Vec::with_capacity(headers.len());
+        for (name, value) in headers {
+            let name = name.as_str();
+            if is_reserved_header(name) {
+                continue;
+            }
+            metadata_bytes += name.len() + value.len();
+            // Non-UTF-8 values cannot be represented as Ruby Strings
+            // faithfully; binary metadata uses `-bin` + base64 per spec.
+            if let Ok(value) = value.to_str() {
+                metadata.push((name.to_owned(), value.to_owned()));
+            }
+        }
+        if metadata_bytes > MAX_METADATA_BYTES {
+            return self.reject(
+                &path,
+                start,
+                Code::ResourceExhausted,
+                "request metadata too large",
+            );
+        }
+
+        let (request_tx, request_rx) = mpsc::channel::<Vec<u8>>(STREAM_BUFFER);
+        let (response_tx, response_rx) = mpsc::channel::<StreamMessage>(STREAM_BUFFER);
+        let request_error = Arc::new(RequestError::default());
+
+        let ctx = Box::new(DispatchContext {
+            callback: self.callback,
+            path: path.clone(),
+            metadata,
+            timeout_ms: timeout.map(|t| t.as_millis()),
+            response_tx,
+            request_rx,
+            request_error: request_error.clone(),
+        });
+
+        match self.request_tx.try_send(ctx) {
+            Ok(()) => {}
+            Err(crossbeam_channel::TrySendError::Full(_)) => {
+                return self.reject(
+                    &path,
+                    start,
+                    Code::ResourceExhausted,
+                    "server busy: dispatch queue full",
+                );
+            }
+            Err(crossbeam_channel::TrySendError::Disconnected(_)) => {
+                return self.reject(&path, start, Code::Unavailable, "server shutting down");
+            }
+        }
+
+        let max_message_size = self.max_message_size;
+        tokio::spawn(feed_request(
+            req.into_body(),
+            request_tx,
+            request_error,
+            max_message_size,
+            gzip_request,
+        ));
+
+        let body = ResponseBody {
+            receiver: Some(response_rx),
+            stats: self.stats.clone(),
+            path,
+            start,
+            compress,
+            deadline: timeout.map(|t| Box::pin(tokio::time::sleep(t))),
+            max_message_size,
+            finished: false,
+        };
+        grpc_response(TonicBody::new(body), compress)
+    }
+}
+
+/// Read the request body and forward each unframed message to the
+/// handler.  Dropping `tx` signals end of stream; errors are recorded
+/// in `error` first so the reader can distinguish them.
+async fn feed_request(
+    body: axum::body::Body,
+    tx: mpsc::Sender<Vec<u8>>,
+    error: Arc<RequestError>,
+    max_message_size: usize,
+    gzip: bool,
+) {
+    let mut reader = MessageReader::new(max_message_size, gzip);
+    let mut body = body;
+    loop {
+        let frame = match body.frame().await {
+            None => break,
+            Some(Ok(frame)) => frame,
+            Some(Err(_)) => {
+                error.set(Code::Cancelled, "request stream reset by client");
+                return;
+            }
+        };
+        let Ok(data) = frame.into_data() else {
+            continue; // request trailers carry nothing for us
+        };
+        reader.push(&data);
+        loop {
+            match reader.try_next() {
+                Ok(Some(message)) => {
+                    if tx.send(message).await.is_err() {
+                        return; // handler finished without reading everything
+                    }
+                }
+                Ok(None) => break,
+                Err(e) => {
+                    error.set(e.code(), e.message());
+                    return;
+                }
+            }
+        }
+    }
+    if let Err(e) = reader.finish() {
+        error.set(e.code(), e.message());
+    }
+}