tonic-rails 0.1.0 → 0.3.0

data/ext/tonic_rails_native/src/body.rs

@@ -7,7 +7,7 @@
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
-use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::task::{Context, Poll};
 use std::time::Instant;
 
@@ -17,8 +17,9 @@ 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::framing::{COMPRESSION_THRESHOLD, append_metadata, frame, status_trailers};
 use crate::gvl::StreamMessage;
+use crate::log::{Level, LogSink, status_name};
 
 /// Request counters, exposed to Ruby via `Server#stats`.
 #[derive(Default)]
@@ -30,12 +31,14 @@ pub struct ServerStats {
     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,
+    pub log: LogSink,
 }
 
 impl ServerStats {
-    pub fn record(&self, path: &str, code: u32, message: &str, start: Instant) {
+    /// Count a finished call.  `transport` marks outcomes decided here in
+    /// Rust (rejections, deadlines, limits) rather than by the Ruby handler;
+    /// only those are logged, since Ruby instruments and logs its own calls.
+    pub fn record(&self, path: &str, code: u32, message: &str, start: Instant, transport: bool) {
         let counter = match Code::from_i32(code as i32) {
             Code::Ok => &self.successful,
             Code::Cancelled => &self.cancelled,
@@ -55,16 +58,15 @@ impl ServerStats {
             _ => {}
         }
 
-        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}");
+        if transport && code != Code::Ok as u32 {
+            let ms = start.elapsed().as_secs_f64() * 1000.0;
+            let status = status_name(code);
+            let line = if message.is_empty() {
+                format!("gRPC {path} {status} in {ms:.1}ms")
             } else {
-                let message = sanitize_for_log(message);
-                eprintln!("tonic-rails: path={path} code={code} duration_us={dur_us} message={message}");
-            }
+                format!("gRPC {path} {status} in {ms:.1}ms: {message}")
+            };
+            self.log.log(Level::for_code(code), &line);
         }
     }
 }
@@ -102,6 +104,8 @@ impl Body for ErrorBody {
 /// channel.  Each `Data` message becomes a gRPC frame; `End` (or the
 /// deadline) produces the trailers.
 pub struct ResponseBody {
+    /// A message already received while waiting for response headers.
+    pub pending: Option<StreamMessage>,
     pub receiver: Option<mpsc::Receiver<StreamMessage>>,
     pub stats: Arc<ServerStats>,
     pub path: String,
@@ -112,12 +116,63 @@ pub struct ResponseBody {
     pub finished: bool,
 }
 
+type Polled = Poll<Option<Result<Frame<Bytes>, Status>>>;
+
 impl ResponseBody {
-    fn finish(&mut self, code: Code, message: &str) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
+    fn finish(&mut self, code: Code, message: &str, transport: bool) -> Polled {
+        self.finish_with(code, message, transport, &[])
+    }
+
+    fn finish_with(
+        &mut self,
+        code: Code,
+        message: &str,
+        transport: bool,
+        metadata: &[(String, String)],
+    ) -> Polled {
         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)))))
+        self.stats
+            .record(&self.path, code as u32, message, self.start, transport);
+        let mut trailers = status_trailers(code as u32, message);
+        let skipped = append_metadata(&mut trailers, metadata);
+        if !skipped.is_empty() {
+            self.stats.log.log(
+                Level::Warn,
+                &format!(
+                    "gRPC {} dropped invalid or reserved trailers: {}",
+                    self.path,
+                    skipped.join(", ")
+                ),
+            );
+        }
+        Poll::Ready(Some(Ok(Frame::trailers(trailers))))
+    }
+
+    /// Turn one handler message into a frame (None: nothing to emit).
+    fn handle(&mut self, message: StreamMessage) -> Option<Polled> {
+        match message {
+            StreamMessage::Data(bytes) => {
+                if bytes.len() > self.max_message_size {
+                    return Some(self.finish(
+                        Code::ResourceExhausted,
+                        "response message exceeds max_message_size",
+                        true,
+                    ));
+                }
+                let compress = self.compress && bytes.len() > COMPRESSION_THRESHOLD;
+                Some(Poll::Ready(Some(Ok(Frame::data(frame(&bytes, compress))))))
+            }
+            StreamMessage::End { status, trailers } => Some(match status {
+                None => self.finish_with(Code::Ok, "", false, &trailers),
+                Some((code, message)) => {
+                    self.finish_with(Code::from_i32(code as i32), &message, false, &trailers)
+                }
+            }),
+            // Headers are only valid first and are consumed by the
+            // dispatcher; a late one has nowhere to go.
+            StreamMessage::Headers(_) => None,
+        }
     }
 }
 
@@ -125,39 +180,36 @@ 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);
-        };
+    fn poll_frame(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Polled {
+        if let Some(message) = self.pending.take()
+            && let Some(poll) = self.handle(message)
+        {
+            return poll;
+        }
 
         // 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",
-                    );
+        loop {
+            let Some(receiver) = self.receiver.as_mut() else {
+                return Poll::Ready(None);
+            };
+            match receiver.poll_recv(cx) {
+                Poll::Ready(Some(message)) => {
+                    if let Some(poll) = self.handle(message) {
+                        return poll;
+                    }
                 }
-                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", true);
+                }
+                Poll::Pending => break,
             }
-            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");
+            return self.finish(Code::DeadlineExceeded, "deadline exceeded", true);
         }
         Poll::Pending
     }
