itsi 0.1.14 → 0.1.19

data/crates/itsi_server/src/ruby_types/itsi_body_proxy/mod.rs

@@ -12,10 +12,12 @@ use std::sync::{
 };
 use tokio_stream::StreamExt;
 
+use crate::server::size_limited_incoming::SizeLimitedIncoming;
+
 #[magnus::wrap(class = "Itsi::BodyProxy", free_immediately, size)]
 #[derive(Debug, Clone)]
 pub struct ItsiBodyProxy {
-    pub incoming: Arc<Mutex<BodyDataStream<Incoming>>>,
+    pub incoming: Arc<Mutex<BodyDataStream<SizeLimitedIncoming<Incoming>>>>,
     pub closed: Arc<AtomicBool>,
     pub buf: Arc<Mutex<Vec<u8>>>,
 }
@@ -34,7 +36,7 @@ impl ItsiBody {
     }
 }
 impl ItsiBodyProxy {
-    pub fn new(incoming: Incoming) -> Self {
+    pub fn new(incoming: SizeLimitedIncoming<Incoming>) -> Self {
         ItsiBodyProxy {
             incoming: Arc::new(Mutex::new(incoming.into_data_stream())),
             closed: Arc::new(AtomicBool::new(false)),
@@ -50,7 +52,7 @@ impl ItsiBodyProxy {
             if let Some(chunk) = block_on(stream.next()) {
                 let chunk = chunk.map_err(|err| {
                     magnus::Error::new(
-                        magnus::exception::exception(),
+                        magnus::exception::standard_error(),
                         format!("Error reading body {:?}", err),
                     )
                 })?;
@@ -82,7 +84,7 @@ impl ItsiBodyProxy {
             if let Some(chunk) = block_on(stream.next()) {
                 let chunk = chunk.map_err(|err| {
                     magnus::Error::new(
-                        magnus::exception::exception(),
+                        magnus::exception::standard_error(),
                         format!("Error reading body {:?}", err),
                     )
                 })?;
@@ -107,7 +109,7 @@ impl ItsiBodyProxy {
         while let Some(chunk) = block_on(stream.next()) {
             let chunk = chunk.map_err(|err| {
                 magnus::Error::new(
-                    magnus::exception::exception(),
+                    magnus::exception::standard_error(),
                     format!("Error reading body {:?}", err),
                 )
             })?;
@@ -129,7 +131,7 @@ impl ItsiBodyProxy {
     fn verify_open(&self) -> MagnusResult<()> {
         if self.closed.load(atomic::Ordering::SeqCst) {
             return Err(magnus::Error::new(
-                magnus::exception::exception(),
+                magnus::exception::standard_error(),
                 "Body stream is closed",
             ));
         }

data/crates/itsi_server/src/ruby_types/itsi_grpc_call.rs

@@ -0,0 +1,344 @@
+use super::itsi_grpc_response_stream::ItsiGrpcResponseStream;
+use crate::prelude::*;
+use crate::server::http_message_types::{HttpRequest, HttpResponse};
+use crate::server::{byte_frame::ByteFrame, request_job::RequestJob};
+use crate::services::itsi_http_service::HttpRequestContext;
+use async_compression::futures::bufread::{GzipDecoder, GzipEncoder, ZlibDecoder, ZlibEncoder};
+use bytes::Bytes;
+use derive_more::Debug;
+use futures::{executor::block_on, io::Cursor, AsyncReadExt};
+use http::{request::Parts, Response, StatusCode};
+use http_body_util::{combinators::BoxBody, BodyExt, Empty};
+use itsi_error::CLIENT_CONNECTION_CLOSED;
+use itsi_rb_helpers::{print_rb_backtrace, HeapValue};
+use itsi_tracing::debug;
+use magnus::{
+    block::Proc,
+    error::{ErrorType, Result as MagnusResult},
+    Error, Symbol,
+};
+use magnus::{
+    value::{LazyId, ReprValue},
+    Ruby, Value,
+};
+use regex::Regex;
+use std::sync::LazyLock;
+use std::{collections::HashMap, sync::Arc, time::Instant};
+use tokio::sync::mpsc::{self};
+
+static ID_MESSAGE: LazyId = LazyId::new("message");
+static MIN_GZIP_SIZE: u32 = 128;
+static MIN_DEFLATE_SIZE: u32 = 128;
+static METHOD_NAME_REGEX: LazyLock<Regex> =
+    LazyLock::new(|| Regex::new(r"([a-z])([A-Z])").expect("Failed to compile regex"));
+
+#[derive(Debug)]
+#[magnus::wrap(class = "Itsi::GrpcCall", free_immediately, size)]
+pub struct ItsiGrpcCall {
+    pub parts: Parts,
+    pub start: Instant,
+    pub compression_in: CompressionAlgorithm,
+    pub compression_out: CompressionAlgorithm,
+    #[debug(skip)]
+    pub context: HttpRequestContext,
+    #[debug(skip)]
+    pub stream: ItsiGrpcResponseStream,
+}
+
+#[derive(Debug, Clone)]
+pub enum CompressionAlgorithm {
+    None,
+    Deflate,
+    Gzip,
+}
+
+impl ItsiGrpcCall {
+    pub fn service_name(&self) -> MagnusResult<String> {
+        let path = self.parts.uri.path();
+        Ok(path.split('/').nth_back(1).unwrap().to_string())
+    }
+
+    pub fn method_name(&self) -> MagnusResult<Symbol> {
+        let path = self.parts.uri.path();
+        let method_name = path.split('/').nth_back(0).unwrap();
+        let snake_case_method_name = METHOD_NAME_REGEX
+            .replace_all(method_name, "${1}_${2}")
+            .to_lowercase();
+        Ok(Symbol::new(snake_case_method_name))
+    }
+
+    pub fn stream(&self) -> MagnusResult<ItsiGrpcResponseStream> {
+        Ok(self.stream.clone())
+    }
+
+    pub fn timeout(&self) -> MagnusResult<Option<f64>> {
+        let timeout_str = self
+            .parts
+            .headers
+            .get("grpc-timeout")
+            .and_then(|hv| hv.to_str().ok())
+            .unwrap_or("");
+        Ok(parse_grpc_timeout(timeout_str).ok())
+    }
+
+    pub fn is_cancelled(&self) -> MagnusResult<bool> {
+        self.stream.is_cancelled()
+    }
+
+    pub fn add_headers(&self, headers: HashMap<Bytes, Vec<Bytes>>) -> MagnusResult<()> {
+        self.stream.add_headers(headers)
+    }
+
+    pub fn content_type_str(&self) -> &str {
+        self.parts
+            .headers
+            .get("Content-Type")
+            .and_then(|hv| hv.to_str().ok())
+            .unwrap_or("application/x-www-form-urlencoded")
+    }
+
+    pub fn is_json(&self) -> bool {
+        self.content_type_str() == "application/json"
+    }
+
+    pub fn process(self, ruby: &Ruby, app_proc: Arc<HeapValue<Proc>>) -> magnus::error::Result<()> {
+        let response = self.stream.clone();
+        let result = app_proc.call::<_, Value>((self,));
+        if let Err(err) = result {
+            Self::internal_error(ruby, response, err);
+        }
+        Ok(())
+    }
+
+    pub fn internal_error(_ruby: &Ruby, stream: ItsiGrpcResponseStream, err: Error) {
+        if let Some(rb_err) = err.value() {
+            print_rb_backtrace(rb_err);
+            stream.internal_server_error(err.to_string());
+        } else {
+            stream.internal_server_error(err.to_string());
+        }
+    }
+
+    pub(crate) async fn process_request(
+        app: Arc<HeapValue<Proc>>,
+        hyper_request: HttpRequest,
+        context: &HttpRequestContext,
+        nonblocking: bool,
+    ) -> itsi_error::Result<HttpResponse> {
+        let (request, mut receiver) = ItsiGrpcCall::new(hyper_request, context).await;
+        let shutdown_channel = context.service.shutdown_channel.clone();
+        let response_stream = request.stream.clone();
+        let sender = if nonblocking {
+            &context.nonblocking_sender
+        } else {
+            &context.sender
+        };
+        match sender
+            .send(RequestJob::ProcessGrpcRequest(request, app))
+            .await
+        {
+            Err(err) => {
+                error!("Error occurred: {}", err);
+                let mut response = Response::new(BoxBody::new(Empty::new()));
+                *response.status_mut() = StatusCode::BAD_REQUEST;
+                Ok(response)
+            }
+            _ => match receiver.recv().await {
+                Some(first_frame) => Ok(response_stream
+                    .build_response(first_frame, receiver, shutdown_channel)
+                    .await),
+                None => Ok(Response::new(BoxBody::new(Empty::new()))),
+            },
+        }
+    }
+
+    pub fn is_connection_closed_err(ruby: &Ruby, err: &Error) -> bool {
+        match err.error_type() {
+            ErrorType::Jump(_) => false,
+            ErrorType::Error(_, _) => false,
+            ErrorType::Exception(exception) => {
+                exception.is_kind_of(ruby.exception_eof_error())
+                    && err
+                        .value()
+                        .map(|v| {
+                            v.funcall::<_, _, String>(*ID_MESSAGE, ())
+                                .unwrap_or("".to_string())
+                                .eq(CLIENT_CONNECTION_CLOSED)
+                        })
+                        .unwrap_or(false)
+            }
+        }
+    }
+
+    pub fn should_compress_output(&self, message_size: u32) -> bool {
+        match self.compression_out {
+            CompressionAlgorithm::Gzip => message_size > MIN_GZIP_SIZE,
+            CompressionAlgorithm::Deflate => message_size > MIN_DEFLATE_SIZE,
+            CompressionAlgorithm::None => false,
+        }
+    }
+
+    pub fn compress_output(&self, bytes: Bytes) -> MagnusResult<Bytes> {
+        match self.compression_out {
+            CompressionAlgorithm::Gzip => Self::compress_gzip(bytes),
+            CompressionAlgorithm::Deflate => Self::compress_deflate(bytes),
+            CompressionAlgorithm::None => Ok(bytes),
+        }
+    }
+
+    pub fn decompress_input(&self, bytes: Bytes) -> MagnusResult<Bytes> {
+        match self.compression_in {
+            CompressionAlgorithm::Gzip => Self::decompress_gzip(bytes),
+            CompressionAlgorithm::Deflate => Self::decompress_deflate(bytes),
+            CompressionAlgorithm::None => Ok(bytes),
+        }
+    }
+
+    fn decompress_deflate(input: Bytes) -> MagnusResult<Bytes> {
+        let cursor = Cursor::new(input);
+        let mut decoder = ZlibDecoder::new(cursor);
+
+        let result = block_on(async {
+            let mut output = Vec::new();
+            decoder.read_to_end(&mut output).await?;
+            Ok(Bytes::from(output))
+        })
+        .map_err(|e: std::io::Error| {
+            Error::new(
+                magnus::exception::standard_error(),
+                format!("deflate decompression failed: {}", e),
+            )
+        })?;
+
+        Ok(result)
+    }
+
+    fn decompress_gzip(input: Bytes) -> MagnusResult<Bytes> {
+        let cursor = Cursor::new(input);
+        let mut decoder = GzipDecoder::new(cursor);
+
+        let result = block_on(async {
+            let mut output = Vec::new();
+            decoder.read_to_end(&mut output).await?;
+            Ok(Bytes::from(output))
+        })
+        .map_err(|e: std::io::Error| {
+            Error::new(
+                magnus::exception::standard_error(),
+                format!("gzip decompression failed: {}", e),
+            )
+        })?;
+
+        Ok(result)
+    }
+
+    fn compress_gzip(input: Bytes) -> MagnusResult<Bytes> {
+        let mut output = Vec::with_capacity(input.len() / 2);
+        let cursor = Cursor::new(input);
+        let mut encoder = GzipEncoder::new(cursor);
+
+        let result = block_on(async {
+            encoder.read_to_end(&mut output).await?;
+            Ok::<Bytes, std::io::Error>(output.into())
+        })
+        .map_err(|e| {
+            Error::new(
+                magnus::exception::standard_error(),
+                format!("gzip compression failed: {e}"),
+            )
+        })?;
+
+        Ok(result)
+    }
+
+    fn compress_deflate(input: Bytes) -> MagnusResult<Bytes> {
+        let mut output = Vec::with_capacity(input.len() / 2);
+        let cursor = Cursor::new(input);
+        let mut encoder = ZlibEncoder::new(cursor);
+
+        let result = block_on(async {
+            encoder.read_to_end(&mut output).await?;
+            Ok::<Bytes, std::io::Error>(output.into())
+        })
+        .map_err(|e| {
+            Error::new(
+                magnus::exception::standard_error(),
+                format!("deflate compression failed: {e}"),
+            )
+        })?;
+
+        Ok(result)
+    }
+
+    pub(crate) async fn new(
+        request: HttpRequest,
+        context: &HttpRequestContext,
+    ) -> (ItsiGrpcCall, mpsc::Receiver<ByteFrame>) {
+        let (parts, body) = request.into_parts();
+        let response_channel = mpsc::channel::<ByteFrame>(100);
+        let compression_in: CompressionAlgorithm = match parts.headers.get("grpc-encoding") {
+            Some(encoding) => match encoding.to_str() {
+                Ok(encoding) => match encoding {
+                    "gzip" => CompressionAlgorithm::Gzip,
+                    "deflate" => CompressionAlgorithm::Deflate,
+                    _ => CompressionAlgorithm::None,
+                },
+                Err(_) => CompressionAlgorithm::None,
+            },
+            None => CompressionAlgorithm::None,
+        };
+        let compression_out: CompressionAlgorithm = match parts.headers.get("grpc-accept-encoding")
+        {
+            Some(accept_encoding) => match accept_encoding.to_str() {
+                Ok(accept_encoding) => {
+                    let encodings: Vec<&str> =
+                        accept_encoding.split(',').map(|s| s.trim()).collect();
+                    if encodings.contains(&"gzip") {
+                        CompressionAlgorithm::Gzip
+                    } else if encodings.contains(&"deflate") {
+                        CompressionAlgorithm::Deflate
+                    } else {
+                        CompressionAlgorithm::None
+                    }
+                }
+                Err(_) => CompressionAlgorithm::None,
+            },
+            None => CompressionAlgorithm::None,
+        };
+        (
+            Self {
+                context: context.clone(),
+                start: Instant::now(),
+                compression_out: compression_out.clone(),
+                compression_in,
+                parts,
+                stream: ItsiGrpcResponseStream::new(
+                    compression_out,
+                    response_channel.0,
+                    body.into_data_stream(),
+                )
+                .await,
+            },
+            response_channel.1,
+        )
+    }
+}
+
+fn parse_grpc_timeout(timeout_str: &str) -> Result<f64> {
+    if timeout_str.len() < 2 {
+        return Err("Timeout string too short".into());
+    }
+    let (value_str, unit) = timeout_str.split_at(timeout_str.len() - 1);
+    let value: u64 = value_str.parse().map_err(|_| "Invalid timeout value")?;
+    let duration_secs = match unit {
+        "n" => value as f64 / 1_000_000_000.0, // nanoseconds
+        "u" => value as f64 / 1_000_000.0,     // microseconds
+        "m" => value as f64 / 1_000.0,         // milliseconds
+        "S" => value as f64,                   // seconds
+        "M" => value as f64 * 60.0,            // minutes
+        "H" => value as f64 * 3600.0,          // hours
+        _ => return Err("Invalid timeout unit".into()),
+    };
+
+    Ok(duration_secs)
+}

data/crates/itsi_server/src/ruby_types/{itsi_grpc_stream → itsi_grpc_response_stream}/mod.rs

@@ -1,78 +1,57 @@
-use std::{collections::HashMap, sync::Arc};
-
-use crate::server::{
-    byte_frame::ByteFrame, serve_strategy::single_mode::RunningPhase, types::HttpResponse,
-};
+use super::itsi_grpc_call::CompressionAlgorithm;
+use crate::prelude::*;
+use crate::server::http_message_types::HttpResponse;
+use crate::server::size_limited_incoming::SizeLimitedIncoming;
+use crate::server::{byte_frame::ByteFrame, serve_strategy::single_mode::RunningPhase};
 use bytes::Bytes;
 use derive_more::Debug;
-use futures::{executor::block_on, stream::unfold};
+use futures::stream::unfold;
+use http::Version;
 use http::{
-    header::{HeaderName, HeaderValue, CONTENT_TYPE},
+    header::{HeaderName, HeaderValue},
     HeaderMap, Response,
 };
 use http_body_util::{combinators::BoxBody, BodyDataStream, BodyExt, Empty, Full, StreamBody};
 use hyper::body::{Frame, Incoming};
 use magnus::error::Result as MagnusResult;
+use nix::unistd::pipe;
 use parking_lot::Mutex;
-use tokio::sync::{
-    mpsc::{Receiver, Sender},
-    oneshot, watch,
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::{
+    collections::HashMap,
+    os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd},
+    sync::Arc,
+};
+use tokio::{
+    spawn,
+    sync::{
+        mpsc::{self, Sender},
+        oneshot, watch,
+    },
 };
 use tokio_stream::{wrappers::ReceiverStream, StreamExt};
-use tracing::{error, info, warn};
 
 #[derive(Debug, Clone)]
-#[magnus::wrap(class = "Itsi::GrpcStream", free_immediately, size)]
-pub struct ItsiGrpcStream {
-    pub inner: Arc<Mutex<ItsiGrpcStreamInner>>,
+#[magnus::wrap(class = "Itsi::GrpcResponseStream", free_immediately, size)]
+pub struct ItsiGrpcResponseStream {
+    pub inner: Arc<Mutex<ItsiGrpcResponseStreamInner>>,
+    pub cancelled: Arc<AtomicBool>,
 }
 
 #[derive(Debug)]
-pub struct ItsiGrpcStreamInner {
-    pub body: BodyDataStream<Incoming>,
+pub struct ItsiGrpcResponseStreamInner {
+    pub incoming_reader: Option<OwnedFd>,
     pub buf: Vec<u8>,
     pub response_sender: Sender<ByteFrame>,
     pub response: Option<HttpResponse>,
+    pub response_headers: HeaderMap,
     trailer_tx: oneshot::Sender<HeaderMap>,
     trailer_rx: Option<oneshot::Receiver<HeaderMap>>,
 }
 
-impl ItsiGrpcStreamInner {
-    pub fn read(&mut self, bytes: usize) -> MagnusResult<Bytes> {
-        let stream = &mut self.body;
-        let buf = &mut self.buf;
-        let mut result = Vec::with_capacity(bytes);
-
-        info!("Entering read with {:?}. Current buf is {:?}", bytes, buf);
-
-        // First, use any data already in the buffer
-        if !buf.is_empty() {
-            let remaining = bytes.min(buf.len());
-            result.extend_from_slice(&buf[..remaining]);
-            buf.drain(..remaining);
-        }
-
-        while result.len() < bytes {
-            if let Some(chunk) = block_on(stream.next()) {
-                let chunk = chunk.map_err(|err| {
-                    magnus::Error::new(
-                        magnus::exception::exception(),
-                        format!("Error reading body {:?}", err),
-                    )
-                })?;
-                let remaining = bytes - result.len();
-                if chunk.len() > remaining {
-                    result.extend_from_slice(&chunk[..remaining]);
-                    buf.extend_from_slice(&chunk[remaining..]);
-                } else {
-                    result.extend_from_slice(&chunk);
-                }
-            } else {
-                break;
-            }
-        }
-
-        Ok(result.into())
+impl ItsiGrpcResponseStreamInner {
+    pub fn reader(&mut self) -> MagnusResult<i32> {
+        Ok(self.incoming_reader.take().unwrap().into_raw_fd())
     }
 
     pub fn write(&mut self, bytes: Bytes) -> MagnusResult<()> {
@@ -80,8 +59,8 @@ impl ItsiGrpcStreamInner {
             .blocking_send(ByteFrame::Data(bytes))
             .map_err(|err| {
                 magnus::Error::new(
-                    magnus::exception::exception(),
-                    format!("Error writing body {:?}", err),
+                    magnus::exception::io_error(),
+                    format!("Trying to write to closed stream: {:?}", err),
                 )
             })?;
         Ok(())
@@ -101,39 +80,97 @@ impl ItsiGrpcStreamInner {
         let trailer_tx = std::mem::replace(&mut self.trailer_tx, oneshot::channel().0);
         trailer_tx.send(header_map).map_err(|err| {
             magnus::Error::new(
-                magnus::exception::exception(),
+                magnus::exception::standard_error(),
                 format!("Error sending trailers {:?}", err),
             )
         })?;
-        self.response_sender
-            .blocking_send(ByteFrame::Empty)
-            .map_err(|err| {
-                magnus::Error::new(
-                    magnus::exception::exception(),
-                    format!("Error flushing {:?}", err),
-                )
+        Ok(())
+    }
+
+    pub fn close(&mut self) -> MagnusResult<()> {
+        self.response_sender.blocking_send(ByteFrame::Empty).ok();
+        Ok(())
+    }
+
+    pub fn add_headers(&mut self, headers: HashMap<Bytes, Vec<Bytes>>) -> MagnusResult<()> {
+        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) };
+                self.response_headers.insert(&header_name, header_value);
+            }
+        }
+
         Ok(())
     }
 }
 
-impl ItsiGrpcStream {
-    pub fn new(response_sender: Sender<ByteFrame>, body: BodyDataStream<Incoming>) -> Self {
+impl ItsiGrpcResponseStream {
+    pub async fn new(
+        compression_out: CompressionAlgorithm,
+        response_sender: Sender<ByteFrame>,
+        mut body: BodyDataStream<SizeLimitedIncoming<Incoming>>,
+    ) -> Self {
         let (trailer_tx, trailer_rx) = oneshot::channel::<HeaderMap>();
-        ItsiGrpcStream {
-            inner: Arc::new(Mutex::new(ItsiGrpcStreamInner {
-                body,
+        let (pipe_read, pipe_write) = pipe().unwrap();
+
+        nix::fcntl::fcntl(
+            pipe_read.as_raw_fd(),
+            nix::fcntl::FcntlArg::F_SETFL(nix::fcntl::OFlag::O_NONBLOCK),
+        )
+        .unwrap();
+
+        nix::fcntl::fcntl(
+            pipe_write.as_raw_fd(),
+            nix::fcntl::FcntlArg::F_SETFL(nix::fcntl::OFlag::O_NONBLOCK),
+        )
+        .unwrap();
+
+        let pipe_raw_fd = pipe_write.into_raw_fd();
+
+        let cancelled = Arc::new(AtomicBool::new(false));
+        let cancelled_clone = cancelled.clone();
+        spawn(async move {
+            use std::io::Write;
+            let mut write_end = unsafe { std::fs::File::from_raw_fd(pipe_raw_fd) };
+            while let Some(Ok(body)) = body.next().await {
+                write_end.write_all(&body).unwrap();
+            }
+            cancelled_clone.store(true, Ordering::SeqCst);
+        });
+
+        let mut response_headers = HeaderMap::new();
+
+        match compression_out {
+            CompressionAlgorithm::None => (),
+            CompressionAlgorithm::Deflate => {
+                response_headers.insert("grpc-encoding", "deflate".parse().unwrap());
+            }
+            CompressionAlgorithm::Gzip => {
+                response_headers.insert("grpc-encoding", "gzip".parse().unwrap());
+            }
+        }
+        ItsiGrpcResponseStream {
+            inner: Arc::new(Mutex::new(ItsiGrpcResponseStreamInner {
                 buf: Vec::new(),
+                response_headers,
+                incoming_reader: Some(pipe_read),
                 response_sender,
                 response: Some(Response::new(BoxBody::new(Empty::new()))),
                 trailer_tx,
                 trailer_rx: Some(trailer_rx),
             })),
+            cancelled,
         }
     }
 
-    pub fn read(&self, bytes: usize) -> MagnusResult<Bytes> {
-        self.inner.lock().read(bytes)
+    pub fn reader(&self) -> MagnusResult<i32> {
+        self.inner.lock().reader()
     }
 
     pub fn write(&self, bytes: Bytes) -> MagnusResult<()> {
@@ -144,21 +181,32 @@ impl ItsiGrpcStream {
         self.inner.lock().flush()
     }
 
+    pub fn is_cancelled(&self) -> MagnusResult<bool> {
+        Ok(self.cancelled.load(Ordering::SeqCst))
+    }
+
     pub fn send_trailers(&self, trailers: HashMap<String, String>) -> MagnusResult<()> {
         self.inner.lock().send_trailers(trailers)
     }
 
+    pub fn close(&self) -> MagnusResult<()> {
+        self.inner.lock().close()
+    }
+
+    pub fn add_headers(&self, headers: HashMap<Bytes, Vec<Bytes>>) -> MagnusResult<()> {
+        self.inner.lock().add_headers(headers)
+    }
+
     pub async fn build_response(
         &self,
         first_frame: ByteFrame,
-        receiver: Receiver<ByteFrame>,
+        receiver: mpsc::Receiver<ByteFrame>,
         shutdown_rx: watch::Receiver<RunningPhase>,
     ) -> HttpResponse {
         let mut response = self.inner.lock().response.take().unwrap();
         let rx = self.inner.lock().trailer_rx.take().unwrap();
-        response
-            .headers_mut()
-            .append(CONTENT_TYPE, "application/grpc".parse().unwrap());
+        *response.version_mut() = Version::HTTP_2;
+        *response.headers_mut() = self.inner.lock().response_headers.clone();
         *response.body_mut() = if matches!(first_frame, ByteFrame::Empty) {
             BoxBody::new(Empty::new())
         } else if matches!(first_frame, ByteFrame::End(_)) {