Skip to main content

tako_rs_core/grpc/
message.rs

1//! Unary gRPC request extractor (`GrpcRequest`) and response responder
2//! (`GrpcResponse`) over a single length-prefixed protobuf frame.
3
4use http::StatusCode;
5use http_body_util::BodyExt;
6use prost::Message;
7
8use super::GrpcError;
9use super::framing::MAX_GRPC_MESSAGE_SIZE;
10use super::framing::grpc_encode;
11use super::status::GrpcStatusCode;
12use super::status::build_grpc_error_response;
13use crate::body::TakoBody;
14use crate::extractors::FromRequest;
15use crate::responder::Responder;
16use crate::types::Request;
17use crate::types::Response;
18
19/// gRPC request extractor.
20///
21/// Extracts and decodes a gRPC-framed protobuf message from the request body.
22/// Validates that the content-type is `application/grpc`.
23pub struct GrpcRequest<T: Message + Default> {
24  /// The decoded protobuf message.
25  pub message: T,
26}
27
28impl<'a, T> FromRequest<'a> for GrpcRequest<T>
29where
30  T: Message + Default + Send + 'static,
31{
32  type Error = GrpcError;
33
34  fn from_request(
35    req: &'a mut Request,
36  ) -> impl core::future::Future<Output = core::result::Result<Self, Self::Error>> + Send + 'a {
37    async move {
38      // Validate content-type
39      let ct = req
40        .headers()
41        .get(http::header::CONTENT_TYPE)
42        .and_then(|v| v.to_str().ok())
43        .unwrap_or("");
44
45      if !ct.starts_with("application/grpc") {
46        return Err(GrpcError::InvalidContentType);
47      }
48
49      // Read body
50      let body_bytes = req
51        .body_mut()
52        .collect()
53        .await
54        .map_err(|e| GrpcError::BodyReadError(e.to_string()))?
55        .to_bytes();
56
57      // Decode gRPC frame: 1 byte compressed + 4 bytes length + message
58      if body_bytes.len() < 5 {
59        return Err(GrpcError::InvalidFrame);
60      }
61
62      if body_bytes[0] != 0 {
63        return Err(GrpcError::CompressionUnsupported);
64      }
65      let msg_len =
66        u32::from_be_bytes([body_bytes[1], body_bytes[2], body_bytes[3], body_bytes[4]]) as usize;
67
68      if msg_len > MAX_GRPC_MESSAGE_SIZE {
69        return Err(GrpcError::MessageTooLarge);
70      }
71      if body_bytes.len() < 5 + msg_len {
72        return Err(GrpcError::InvalidFrame);
73      }
74
75      let message = T::decode(&body_bytes[5..5 + msg_len])
76        .map_err(|e| GrpcError::DecodeError(e.to_string()))?;
77
78      Ok(GrpcRequest { message })
79    }
80  }
81}
82
83/// gRPC response wrapper.
84///
85/// Encodes a protobuf message with gRPC framing and sets appropriate headers.
86pub struct GrpcResponse<T: Message> {
87  /// The response message (None for error-only responses).
88  message: Option<T>,
89  /// gRPC status code.
90  status: GrpcStatusCode,
91  /// Optional error message.
92  error_message: Option<String>,
93}
94
95impl<T: Message> GrpcResponse<T> {
96  /// Creates a successful gRPC response with the given message.
97  pub fn ok(message: T) -> Self {
98    Self {
99      message: Some(message),
100      status: GrpcStatusCode::Ok,
101      error_message: None,
102    }
103  }
104
105  /// Creates an error gRPC response with the given status and message.
106  pub fn error(status: GrpcStatusCode, message: impl Into<String>) -> Self {
107    Self {
108      message: None,
109      status,
110      error_message: Some(message.into()),
111    }
112  }
113}
114
115impl<T: Message> Responder for GrpcResponse<T> {
116  fn into_response(self) -> Response {
117    if self.status != GrpcStatusCode::Ok {
118      return build_grpc_error_response(self.status, self.error_message.as_deref().unwrap_or(""));
119    }
120
121    let body_bytes = match self.message {
122      Some(msg) => grpc_encode(&msg),
123      None => Vec::new(),
124    };
125
126    let mut resp = Response::new(TakoBody::from(body_bytes));
127    *resp.status_mut() = StatusCode::OK;
128    resp.headers_mut().insert(
129      http::header::CONTENT_TYPE,
130      http::HeaderValue::from_static("application/grpc"),
131    );
132    // gRPC uses trailers for status. Since we're using HTTP/1.1-compatible
133    // responses, we put the status in headers as a fallback.
134    if let Ok(val) = http::HeaderValue::from_str(&(self.status as u8).to_string()) {
135      resp.headers_mut().insert("grpc-status", val);
136    }
137    resp
138  }
139}