From 2ebad2d7787e3971cb2929364802ef136d8b443c Mon Sep 17 00:00:00 2001 From: Lucio Franco Date: Sat, 17 Aug 2019 23:31:32 -0400 Subject: [PATCH] Add client, client codegen, helloworld and routeguide client examples --- tonic-examples/Cargo.toml | 11 +- tonic-examples/src/helloworld/client.rs | 32 +++ tonic-examples/src/routeguide/client.rs | 66 +++++ tonic-examples/src/routeguide/server.rs | 27 +- tonic-macros/src/client.rs | 92 +++++++ tonic-macros/src/lib.rs | 32 +++ tonic-macros/src/service.rs | 2 +- tonic/src/body.rs | 61 ++++- tonic/src/client/grpc.rs | 157 ++++++++++++ tonic/src/client/mod.rs | 3 + tonic/src/codec.rs | 314 ------------------------ tonic/src/codec/decode.rs | 155 ++++++++++++ tonic/src/codec/encode.rs | 40 +++ tonic/src/codec/mod.rs | 23 ++ tonic/src/codec/prost.rs | 76 ++++++ tonic/src/error.rs | 2 +- tonic/src/lib.rs | 7 +- tonic/src/server/grpc.rs | 6 +- tonic/src/service.rs | 37 +++ tonic/tests/h2.rs | 24 +- tonic/tests/server.rs | 25 +- tower-h2/src/add_origin.rs | 48 ++++ tower-h2/src/client.rs | 2 +- tower-h2/src/flush.rs | 2 +- tower-h2/src/lib.rs | 2 + tower-h2/src/recv_body.rs | 1 - 26 files changed, 883 insertions(+), 364 deletions(-) create mode 100644 tonic-examples/src/helloworld/client.rs create mode 100644 tonic-examples/src/routeguide/client.rs create mode 100644 tonic-macros/src/client.rs create mode 100644 tonic/src/client/grpc.rs create mode 100644 tonic/src/client/mod.rs delete mode 100644 tonic/src/codec.rs create mode 100644 tonic/src/codec/decode.rs create mode 100644 tonic/src/codec/encode.rs create mode 100644 tonic/src/codec/mod.rs create mode 100644 tonic/src/codec/prost.rs create mode 100644 tonic/src/service.rs create mode 100644 tower-h2/src/add_origin.rs diff --git a/tonic-examples/Cargo.toml b/tonic-examples/Cargo.toml index 92eddf1..778ddce 100644 --- a/tonic-examples/Cargo.toml +++ b/tonic-examples/Cargo.toml @@ -10,14 +10,18 @@ edition = "2018" name = "helloworld-server" path = "src/helloworld/server.rs" -# [[bin]] -# name = "helloworld-client" -# path = "src/helloworld/client.rs" +[[bin]] +name = "helloworld-client" +path = "src/helloworld/client.rs" [[bin]] name = "routeguide-server" path = "src/routeguide/server.rs" +[[bin]] +name = "routeguide-client" +path = "src/routeguide/client.rs" + [dependencies] tonic = { path = "../tonic" } tower-h2 = { path = "../tower-h2" } @@ -29,6 +33,7 @@ bytes = "0.4" serde_json = "1.0" serde = { version = "1.0", features = ["derive"] } async-stream = "0.1" +http = "0.1" [build-dependencies] tonic-build = { path = "../tonic-build" } diff --git a/tonic-examples/src/helloworld/client.rs b/tonic-examples/src/helloworld/client.rs new file mode 100644 index 0000000..14ac4d9 --- /dev/null +++ b/tonic-examples/src/helloworld/client.rs @@ -0,0 +1,32 @@ +#![feature(async_await)] + +use tokio::net::TcpStream; +use tower_h2::{add_origin::AddOrigin, Connection}; + +pub mod hello_world { + include!(concat!(env!("OUT_DIR"), "/helloworld.rs")); + tonic::client!(service = "helloworld.Greeter", proto = "self"); +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let addr = "[::1]:50051".parse()?; + let io = TcpStream::connect(&addr).await?; + + let origin = http::Uri::from_shared(format!("http://{}", addr).into()).unwrap(); + + let svc = Connection::handshake(io).await?; + let svc = AddOrigin::new(svc, origin); + + let mut client = hello_world::GreeterClient::new(svc); + + let request = tonic::Request::new(hello_world::HelloRequest { + name: "hello".into(), + }); + + let response = client.say_hello(request).await?; + + println!("RESPONSE={:?}", response); + + Ok(()) +} diff --git a/tonic-examples/src/routeguide/client.rs b/tonic-examples/src/routeguide/client.rs new file mode 100644 index 0000000..25c2800 --- /dev/null +++ b/tonic-examples/src/routeguide/client.rs @@ -0,0 +1,66 @@ +#![feature(async_await)] + +use route_guide::{Point, RouteNote}; +use std::time::{Duration, Instant}; +use tokio::{net::TcpStream, timer::Interval}; +use tonic::Request; +use tower_h2::{add_origin::AddOrigin, Connection}; +use futures::TryStreamExt; + +mod route_guide { + include!(concat!(env!("OUT_DIR"), "/routeguide.rs")); + tonic::client!(service = "routeguide.RouteGuide", proto = "self"); +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let addr = "[::1]:10000".parse()?; + let io = TcpStream::connect(&addr).await?; + + let origin = http::Uri::from_shared(format!("http://{}", addr).into()).unwrap(); + + let svc = Connection::handshake(io).await?; + let svc = AddOrigin::new(svc, origin); + + let mut client = route_guide::RouteGuideClient::new(svc); + + let start = Instant::now(); + + let response = client + .get_feature(Request::new(Point { + latitude: 409146138, + longitude: -746188906, + })) + .await?; + + println!("FEATURE = {:?}", response); + + let outbound = async_stream::try_stream! { + let mut interval = Interval::new_interval(Duration::from_secs(1)); + + while let Some(time) = interval.next().await { + let elapsed = time.duration_since(start); + let note = RouteNote { + location: Some(Point { + latitude: 409146138 + elapsed.as_secs() as i32, + longitude: -746188906, + }), + message: format!("at {:?}", elapsed), + }; + + yield note; + } + }; + + let request = Request::new(outbound); + + let response = client.route_chat(request).await?; + + let mut inbound = response.into_inner(); + + while let Some(note) = inbound.try_next().await? { + println!("NOTE = {:?}", note); + } + + Ok(()) +} diff --git a/tonic-examples/src/routeguide/server.rs b/tonic-examples/src/routeguide/server.rs index 6abc11e..94209bf 100644 --- a/tonic-examples/src/routeguide/server.rs +++ b/tonic-examples/src/routeguide/server.rs @@ -3,19 +3,22 @@ mod data; use futures::{Stream, StreamExt}; -use tokio::{net::TcpListener, sync::{mpsc, Lock}}; -use tonic::{Request, Response, Status}; -use tower_h2::Server; -use std::sync::Arc; use std::collections::HashMap; use std::hash::{Hash, Hasher}; +use std::sync::Arc; use std::time::Instant; +use tokio::{ + net::TcpListener, + sync::{mpsc, Lock}, +}; +use tonic::{Request, Response, Status}; +use tower_h2::Server; pub mod routeguide { include!(concat!(env!("OUT_DIR"), "/routeguide.rs")); } -use routeguide::{Point, Rectangle, Feature, RouteNote, RouteSummary}; +use routeguide::{Feature, Point, Rectangle, RouteNote, RouteSummary}; #[derive(Debug)] pub struct RouteGuide { @@ -51,7 +54,7 @@ impl RouteGuide { &self, request: Request, ) -> Result>>, Status> { - use std::thread; + use std::thread; println!("ListFeatures = {:?}", request); @@ -70,22 +73,20 @@ impl RouteGuide { println!(" /// done sending"); }); - Ok(Response::new(rx)) } pub async fn record_route( &self, request: Request>>, - ) -> Result, Status> - { + ) -> Result, Status> { println!("RecordRoute"); - + let stream = request.into_inner(); - // Pin the inbound stream to the stack so that we can call next on it + // Pin the inbound stream to the stack so that we can call next on it futures::pin_mut!(stream); - + let mut summary = RouteSummary::default(); let mut last_point = None; let now = Instant::now(); @@ -127,8 +128,6 @@ impl RouteGuide { let stream = request.into_inner(); let mut state = self.state.clone(); - - let output = async_stream::try_stream! { futures::pin_mut!(stream); diff --git a/tonic-macros/src/client.rs b/tonic-macros/src/client.rs new file mode 100644 index 0000000..5d302f6 --- /dev/null +++ b/tonic-macros/src/client.rs @@ -0,0 +1,92 @@ +use super::{Method, Service}; +use proc_macro2::TokenStream; +use quote::{format_ident, quote}; +use syn::Path; + +pub(crate) fn generate(service: Service, proto: String) -> TokenStream { + let mut stream = TokenStream::new(); + + for method in &service.methods { + let path = format!( + "/{}.{}/{}", + service.package, service.proto_name, method.proto_name + ); + + let method = match (method.client_streaming, method.server_streaming) { + (false, false) => generate_unary(method, &proto, path), + (false, true) => generate_server_streaming(method, &proto, path), + (true, false) => generate_client_streaming(method, &proto, path), + (true, true) => generate_streaming(method, &proto, path), + }; + + stream.extend(method); + } + + stream +} + +fn generate_unary(method: &Method, proto: &str, path: String) -> TokenStream { + let ident = format_ident!("{}", method.name); + let request: Path = syn::parse_str(&format!("{}::{}", proto, method.input_type)).unwrap(); + let response: Path = syn::parse_str(&format!("{}::{}", proto, method.output_type)).unwrap(); + + quote! { + pub async fn #ident (&mut self, request: tonic::Request<#request>) + -> Result, tonic::Status> { + let codec = tonic::codec::ProstCodec::new(); + let path = http::uri::PathAndQuery::from_static(#path); + self.inner.unary(request, path, codec).await + } + } +} + +fn generate_server_streaming(method: &Method, proto: &str, path: String) -> TokenStream { + let ident = format_ident!("{}", method.name); + let request: Path = syn::parse_str(&format!("{}::{}", proto, method.input_type)).unwrap(); + let response: Path = syn::parse_str(&format!("{}::{}", proto, method.output_type)).unwrap(); + + quote! { + pub async fn #ident (&mut self, request: tonic::Request<#request>) + -> Result>, tonic::Status> { + let codec = tonic::codec::ProstCodec::new(); + let path = http::uri::PathAndQuery::from_static(#path); + self.inner.server_streaming(request, path, codec).await + } + } +} + +fn generate_client_streaming(method: &Method, proto: &str, path: String) -> TokenStream { + let ident = format_ident!("{}", method.name); + let request: Path = syn::parse_str(&format!("{}::{}", proto, method.input_type)).unwrap(); + let response: Path = syn::parse_str(&format!("{}::{}", proto, method.output_type)).unwrap(); + + quote! { + pub async fn #ident (&mut self, request: tonic::Request) + -> Result, tonic::Status> + where S: tonic::_codegen::Stream> + Send + 'static, + { + let codec = tonic::codec::ProstCodec::new(); + let path = http::uri::PathAndQuery::from_static(#path); + let request = request.map(|s| Box::pin(s)); + self.inner.client_streaming(request, path, codec).await + } + } +} + +fn generate_streaming(method: &Method, proto: &str, path: String) -> TokenStream { + let ident = format_ident!("{}", method.name); + let request: Path = syn::parse_str(&format!("{}::{}", proto, method.input_type)).unwrap(); + let response: Path = syn::parse_str(&format!("{}::{}", proto, method.output_type)).unwrap(); + + quote! { + pub async fn #ident (&mut self, request: tonic::Request) + -> Result>, tonic::Status> + where S: tonic::_codegen::Stream> + Send + 'static, + { + let codec = tonic::codec::ProstCodec::new(); + let path = http::uri::PathAndQuery::from_static(#path); + let request = request.map(|s| Box::pin(s)); + self.inner.streaming(request, path, codec).await + } + } +} diff --git a/tonic-macros/src/lib.rs b/tonic-macros/src/lib.rs index eb65f52..79e1610 100644 --- a/tonic-macros/src/lib.rs +++ b/tonic-macros/src/lib.rs @@ -3,12 +3,44 @@ extern crate proc_macro; +mod client; mod service; use proc_macro::TokenStream; +use quote::quote; use serde::Deserialize; use syn::{AttributeArgs, ItemImpl}; +#[proc_macro] +pub fn client(attr: TokenStream) -> TokenStream { + let args = syn::parse_macro_input!(attr as AttributeArgs); + let (service, proto_path) = load_service(args); + + let service_ident = quote::format_ident!("{}Client", service.name); + let methods = client::generate(service, proto_path); + + let output = quote! { + pub struct #service_ident { + inner: tonic::client::Grpc, + } + + impl #service_ident + where T: tonic::GrpcService, + T::ResponseBody: tonic::body::Body + tonic::_codegen::HttpBody + Send + 'static, + ::Error: Into + Send, + ::Data: Send, { + pub fn new(inner: T) -> Self { + let inner = tonic::client::Grpc::new(inner); + Self { inner } + } + + #methods + } + }; + + TokenStream::from(output) +} + #[proc_macro_attribute] pub fn server(attr: TokenStream, item: TokenStream) -> TokenStream { let mut original = item.clone(); diff --git a/tonic-macros/src/service.rs b/tonic-macros/src/service.rs index c74f535..32d1477 100644 --- a/tonic-macros/src/service.rs +++ b/tonic-macros/src/service.rs @@ -272,7 +272,7 @@ fn generate_client_streaming( struct #service_ident(pub std::sync::Arc<#service_impl>); impl tonic::server::ClientStreamingService for #service_ident - where S: Stream> + Unpin + Send + 'static { + where S: tonic::_codegen::Stream> + Unpin + Send + 'static { type Response = #response; type Future = BoxFuture, tonic::Status>; diff --git a/tonic/src/body.rs b/tonic/src/body.rs index c27c897..53d824e 100644 --- a/tonic/src/body.rs +++ b/tonic/src/body.rs @@ -1,14 +1,63 @@ -use crate::{Code, Status}; -use bytes::{Bytes, IntoBuf}; +use crate::{Code, Error, Status}; +use bytes::{Buf, Bytes, IntoBuf}; use futures_core::{Stream, TryStream}; use futures_util::{ready, TryStreamExt}; use http::HeaderMap; -use http_body::Body; +use http_body::Body as HttpBody; use std::pin::Pin; use std::task::{Context, Poll}; pub type BytesBuf = ::Buf; +pub trait Body: sealed::Sealed { + type Data: Buf; + type Error: Into; + + fn is_end_stream(&self) -> bool; + + fn poll_data(&mut self, cx: &mut Context<'_>) -> Poll>>; + + fn poll_trailers( + &mut self, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>; +} + +impl Body for T +where + T: HttpBody, + T::Error: Into, +{ + type Data = T::Data; + type Error = T::Error; + + fn is_end_stream(&self) -> bool { + HttpBody::is_end_stream(self) + } + + fn poll_data(&mut self, cx: &mut Context<'_>) -> Poll>> { + HttpBody::poll_data(self, cx) + } + + fn poll_trailers( + &mut self, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>> { + HttpBody::poll_trailers(self, cx) + } +} + +impl sealed::Sealed for T +where + T: HttpBody, + T::Error: Into, +{ +} + +mod sealed { + pub trait Sealed {} +} + pub struct BoxBody { inner: Box + Send>, } @@ -25,7 +74,7 @@ impl BoxBody { } } -impl Body for BoxBody { +impl HttpBody for BoxBody { type Data = BytesBuf; type Error = Status; @@ -72,7 +121,7 @@ impl BoxAsyncBody { } } -impl Body for BoxAsyncBody { +impl HttpBody for BoxAsyncBody { type Data = BytesBuf; type Error = Status; @@ -114,7 +163,7 @@ where } } -impl Body for AsyncBody +impl HttpBody for AsyncBody where S: Stream> + Unpin, { diff --git a/tonic/src/client/grpc.rs b/tonic/src/client/grpc.rs new file mode 100644 index 0000000..98d9152 --- /dev/null +++ b/tonic/src/client/grpc.rs @@ -0,0 +1,157 @@ +use crate::{ + body::{Body, BoxAsyncBody}, + codec::{decode, encode, Codec, Streaming}, + Code, GrpcService, Request, Response, Status, +}; +use futures_core::Stream; +use futures_util::{future, stream, TryStreamExt}; +use http::{ + header::{HeaderValue, CONTENT_TYPE, TE}, + uri::{Parts, PathAndQuery, Uri}, +}; +use http_body::Body as HttpBody; + +pub struct Grpc { + inner: T, +} + +impl Grpc { + pub fn new(inner: T) -> Self { + Self { inner } + } + + pub async fn unary( + &mut self, + request: Request, + path: PathAndQuery, + codec: C, + ) -> Result, Status> + where + T: GrpcService, + T::ResponseBody: Body + HttpBody + Send + 'static, + ::Error: Into + Send, + ::Data: Send, + C: Codec, + C::Encoder: Send + 'static, + C::Decoder: Send + 'static, + M1: Send + 'static, + M2: Send + Unpin + 'static, + { + let request = request.map(|m| stream::once(future::ok(m))); + self.client_streaming(request, path, codec).await + } + + pub async fn client_streaming( + &mut self, + request: Request, + path: PathAndQuery, + codec: C, + ) -> Result, Status> + where + T: GrpcService, + T::ResponseBody: Body + HttpBody + Send + 'static, + ::Error: Into + Send, + ::Data: Send, + S: Stream> + Send + 'static, + C: Codec, + C::Encoder: Send + 'static, + C::Decoder: Send + 'static, + M1: Send, + M2: Send + Unpin + 'static, + { + let response = self.streaming(request, path, codec).await?; + + // TODO: use response to parts + let mut body = response.into_inner(); + let message = body + .try_next() + .await? + .ok_or(Status::new(Code::Internal, "Missing response message."))?; + + Ok(Response::new(message)) + } + + pub async fn server_streaming( + &mut self, + request: Request, + path: PathAndQuery, + codec: C, + ) -> Result>, Status> + where + T: GrpcService, + T::ResponseBody: Body + HttpBody + Send + 'static, + ::Error: Into + Send, + ::Data: Send, + C: Codec, + C::Encoder: Send + 'static, + C::Decoder: Send + 'static, + M1: Send + 'static, + M2: Send + Unpin + 'static, + { + let request = request.map(|m| stream::once(future::ok(m))); + self.streaming(request, path, codec).await + } + + pub async fn streaming( + &mut self, + request: Request, + path: PathAndQuery, + mut codec: C, + ) -> Result>, Status> + where + T: GrpcService, + T::ResponseBody: Body + HttpBody + Send + 'static, + ::Error: Into + Send, + ::Data: Send, + S: Stream> + Send + 'static, + C: Codec, + C::Encoder: Send + 'static, + C::Decoder: Send + 'static, + M1: Send, + M2: Send + Unpin + 'static, + { + let mut parts = Parts::default(); + parts.path_and_query = Some(path); + + let uri = Uri::from_parts(parts).expect("path_and_query only is valid Uri"); + + let request = request + .map(|s| encode(codec.encoder(), Box::pin(s))) + .map(BoxAsyncBody::new_try); + + let mut request = request.into_http(uri); + + // Add the gRPC related HTTP headers + request + .headers_mut() + .insert(TE, HeaderValue::from_static("trailers")); + + // Set the content type + // TODO: Don't hard code this here + let content_type = ::CONTENT_TYPE; + request + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static(content_type)); + + let response = self + .inner + .call(request) + .await + .map_err(|err| Status::from_error(&*(err.into())))?; + + let status_code = response.status(); + let trailers_only_status = Status::from_header_map(response.headers()); + + if let Some(status) = trailers_only_status { + if status.code() != Code::Ok { + return Err(status); + } + } + + let response = response + .map(|b| decode(codec.decoder(), b).into_stream()) + .map(Streaming::new); + + Ok(Response::from_http(response)) + } +} diff --git a/tonic/src/client/mod.rs b/tonic/src/client/mod.rs new file mode 100644 index 0000000..d9a4942 --- /dev/null +++ b/tonic/src/client/mod.rs @@ -0,0 +1,3 @@ +mod grpc; + +pub use self::grpc::Grpc; diff --git a/tonic/src/codec.rs b/tonic/src/codec.rs deleted file mode 100644 index 8ca3d83..0000000 --- a/tonic/src/codec.rs +++ /dev/null @@ -1,314 +0,0 @@ -use crate::{body::BytesBuf, Code, Status}; -use async_stream::stream; -use bytes::{Buf, BufMut, BytesMut, IntoBuf}; -use futures_core::{Stream, TryStream}; -use futures_util::{future, StreamExt}; -use http_body::Body; -use prost::Message; -use std::marker::PhantomData; -use std::pin::Pin; -use tokio_codec::{Decoder, Encoder}; -use tracing::{debug, trace}; - -pub trait Codec { - type Encode; - type Decode; - - type Encoder: Encoder; - type Decoder: Decoder; - - const CONTENT_TYPE: &'static str; - - fn encoder(&mut self) -> Self::Encoder; - fn decoder(&mut self) -> Self::Decoder; -} - -pub struct Streaming { - inner: Pin> + Send + 'static>>, -} - -impl Streaming { - pub fn new(inner: impl Stream> + Send + 'static) -> Self { - let inner = Box::pin(inner); - Self { inner } - } -} - -use std::task::{Context, Poll}; -impl Stream for Streaming { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.inner).poll_next(cx) - } -} - -pub fn encode(mut encoder: T, mut source: U) -> impl TryStream -where - T: Encoder, - U: Stream> + Unpin, -{ - stream! { - let mut buf = BytesMut::with_capacity(1024); - - loop { - match source.next().await { - Some(Ok(item)) => { - buf.reserve(5); - unsafe { - buf.advance_mut(5); - } - encoder.encode(item, &mut buf).map_err(drop).unwrap(); - - // now that we know length, we can write the header - let len = buf.len() - 5; - assert!(len <= ::std::u32::MAX as usize); - { - let mut cursor = ::std::io::Cursor::new(&mut buf[..5]); - cursor.put_u8(0); // byte must be 0, reserve doesn't auto-zero - cursor.put_u32_be(len as u32); - } - - yield Ok(buf.split_to(len + 5).freeze().into_buf()); - }, - Some(Err(status)) => yield Err(status), - None => break, - } - } - } -} - -pub fn decode( - mut decoder: T, - mut source: B, -) -> impl TryStream + 'static -where - T: Decoder + 'static, - T::Item: Unpin + 'static, - B: Body + 'static, - B::Error: Into, -{ - stream! { - let mut buf = BytesMut::with_capacity(1024); - let mut state = State::ReadHeader; - - loop { - // TODO: use try_stream! and ? - if let Some(item) = decode_chunk(&mut decoder, &mut buf, &mut state).unwrap() { - yield Ok(item); - } - - let chunk = match future::poll_fn(|cx| source.poll_data(cx)).await { - Some(Ok(d)) => Some(d), - Some(Err(e)) => { - let err = e.into(); - debug!("decoder inner stream error: {:?}", err); - let status = Status::from_error(&*err); - yield Err(status); - break; - }, - None => None, - }; - - if let Some(data)= chunk { - buf.put(data); - } else { - if buf.has_remaining_mut() { - trace!("unexpected EOF decoding stream"); - yield Err(Status::new( - Code::Internal, - "Unexpected EOF decoding stream.".to_string(), - )); - } else { - break; - } - } - - // TODO: poll_trailers for Response status code - } - } -} - -fn decode_chunk( - decoder: &mut T, - buf1: &mut BytesMut, - state: &mut State, -) -> Result, Status> -where - T: Decoder, -{ - let mut buf = (&buf1[..]).into_buf(); - - if let State::ReadHeader = state { - println!("reading header"); - if buf.remaining() < 5 { - return Ok(None); - } - - let is_compressed = match buf.get_u8() { - 0 => false, - 1 => { - trace!("message compressed, compression not supported yet"); - return Err(crate::Status::new( - crate::Code::Unimplemented, - "Message compressed, compression not supported yet.".to_string(), - )); - } - f => { - trace!("unexpected compression flag"); - return Err(crate::Status::new( - crate::Code::Internal, - format!("Unexpected compression flag: {}", f), - )); - } - }; - let len = buf.get_u32_be() as usize; - - *state = State::ReadBody { - compression: is_compressed, - len, - } - } - - if let State::ReadBody { len, .. } = state { - println!("reading body"); - if buf.remaining() < *len { - return Ok(None); - } - - // advance past the header - buf1.advance(5); - - match decoder.decode(buf1) { - Ok(Some(msg)) => { - *state = State::ReadHeader; - return Ok(Some(msg)); - } - Ok(None) => return Ok(None), - Err(e) => { - return Err(e); - } - } - } - - Ok(None) -} - -#[derive(Debug, Clone)] -pub struct ProstCodec { - _pd: PhantomData<(T, U)>, -} - -impl ProstCodec { - pub fn new() -> Self { - Self { _pd: PhantomData } - } -} - -impl Codec for ProstCodec -where - T: Message, - U: Message + Default, -{ - type Encode = T; - type Decode = U; - - type Encoder = ProstEncoder; - type Decoder = ProstDecoder; - - const CONTENT_TYPE: &'static str = "application/groc+proto"; - - fn encoder(&mut self) -> Self::Encoder { - ProstEncoder(PhantomData) - } - - fn decoder(&mut self) -> Self::Decoder { - ProstDecoder(PhantomData) - } -} - -pub struct ProstEncoder(PhantomData); - -impl Encoder for ProstEncoder { - type Item = T; - type Error = Status; - - fn encode(&mut self, item: Self::Item, buf: &mut BytesMut) -> Result<(), Self::Error> { - let len = item.encoded_len(); - - if buf.remaining_mut() < len { - buf.reserve(len); - } - - item.encode(buf) - .map_err(|_| unreachable!("Message only errors if not enough space")) - } -} - -pub struct ProstDecoder(PhantomData); - -impl Decoder for ProstDecoder { - type Item = U; - type Error = Status; - - fn decode(&mut self, buf: &mut BytesMut) -> Result, Self::Error> { - Message::decode(buf.take()) - .map(Option::Some) - .map_err(from_decode_error) - } -} - -fn from_decode_error(error: prost::DecodeError) -> crate::Status { - // Map Protobuf parse errors to an INTERNAL status code, as per - // https://github.com/grpc/grpc/blob/master/doc/statuscodes.md - Status::new(Code::Internal, error.to_string()) -} - -#[derive(Default)] -pub struct UnitCodec; - -impl Codec for UnitCodec { - type Encode = (); - type Decode = (); - - type Encoder = UnitEncoder; - type Decoder = UnitDecoder; - - const CONTENT_TYPE: &'static str = "()"; - - fn encoder(&mut self) -> Self::Encoder { - UnitEncoder - } - - fn decoder(&mut self) -> Self::Decoder { - UnitDecoder - } -} - -pub struct UnitEncoder; - -impl Encoder for UnitEncoder { - type Item = (); - type Error = crate::Status; - - fn encode(&mut self, _item: Self::Item, _buf: &mut BytesMut) -> Result<(), Self::Error> { - unimplemented!() - } -} - -pub struct UnitDecoder; - -impl Decoder for UnitDecoder { - type Item = (); - type Error = Status; - - fn decode(&mut self, _buf: &mut BytesMut) -> Result, Self::Error> { - Ok(Some(())) - } -} - -#[derive(Debug)] -enum State { - ReadHeader, - ReadBody { compression: bool, len: usize }, -} diff --git a/tonic/src/codec/decode.rs b/tonic/src/codec/decode.rs new file mode 100644 index 0000000..d20db81 --- /dev/null +++ b/tonic/src/codec/decode.rs @@ -0,0 +1,155 @@ +use crate::{Code, Status}; +use bytes::{Buf, BufMut, BytesMut, IntoBuf}; +use futures_core::{Stream, TryStream}; +use futures_util::future; +use http::StatusCode; +use http_body::Body; +use std::pin::Pin; +use tokio_codec::Decoder; +use tracing::{debug, trace}; + +pub struct Streaming { + inner: Pin> + Send + 'static>>, +} + +impl Streaming { + pub fn new(inner: impl Stream> + Send + 'static) -> Self { + let inner = Box::pin(inner); + Self { inner } + } +} + +use std::task::{Context, Poll}; +impl Stream for Streaming { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_next(cx) + } +} + +#[derive(Debug)] +enum State { + ReadHeader, + ReadBody { compression: bool, len: usize }, +} + +enum Direction { + Request, + Response(StatusCode), + EmptyResponse, +} + +pub fn decode( + mut decoder: T, + mut source: B, +) -> impl TryStream + 'static +where + T: Decoder + 'static, + T::Item: Unpin + 'static, + B: Body + 'static, + B::Error: Into, +{ + async_stream::stream! { + let mut buf = BytesMut::with_capacity(1024 * 1024); + let mut state = State::ReadHeader; + + loop { + // TODO: use try_stream! and ? + if let Some(item) = decode_chunk(&mut decoder, &mut buf, &mut state).unwrap() { + yield Ok(item); + } + + let chunk = match future::poll_fn(|cx| source.poll_data(cx)).await { + Some(Ok(d)) => Some(d), + Some(Err(e)) => { + let err = e.into(); + debug!("decoder inner stream error: {:?}", err); + let status = Status::from_error(&*err); + yield Err(status); + break; + }, + None => None, + }; + + if let Some(data) = chunk { + buf.put(data); + } else { + if buf.has_remaining_mut() { + trace!("unexpected EOF decoding stream"); + yield Err(Status::new( + Code::Internal, + "Unexpected EOF decoding stream.".to_string(), + )); + } else { + break; + } + } + + // TODO: poll_trailers for Response status code + } + } +} + +fn decode_chunk( + decoder: &mut T, + buf1: &mut BytesMut, + state: &mut State, +) -> Result, Status> +where + T: Decoder, +{ + let mut buf = (&buf1[..]).into_buf(); + + if let State::ReadHeader = state { + if buf.remaining() < 5 { + return Ok(None); + } + + let is_compressed = match buf.get_u8() { + 0 => false, + 1 => { + trace!("message compressed, compression not supported yet"); + return Err(crate::Status::new( + crate::Code::Unimplemented, + "Message compressed, compression not supported yet.".to_string(), + )); + } + f => { + trace!("unexpected compression flag"); + return Err(crate::Status::new( + crate::Code::Internal, + format!("Unexpected compression flag: {}", f), + )); + } + }; + let len = buf.get_u32_be() as usize; + + *state = State::ReadBody { + compression: is_compressed, + len, + } + } + + if let State::ReadBody { len, .. } = state { + if buf.remaining() < *len { + return Ok(None); + } + + // advance past the header + buf1.advance(5); + + match decoder.decode(buf1) { + Ok(Some(msg)) => { + *state = State::ReadHeader; + return Ok(Some(msg)); + } + Ok(None) => return Ok(None), + Err(e) => { + return Err(e); + } + } + } + + Ok(None) +} diff --git a/tonic/src/codec/encode.rs b/tonic/src/codec/encode.rs new file mode 100644 index 0000000..ff271d9 --- /dev/null +++ b/tonic/src/codec/encode.rs @@ -0,0 +1,40 @@ +use crate::{body::BytesBuf, Status}; +use bytes::{BufMut, BytesMut, IntoBuf}; +use futures_core::{Stream, TryStream}; +use futures_util::StreamExt; +use tokio_codec::Encoder; + +pub fn encode(mut encoder: T, mut source: U) -> impl TryStream +where + T: Encoder, + U: Stream> + Unpin, +{ + async_stream::stream! { + let mut buf = BytesMut::with_capacity(1024); + + loop { + match source.next().await { + Some(Ok(item)) => { + buf.reserve(5); + unsafe { + buf.advance_mut(5); + } + encoder.encode(item, &mut buf).map_err(drop).unwrap(); + + // now that we know length, we can write the header + let len = buf.len() - 5; + assert!(len <= std::u32::MAX as usize); + { + let mut cursor = std::io::Cursor::new(&mut buf[..5]); + cursor.put_u8(0); // byte must be 0, reserve doesn't auto-zero + cursor.put_u32_be(len as u32); + } + + yield Ok(buf.split_to(len + 5).freeze().into_buf()); + }, + Some(Err(status)) => yield Err(status), + None => break, + } + } + } +} diff --git a/tonic/src/codec/mod.rs b/tonic/src/codec/mod.rs new file mode 100644 index 0000000..352bd95 --- /dev/null +++ b/tonic/src/codec/mod.rs @@ -0,0 +1,23 @@ +mod decode; +mod encode; +mod prost; + +pub use self::decode::{decode, Streaming}; +pub use self::encode::encode; +pub use self::prost::ProstCodec; + +use crate::Status; +use tokio_codec::{Decoder, Encoder}; + +pub trait Codec { + type Encode; + type Decode; + + type Encoder: Encoder; + type Decoder: Decoder; + + const CONTENT_TYPE: &'static str; + + fn encoder(&mut self) -> Self::Encoder; + fn decoder(&mut self) -> Self::Decoder; +} diff --git a/tonic/src/codec/prost.rs b/tonic/src/codec/prost.rs new file mode 100644 index 0000000..177a582 --- /dev/null +++ b/tonic/src/codec/prost.rs @@ -0,0 +1,76 @@ +use super::Codec; +use crate::{Code, Status}; +use bytes::{BufMut, BytesMut}; +use prost::Message; +use std::marker::PhantomData; +use tokio_codec::{Decoder, Encoder}; + +#[derive(Debug, Clone)] +pub struct ProstCodec { + _pd: PhantomData<(T, U)>, +} + +impl ProstCodec { + pub fn new() -> Self { + Self { _pd: PhantomData } + } +} + +impl Codec for ProstCodec +where + T: Message, + U: Message + Default, +{ + type Encode = T; + type Decode = U; + + type Encoder = ProstEncoder; + type Decoder = ProstDecoder; + + const CONTENT_TYPE: &'static str = "application/groc+proto"; + + fn encoder(&mut self) -> Self::Encoder { + ProstEncoder(PhantomData) + } + + fn decoder(&mut self) -> Self::Decoder { + ProstDecoder(PhantomData) + } +} + +pub struct ProstEncoder(PhantomData); + +impl Encoder for ProstEncoder { + type Item = T; + type Error = Status; + + fn encode(&mut self, item: Self::Item, buf: &mut BytesMut) -> Result<(), Self::Error> { + let len = item.encoded_len(); + + if buf.remaining_mut() < len { + buf.reserve(len); + } + + item.encode(buf) + .map_err(|_| unreachable!("Message only errors if not enough space")) + } +} + +pub struct ProstDecoder(PhantomData); + +impl Decoder for ProstDecoder { + type Item = U; + type Error = Status; + + fn decode(&mut self, buf: &mut BytesMut) -> Result, Self::Error> { + Message::decode(buf.take()) + .map(Option::Some) + .map_err(from_decode_error) + } +} + +fn from_decode_error(error: prost::DecodeError) -> crate::Status { + // Map Protobuf parse errors to an INTERNAL status code, as per + // https://github.com/grpc/grpc/blob/master/doc/statuscodes.md + Status::new(Code::Internal, error.to_string()) +} diff --git a/tonic/src/error.rs b/tonic/src/error.rs index 9ad0196..ec8f9e5 100644 --- a/tonic/src/error.rs +++ b/tonic/src/error.rs @@ -1,7 +1,7 @@ use std::fmt; #[allow(dead_code)] -pub(crate) type Error = Box; +pub type Error = Box; #[derive(Debug)] #[allow(dead_code)] diff --git a/tonic/src/lib.rs b/tonic/src/lib.rs index eb7078c..12c2908 100644 --- a/tonic/src/lib.rs +++ b/tonic/src/lib.rs @@ -4,6 +4,7 @@ //! gRPC implementation pub mod body; +pub mod client; pub mod codec; #[doc(hidden)] pub mod error; @@ -12,13 +13,15 @@ pub mod server; mod request; mod response; +mod service; mod status; pub use body::{BoxAsyncBody, BoxBody}; pub use request::Request; pub use response::Response; +pub use service::GrpcService; pub use status::{Code, Status}; -pub use tonic_macros::server; +pub use tonic_macros::{client, server}; pub(crate) use error::Error; @@ -35,7 +38,9 @@ pub trait GrpcInnerService { #[doc(hidden)] pub mod _codegen { + pub use futures_core::Stream; pub use futures_util::future::{ok, Ready}; + pub use http_body::Body as HttpBody; pub use std::future::Future; pub use std::pin::Pin; pub use std::task::{Context, Poll}; diff --git a/tonic/src/server/grpc.rs b/tonic/src/server/grpc.rs index f73f6f1..6478cfa 100644 --- a/tonic/src/server/grpc.rs +++ b/tonic/src/server/grpc.rs @@ -83,7 +83,7 @@ where self.map_response(response).map(BoxAsyncBody::new_try) } -//BoxStream, + //BoxStream, pub async fn client_streaming( &mut self, mut service: S, @@ -154,9 +154,7 @@ where B::Error: Into + Send, { Request::from_http( - request.map(|b| { - Streaming::new(decode(self.codec.decoder(), b).into_stream()) - }), + request.map(|b| Streaming::new(decode(self.codec.decoder(), b).into_stream())), ) } diff --git a/tonic/src/service.rs b/tonic/src/service.rs new file mode 100644 index 0000000..aa28015 --- /dev/null +++ b/tonic/src/service.rs @@ -0,0 +1,37 @@ +use crate::body::Body; +use http::{Request, Response}; +use http_body::Body as HttpBody; +use std::future::Future; +use std::task::{Context, Poll}; +use tower_service::Service; + +pub trait GrpcService { + type ResponseBody: Body + HttpBody; + type Error: Into; + + type Future: Future, Self::Error>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll>; + + fn call(&mut self, request: Request) -> Self::Future; +} + +impl GrpcService for T +where + T: Service, Response = Response>, + T::Error: Into, + ResBody: Body + HttpBody, + ::Error: Into, +{ + type ResponseBody = ResBody; + type Error = T::Error; + type Future = T::Future; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + Service::poll_ready(self, cx) + } + + fn call(&mut self, request: Request) -> Self::Future { + Service::call(self, request) + } +} diff --git a/tonic/tests/h2.rs b/tonic/tests/h2.rs index 1894cd2..d118be4 100644 --- a/tonic/tests/h2.rs +++ b/tonic/tests/h2.rs @@ -1,14 +1,14 @@ #![feature(async_await, type_alias_impl_trait)] -use futures_util::future; use futures_core::Stream; -use std::pin::Pin; +use futures_util::future; use std::future::Future; +use std::pin::Pin; use std::task::{Context, Poll}; use tokio::net::TcpListener; use tonic::{ body, - server::{Grpc, UnaryService, ClientStreamingService}, + server::{ClientStreamingService, Grpc, UnaryService}, Request, Response, Status, }; use tower_h2::{RecvBody, Server}; @@ -47,15 +47,20 @@ impl UnaryService for SayHello { struct SayHelloStream; -impl ClientStreamingService for SayHelloStream -where S: Stream> + Unpin + Send + 'static { +impl ClientStreamingService for SayHelloStream +where + S: Stream> + Unpin + Send + 'static, +{ type Response = HelloReply; // type Future = impl Future, Status>>; - type Future = Pin, Status>> + Send + 'static>>; + type Future = + Pin, Status>> + Send + 'static>>; fn call(&mut self, _req: Request) -> Self::Future { - let fut = async move { - Ok(Response::new(HelloReply { message: "hello".into()})) + let fut = async move { + Ok(Response::new(HelloReply { + message: "hello".into(), + })) }; Box::pin(fut) } @@ -116,9 +121,8 @@ impl Service> for Svc { Box::pin(fut) } - _ => unimplemented!() + _ => unimplemented!(), } - } } diff --git a/tonic/tests/server.rs b/tonic/tests/server.rs index 15ffed4..b22554d 100644 --- a/tonic/tests/server.rs +++ b/tonic/tests/server.rs @@ -1,13 +1,13 @@ #![feature(async_await, type_alias_impl_trait)] +use futures_core::Stream; use std::future::Future; +use std::pin::Pin; use std::task::{Context, Poll}; use tokio_buf::BufStream; use tonic::codec::ProstCodec; use tonic::server::*; use tonic::{Request, Response, Status}; -use std::pin::Pin; -use futures_core::Stream; #[derive(Clone, PartialEq, prost::Message)] pub struct HelloRequest { @@ -28,20 +28,31 @@ impl UnaryService for SayHello { type Future = impl Future, Status>>; fn call(&mut self, _request: Request) -> Self::Future { - async move { Ok(Response::new(HelloReply { message: "hello".into()})) } + async move { + Ok(Response::new(HelloReply { + message: "hello".into(), + })) + } } } struct SayHelloStream; -impl ClientStreamingService for SayHelloStream -where S: Stream> + Unpin + Send + 'static { +impl ClientStreamingService for SayHelloStream +where + S: Stream> + Unpin + Send + 'static, +{ type Response = HelloReply; // type Future = impl Future, Status>>; - type Future = Pin, Status>> + Send + 'static>>; + type Future = + Pin, Status>> + Send + 'static>>; fn call(&mut self, _: Request) -> Self::Future { - let fut = async move { Ok(Response::new(HelloReply { message: "hello".into()})) }; + let fut = async move { + Ok(Response::new(HelloReply { + message: "hello".into(), + })) + }; Box::pin(fut) } } diff --git a/tower-h2/src/add_origin.rs b/tower-h2/src/add_origin.rs new file mode 100644 index 0000000..0f3b6dd --- /dev/null +++ b/tower-h2/src/add_origin.rs @@ -0,0 +1,48 @@ +use http::{Request, Uri}; +use std::task::{Context, Poll}; +use tower_service::Service; + +#[derive(Debug)] +pub struct AddOrigin { + inner: T, + origin: Uri, +} + +impl AddOrigin { + pub fn new(inner: T, origin: Uri) -> Self { + Self { inner, origin } + } +} + +impl Service> for AddOrigin +where + T: Service>, +{ + type Response = T::Response; + type Error = T::Error; + type Future = T::Future; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: Request) -> Self::Future { + // Split the request into the head and the body. + let (mut head, body) = req.into_parts(); + + // Split the request URI into parts. + let mut uri: http::uri::Parts = head.uri.into(); + let set_uri = self.origin.clone().into_parts(); + + // Update the URI parts, setting hte scheme and authority + uri.scheme = Some(set_uri.scheme.expect("expected scheme").clone()); + uri.authority = Some(set_uri.authority.expect("expected authority").clone()); + + // Update the the request URI + head.uri = http::Uri::from_parts(uri).expect("valid uri"); + + let request = Request::from_parts(head, body); + + self.inner.call(request) + } +} diff --git a/tower-h2/src/client.rs b/tower-h2/src/client.rs index 0c0f902..e0213be 100644 --- a/tower-h2/src/client.rs +++ b/tower-h2/src/client.rs @@ -1,6 +1,6 @@ use crate::{buf::SendBuf, flush::Flush, recv_body::RecvBody}; use futures_util::{future, FutureExt, TryFutureExt}; -use h2::{client::SendRequest, RecvStream}; +use h2::client::SendRequest; use http::{Request, Response}; use http_body::Body; use std::future::Future; diff --git a/tower-h2/src/flush.rs b/tower-h2/src/flush.rs index e7e8e60..4759430 100644 --- a/tower-h2/src/flush.rs +++ b/tower-h2/src/flush.rs @@ -96,7 +96,7 @@ where match ready!(self.h2.poll_capacity(cx)) { Some(Ok(0)) => {} Some(Ok(_)) => break, - Some(Err(e)) => return panic!("error {:?}", e), + Some(Err(e)) => panic!("error {:?}", e), None => { debug!("connection closed early"); // The error shouldn't really matter at this diff --git a/tower-h2/src/lib.rs b/tower-h2/src/lib.rs index 4a8cad4..cdff648 100644 --- a/tower-h2/src/lib.rs +++ b/tower-h2/src/lib.rs @@ -3,6 +3,8 @@ #[macro_use] extern crate log; +pub mod add_origin; + mod buf; mod client; mod error; diff --git a/tower-h2/src/recv_body.rs b/tower-h2/src/recv_body.rs index 64306a5..231fed4 100644 --- a/tower-h2/src/recv_body.rs +++ b/tower-h2/src/recv_body.rs @@ -1,5 +1,4 @@ use bytes::{Buf, Bytes, BytesMut}; -use futures_core::Stream; use futures_util::TryStreamExt; use http_body::Body; use std::task::{Context, Poll};