tako_rs_core/grpc/
message.rs1use 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
19pub struct GrpcRequest<T: Message + Default> {
24 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 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 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 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
83pub struct GrpcResponse<T: Message> {
87 message: Option<T>,
89 status: GrpcStatusCode,
91 error_message: Option<String>,
93}
94
95impl<T: Message> GrpcResponse<T> {
96 pub fn ok(message: T) -> Self {
98 Self {
99 message: Some(message),
100 status: GrpcStatusCode::Ok,
101 error_message: None,
102 }
103 }
104
105 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 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}