use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use bytes::Bytes;
use futures::stream::{LocalBoxStream, StreamExt};
use h2ts_client::{RequestBody, Trailers};
use http::{HeaderMap, HeaderName, HeaderValue};
use http_body::{Body, Frame};
use send_wrapper::SendWrapper;
use crate::client::Client;
#[derive(Clone)]
pub struct TonicService {
client: Client,
}
impl std::fmt::Debug for TonicService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TonicService").finish_non_exhaustive()
}
}
impl TonicService {
pub fn new(client: Client) -> TonicService {
TonicService { client }
}
pub fn client(&self) -> &Client {
&self.client
}
}
impl Client {
pub fn into_tonic(self) -> TonicService {
TonicService::new(self)
}
}
impl tower_service::Service<http::Request<tonic::body::Body>> for TonicService {
type Response = http::Response<ResponseBody>;
type Error = tonic::Status;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, request: http::Request<tonic::body::Body>) -> Self::Future {
let client = self.client.clone();
Box::pin(send(client, request))
}
}
async fn send(
client: Client,
request: http::Request<tonic::body::Body>,
) -> Result<http::Response<ResponseBody>, tonic::Status> {
let (parts, body) = request.into_parts();
let path = parts.uri.path().to_string();
let timeout = grpc_timeout(&parts.headers);
let headers = to_wire_headers(&parts.headers)?;
let body = RequestBody::stream(request_chunks(body));
let (response, deadline) = match timeout {
None => (client.send(&path, headers, body).await?, None),
Some(timeout) => {
use futures::future::{select, Either};
let mut timer = futures_timer::Delay::new(timeout);
let opened = {
let open = client.send(&path, headers, body);
futures::pin_mut!(open);
match select(open, &mut timer).await {
Either::Left((response, _)) => Some(response),
Either::Right(((), _)) => None,
}
};
match opened {
Some(response) => (response?, Some(timer)),
None => return Err(tonic::Status::deadline_exceeded("deadline exceeded")),
}
}
};
let mut builder = http::Response::builder().status(response.status);
for header in &response.raw_headers {
if header.name.starts_with(':') {
continue;
}
builder = builder.header(&header.name, &header.value);
}
let (body, trailers) = response.into_parts();
builder
.body(ResponseBody::new(body.boxed_local(), trailers, deadline))
.map_err(|e| tonic::Status::internal(format!("malformed response headers: {e}")))
}
fn request_chunks(body: tonic::body::Body) -> impl futures::Stream<Item = Vec<u8>> {
futures::stream::unfold(Box::pin(body), |mut body| async move {
loop {
let frame = std::future::poll_fn(|cx| body.as_mut().poll_frame(cx)).await;
match frame {
Some(Ok(frame)) => match frame.into_data() {
Ok(data) => return Some((data.to_vec(), body)),
Err(_trailers) => continue,
},
Some(Err(_)) | None => return None,
}
}
})
}
fn to_wire_headers(headers: &HeaderMap) -> Result<Vec<(String, String)>, tonic::Status> {
headers
.iter()
.map(|(name, value)| {
let value = value.to_str().map_err(|_| {
tonic::Status::internal(format!(
"metadata value for `{name}` is not valid ASCII; \
binary metadata must use a `-bin` key"
))
})?;
Ok((name.as_str().to_string(), value.to_string()))
})
.collect()
}
fn grpc_timeout(headers: &HeaderMap) -> Option<Duration> {
let raw = headers.get("grpc-timeout")?.to_str().ok()?;
let (digits, unit) = raw.split_at_checked(raw.len().checked_sub(1)?)?;
let value: u64 = digits.parse().ok()?;
Some(match unit {
"n" => Duration::from_nanos(value),
"u" => Duration::from_micros(value),
"m" => Duration::from_millis(value),
"S" => Duration::from_secs(value),
"M" => Duration::from_secs(value.checked_mul(60)?),
"H" => Duration::from_secs(value.checked_mul(3600)?),
_ => return None,
})
}
fn to_header_map(headers: std::collections::HashMap<String, String>) -> HeaderMap {
let mut map = HeaderMap::with_capacity(headers.len());
for (name, value) in headers {
if let (Ok(name), Ok(value)) =
(HeaderName::try_from(name), HeaderValue::try_from(value))
{
map.append(name, value);
}
}
map
}
pub struct ResponseBody(SendWrapper<Inner>);
struct Inner {
body: LocalBoxStream<'static, Result<Vec<u8>, h2ts_client::H2Error>>,
trailers: Trailers,
deadline: Option<futures_timer::Delay>,
ended: bool,
}
impl std::fmt::Debug for ResponseBody {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ResponseBody").field("ended", &self.0.ended).finish_non_exhaustive()
}
}
impl ResponseBody {
fn new(
body: LocalBoxStream<'static, Result<Vec<u8>, h2ts_client::H2Error>>,
trailers: Trailers,
deadline: Option<futures_timer::Delay>,
) -> ResponseBody {
ResponseBody(SendWrapper::new(Inner { body, trailers, deadline, ended: false }))
}
}
impl Body for ResponseBody {
type Data = Bytes;
type Error = tonic::Status;
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Bytes>, Self::Error>>> {
let inner = &mut *self.get_mut().0;
if inner.ended {
return Poll::Ready(None);
}
if let Some(timer) = inner.deadline.as_mut() {
if Pin::new(timer).poll(cx).is_ready() {
inner.ended = true;
return Poll::Ready(Some(Err(tonic::Status::deadline_exceeded(
"deadline exceeded",
))));
}
}
match inner.body.poll_next_unpin(cx) {
Poll::Ready(Some(Ok(chunk))) => Poll::Ready(Some(Ok(Frame::data(Bytes::from(chunk))))),
Poll::Ready(Some(Err(e))) => {
inner.ended = true;
Poll::Ready(Some(Err(tonic::Status::unavailable(format!(
"stream failed: {e}"
)))))
}
Poll::Ready(None) => {
inner.ended = true;
match inner.trailers.get() {
Some(trailers) => {
Poll::Ready(Some(Ok(Frame::trailers(to_header_map(trailers)))))
}
None => Poll::Ready(None),
}
}
Poll::Pending => Poll::Pending,
}
}
fn is_end_stream(&self) -> bool {
self.0.ended
}
}
impl From<crate::Status> for tonic::Status {
fn from(status: crate::Status) -> tonic::Status {
let mut headers = HeaderMap::new();
for (name, value) in status.metadata.to_headers() {
if let (Ok(name), Ok(value)) =
(HeaderName::try_from(name), HeaderValue::try_from(value))
{
headers.append(name, value);
}
}
tonic::Status::with_metadata(
tonic::Code::from_i32(status.code as i32),
status.message,
tonic::metadata::MetadataMap::from_headers(headers),
)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn headers(pairs: &[(&str, &str)]) -> HeaderMap {
let mut map = HeaderMap::new();
for (name, value) in pairs {
map.append(
HeaderName::try_from(*name).unwrap(),
HeaderValue::try_from(*value).unwrap(),
);
}
map
}
#[test]
fn every_grpc_timeout_unit_is_understood() {
for (raw, expected) in [
("100n", Duration::from_nanos(100)),
("100u", Duration::from_micros(100)),
("100m", Duration::from_millis(100)),
("2S", Duration::from_secs(2)),
("2M", Duration::from_secs(120)),
("2H", Duration::from_secs(7200)),
] {
assert_eq!(grpc_timeout(&headers(&[("grpc-timeout", raw)])), Some(expected), "{raw}");
}
}
#[test]
fn a_missing_or_malformed_timeout_leaves_the_call_unbounded() {
assert_eq!(grpc_timeout(&headers(&[])), None);
for raw in ["", "m", "100", "100x", "-1S", "abcS"] {
assert_eq!(grpc_timeout(&headers(&[("grpc-timeout", raw)])), None, "{raw:?}");
}
}
#[test]
fn tonics_own_encoding_round_trips() {
assert_eq!(
grpc_timeout(&headers(&[("grpc-timeout", "500000u")])),
Some(Duration::from_millis(500))
);
}
#[test]
fn headers_convert_and_binary_metadata_survives_as_base64() {
let mut metadata = tonic::metadata::MetadataMap::new();
metadata.insert("x-request-id", "abc-123".parse().unwrap());
metadata.insert_bin(
"x-trace-bin",
tonic::metadata::MetadataValue::from_bytes(&[0, 1, 250]),
);
let wire = to_wire_headers(&metadata.into_headers()).unwrap();
assert!(wire.contains(&("x-request-id".to_string(), "abc-123".to_string())));
let (_, encoded) = wire.iter().find(|(k, _)| k == "x-trace-bin").expect("the -bin key");
use base64::Engine as _;
assert_eq!(
base64::engine::general_purpose::STANDARD_NO_PAD.decode(encoded).unwrap(),
vec![0, 1, 250]
);
}
#[test]
fn a_non_ascii_metadata_value_is_refused_rather_than_mangled() {
let mut map = HeaderMap::new();
map.append("x-bad", HeaderValue::from_bytes(&[0xff, 0xfe]).unwrap());
let error = to_wire_headers(&map).expect_err("not representable on the wire");
assert_eq!(error.code(), tonic::Code::Internal);
assert!(error.message().contains("x-bad"), "unhelpful message: {}", error.message());
}
#[test]
fn a_status_carries_its_code_message_and_metadata_into_tonic() {
let mut metadata = crate::Metadata::new();
metadata.insert("x-detail", "quota-exhausted");
metadata.insert_bin("x-detail-bin", vec![0, 1, 250]);
let status = tonic::Status::from(crate::Status {
code: crate::Code::FailedPrecondition,
message: "no".into(),
metadata,
});
assert_eq!(status.code(), tonic::Code::FailedPrecondition);
assert_eq!(status.message(), "no");
assert_eq!(status.metadata().get("x-detail").unwrap(), "quota-exhausted");
assert_eq!(
status.metadata().get_bin("x-detail-bin").unwrap().to_bytes().unwrap().as_ref(),
&[0, 1, 250]
);
}
}