@@ -173,12 +225,6 @@ impl Drop for ResponseBody {
         // (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)
-                );
-            }
         }
     }
 }
@@ -197,6 +243,7 @@ mod tests {
 
     fn response_body(rx: mpsc::Receiver<StreamMessage>, stats: Arc<ServerStats>) -> ResponseBody {
         ResponseBody {
+            pending: None,
             receiver: Some(rx),
             stats,
             path: "/test.Svc/Method".into(),
@@ -231,7 +278,11 @@ mod tests {
         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();
+        tx.try_send(StreamMessage::End {
+            status: None,
+            trailers: vec![("x-done".into(), "yes".into())],
+        })
+        .ok();
 
         match poll(&mut body) {
             Poll::Ready(Some(Ok(frame))) => {
@@ -242,7 +293,9 @@ mod tests {
             }
             _ => panic!("expected data"),
         }
-        assert_eq!(trailers(poll(&mut body))["grpc-status"], "0");
+        let t = trailers(poll(&mut body));
+        assert_eq!(t["grpc-status"], "0");
+        assert_eq!(t["x-done"], "yes");
         assert!(matches!(poll(&mut body), Poll::Ready(None)));
         drop(body);
         assert_eq!(stats.successful.load(Ordering::Relaxed), 1);
@@ -263,12 +316,23 @@ mod tests {
     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();
+        tx.try_send(StreamMessage::error(Code::PermissionDenied, "denied"))
+            .ok();
         let t = trailers(poll(&mut body));
         assert_eq!(t["grpc-status"], "7");
         assert_eq!(t["grpc-message"], "denied");
     }
 
+    #[test]
+    fn response_body_starts_with_a_pending_message() {
+        let (tx, rx) = mpsc::channel(4);
+        let mut body = response_body(rx, Arc::new(ServerStats::default()));
+        body.pending = Some(StreamMessage::Data(b"x".to_vec()));
+        tx.try_send(StreamMessage::error(Code::Aborted, "stop")).ok();
+        assert!(matches!(poll(&mut body), Poll::Ready(Some(Ok(f))) if f.is_data()));
+        assert_eq!(trailers(poll(&mut body))["grpc-status"], "10");
+    }
+
     #[test]
     fn response_body_closed_channel_is_internal() {
         let (tx, rx) = mpsc::channel(4);

data/ext/tonic_rails_native/src/dispatch.rs

@@ -21,10 +21,12 @@ use magnus::{Value, value::Opaque};
 use tokio::sync::mpsc;
 use tonic::Code;
 use tonic::body::Body as TonicBody;
+use tonic::transport::server::{TcpConnectInfo, TlsConnectInfo};
 
 use crate::body::{ErrorBody, ResponseBody, ServerStats};
-use crate::framing::{MessageReader, parse_grpc_timeout};
-use crate::gvl::{DispatchContext, RequestError, StreamMessage};
+use crate::framing::{MessageReader, append_metadata, parse_grpc_timeout};
+use crate::gvl::{DispatchContext, Peer, RequestError, StreamMessage};
+use crate::log::Level;
 
 /// Messages buffered per direction before backpressure applies.
 const STREAM_BUFFER: usize = 32;
@@ -62,8 +64,8 @@ impl tower::Service<Request<axum::body::Body>> for DispatchService {
     }
 
     fn call(&mut self, req: Request<axum::body::Body>) -> Self::Future {
-        let response = self.0.handle(req);
-        Box::pin(async move { Ok(response) })
+        let inner = self.0.clone();
+        Box::pin(async move { Ok(inner.handle(req).await) })
     }
 }
 
@@ -83,11 +85,11 @@ fn grpc_response(body: TonicBody, compressed: bool) -> Response<TonicBody> {
 
 impl DispatchInner {
     fn reject(&self, path: &str, start: Instant, code: Code, message: &str) -> Response<TonicBody> {
-        self.stats.record(path, code as u32, message, start);
+        self.stats.record(path, code as u32, message, start, true);
         grpc_response(TonicBody::new(ErrorBody::new(code, message)), false)
     }
 
-    fn handle(&self, req: Request<axum::body::Body>) -> Response<TonicBody> {
+    async 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();
@@ -164,6 +166,8 @@ impl DispatchInner {
             response_tx,
             request_rx,
             request_error: request_error.clone(),
+            log: self.stats.log.clone(),
+            peer: peer_of(&req),
         });
 
         match self.request_tx.try_send(ctx) {
@@ -190,17 +194,72 @@ impl DispatchInner {
             gzip_request,
         ));
 
+        // Response headers wait for the handler's first output, so it can
+        // set initial metadata; unary calls send headers with the response
+        // as usual.  The client's deadline applies while we wait.
+        let mut response_rx = response_rx;
+        let mut deadline = timeout.map(|t| Box::pin(tokio::time::sleep(t)));
+        let first = match deadline.as_mut() {
+            Some(deadline) => tokio::select! {
+                biased;
+                message = response_rx.recv() => Some(message),
+                _ = deadline.as_mut() => None,
+            },
+            None => Some(response_rx.recv().await),
+        };
+        let (headers, pending) = match first {
+            None => return self.reject(&path, start, Code::DeadlineExceeded, "deadline exceeded"),
+            Some(None) => {
+                return self.reject(&path, start, Code::Internal, "handler ended without a response");
+            }
+            Some(Some(StreamMessage::Headers(headers))) => (headers, None),
+            Some(Some(message)) => (Vec::new(), Some(message)),
+        };
+
         let body = ResponseBody {
+            pending,
             receiver: Some(response_rx),
             stats: self.stats.clone(),
             path,
             start,
             compress,
-            deadline: timeout.map(|t| Box::pin(tokio::time::sleep(t))),
+            deadline,
             max_message_size,
             finished: false,
         };
-        grpc_response(TonicBody::new(body), compress)
+        let path = body.path.clone();
+        let mut response = grpc_response(TonicBody::new(body), compress);
+        let skipped = append_metadata(response.headers_mut(), &headers);
+        if !skipped.is_empty() {
+            self.stats.log.log(
+                Level::Warn,
+                &format!(
+                    "gRPC {path} dropped invalid or reserved response headers: {}",
+                    skipped.join(", ")
+                ),
+            );
+        }
+        response
+    }
+}
+
+/// The client's address and, if it presented one over TLS, its
+/// certificate chain.  tonic records both on the request.
+fn peer_of<B>(req: &Request<B>) -> Peer {
+    let address = req
+        .extensions()
+        .get::<TcpConnectInfo>()
+        .and_then(TcpConnectInfo::remote_addr)
+        .map(|addr| addr.to_string());
+    let certificates = req
+        .extensions()
+        .get::<TlsConnectInfo<TcpConnectInfo>>()
+        .and_then(TlsConnectInfo::peer_certs)
+        .map(|chain| chain.iter().map(|der| der.to_vec()).collect())
+        .unwrap_or_default();
+    Peer {
+        address,
+        certificates,
     }
 }
 

data/ext/tonic_rails_native/src/framing.rs

@@ -190,6 +190,35 @@ pub fn sanitize_for_log(message: &str) -> String {
         .collect()
 }
 
+/// Headers a handler may not set: they belong to HTTP/2 or the gRPC
+/// protocol itself.
+fn is_reserved_metadata(name: &str) -> bool {
+    matches!(
+        name,
+        "content-type" | "te" | "user-agent" | "content-length" | "connection" | "transfer-encoding"
+    ) || (name.starts_with("grpc-") && name != "grpc-status-details-bin")
+}
+
+/// Append handler-supplied metadata to `headers`, skipping names or values
+/// that are invalid HTTP/2 headers or reserved by gRPC.  Returns the names
+/// that were skipped so the caller can report them.
+pub fn append_metadata(headers: &mut HeaderMap, metadata: &[(String, String)]) -> Vec<String> {
+    let mut skipped = Vec::new();
+    for (name, value) in metadata {
+        let header = (!is_reserved_metadata(name))
+            .then(|| http::HeaderName::from_bytes(name.as_bytes()).ok())
+            .flatten()
+            .filter(|n| n.as_str() == name.as_str()); // must already be lower-case
+        match (header, HeaderValue::from_str(value)) {
+            (Some(header), Ok(value)) => {
+                headers.append(header, value);
+            }
+            _ => skipped.push(name.clone()),
+        }
+    }
+    skipped
+}
+
 pub fn status_trailers(code: u32, message: &str) -> HeaderMap {
     let mut trailers = HeaderMap::with_capacity(2);
     trailers.insert("grpc-status", HeaderValue::from(code));
@@ -342,6 +371,27 @@ mod tests {
         assert!(t.get("grpc-message").is_none());
     }
 
+    #[test]
+    fn append_metadata_skips_reserved_and_invalid() {
+        let mut headers = HeaderMap::new();
+        let skipped = append_metadata(
+            &mut headers,
+            &[
+                ("x-request-id".into(), "abc".into()),
+                ("x-trace-bin".into(), "AAE=".into()),
+                ("grpc-status-details-bin".into(), "CAU=".into()),
+                ("grpc-status".into(), "0".into()),
+                ("content-type".into(), "text/plain".into()),
+                ("Upper".into(), "x".into()),
+                ("x-bad".into(), "line\nbreak".into()),
+            ],
+        );
+        assert_eq!(headers["x-request-id"], "abc");
+        assert_eq!(headers["x-trace-bin"], "AAE=");
+        assert_eq!(headers["grpc-status-details-bin"], "CAU=");
+        assert_eq!(skipped, ["grpc-status", "content-type", "Upper", "x-bad"]);
+    }
+
     #[test]
     fn sanitize_for_log_strips_controls() {
         assert_eq!(sanitize_for_log("a\nb\tc"), "a b c");