From 2c2b7d790420ea4d3b462a7736507c1ab7e3316e Mon Sep 17 00:00:00 2001 From: Lucio Franco Date: Sun, 29 Sep 2019 18:33:28 -0400 Subject: [PATCH] Add auth and lb example --- tonic-examples/Cargo.toml | 17 ++++ tonic-examples/build.rs | 1 + tonic-examples/proto/echo/echo.proto | 43 +++++++++++ tonic-examples/src/authentication/client.rs | 31 ++++++++ tonic-examples/src/authentication/server.rs | 86 +++++++++++++++++++++ tonic-examples/src/load_balance/client.rs | 29 +++++++ tonic-examples/src/load_balance/server.rs | 77 ++++++++++++++++++ tonic/src/transport/channel.rs | 32 ++++++-- tonic/src/transport/endpoint.rs | 22 +++++- 9 files changed, 332 insertions(+), 6 deletions(-) create mode 100644 tonic-examples/proto/echo/echo.proto create mode 100644 tonic-examples/src/authentication/client.rs create mode 100644 tonic-examples/src/authentication/server.rs create mode 100644 tonic-examples/src/load_balance/client.rs create mode 100644 tonic-examples/src/load_balance/server.rs diff --git a/tonic-examples/Cargo.toml b/tonic-examples/Cargo.toml index 453a180..d9da710 100644 --- a/tonic-examples/Cargo.toml +++ b/tonic-examples/Cargo.toml @@ -20,6 +20,22 @@ path = "src/routeguide/server.rs" name = "routeguide-client" path = "src/routeguide/client.rs" +[[bin]] +name = "authentication-client" +path = "src/authentication/client.rs" + +[[bin]] +name = "authentication-server" +path = "src/authentication/server.rs" + +[[bin]] +name = "load-balance-client" +path = "src/load_balance/client.rs" + +[[bin]] +name = "load-balance-server" +path = "src/load_balance/server.rs" + [dependencies] tonic = { path = "../tonic" } futures-preview = { version = "=0.3.0-alpha.18", default-features = false, features = ["alloc"]} @@ -31,6 +47,7 @@ serde_json = "1.0" serde = { version = "1.0", features = ["derive"] } async-stream = "0.1.1" http = "0.1" +tower = "0.3.0-alpha.1a" [build-dependencies] tonic-build = { path = "../tonic-build" } diff --git a/tonic-examples/build.rs b/tonic-examples/build.rs index fcd3a91..4649bee 100644 --- a/tonic-examples/build.rs +++ b/tonic-examples/build.rs @@ -1,4 +1,5 @@ fn main() { tonic_build::compile_protos("proto/helloworld/helloworld.proto").unwrap(); tonic_build::compile_protos("proto/routeguide/route_guide.proto").unwrap(); + tonic_build::compile_protos("proto/echo/echo.proto").unwrap(); } diff --git a/tonic-examples/proto/echo/echo.proto b/tonic-examples/proto/echo/echo.proto new file mode 100644 index 0000000..49adc34 --- /dev/null +++ b/tonic-examples/proto/echo/echo.proto @@ -0,0 +1,43 @@ +/* + * + * Copyright 2018 gRPC authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + + syntax = "proto3"; + + package grpc.examples.echo; + + // EchoRequest is the request for echo. + message EchoRequest { + string message = 1; + } + + // EchoResponse is the response for echo. + message EchoResponse { + string message = 1; + } + + // Echo is the echo service. + service Echo { + // UnaryEcho is unary echo. + rpc UnaryEcho(EchoRequest) returns (EchoResponse) {} + // ServerStreamingEcho is server side streaming. + rpc ServerStreamingEcho(EchoRequest) returns (stream EchoResponse) {} + // ClientStreamingEcho is client side streaming. + rpc ClientStreamingEcho(stream EchoRequest) returns (EchoResponse) {} + // BidirectionalStreamingEcho is bidi streaming. + rpc BidirectionalStreamingEcho(stream EchoRequest) returns (stream EchoResponse) {} + } \ No newline at end of file diff --git a/tonic-examples/src/authentication/client.rs b/tonic-examples/src/authentication/client.rs new file mode 100644 index 0000000..c1320c8 --- /dev/null +++ b/tonic-examples/src/authentication/client.rs @@ -0,0 +1,31 @@ +pub mod pb { + include!(concat!(env!("OUT_DIR"), "/grpc.examples.echo.rs")); +} + +use http::header::HeaderValue; +use pb::{client::EchoClient, EchoRequest}; +use tonic::transport::Channel; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let channel = Channel::from_static("http://[::1]:50051") + .intercept_headers(|headers| { + headers.insert( + "authorization", + HeaderValue::from_static("Bearer some-secret-token"), + ); + }) + .channel(); + + let mut client = EchoClient::new(channel); + + let request = tonic::Request::new(EchoRequest { + message: "hello".into(), + }); + + let response = client.unary_echo(request).await?; + + println!("RESPONSE={:?}", response); + + Ok(()) +} diff --git a/tonic-examples/src/authentication/server.rs b/tonic-examples/src/authentication/server.rs new file mode 100644 index 0000000..a0c78d8 --- /dev/null +++ b/tonic-examples/src/authentication/server.rs @@ -0,0 +1,86 @@ +pub mod pb { + include!(concat!(env!("OUT_DIR"), "/grpc.examples.echo.rs")); +} + +use pb::{EchoRequest, EchoResponse}; +use std::collections::VecDeque; +use tonic::{body::BoxBody, transport::Server, Request, Response, Status, Streaming}; +use tower::Service; + +type EchoResult = Result, Status>; +type Stream = VecDeque>; + +#[derive(Default)] +pub struct EchoServer; + +#[tonic::async_trait] +impl pb::server::Echo for EchoServer { + async fn unary_echo(&self, request: Request) -> EchoResult { + let message = request.into_inner().message; + Ok(Response::new(EchoResponse { message })) + } + + type ServerStreamingEchoStream = Stream; + + async fn server_streaming_echo( + &self, + _: Request, + ) -> EchoResult { + Err(Status::unimplemented("not implemented")) + } + + async fn client_streaming_echo( + &self, + _: Request>, + ) -> EchoResult { + Err(Status::unimplemented("not implemented")) + } + + type BidirectionalStreamingEchoStream = Stream; + + async fn bidirectional_streaming_echo( + &self, + _: Request>, + ) -> EchoResult { + Err(Status::unimplemented("not implemented")) + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let addr = "[::1]:50051".parse().unwrap(); + let server = EchoServer::default(); + + Server::builder() + .interceptor_fn(move |svc, req| { + let auth_header = req.headers().get("authorization").clone(); + + let authed = if let Some(auth_header) = auth_header { + auth_header == "Bearer some-secret-token" + } else { + false + }; + + let fut = svc.call(req); + + async move { + if authed { + fut.await + } else { + // Cancel the inner future since we never await it + // the IO never gets registered. + drop(fut); + let res = http::Response::builder() + .header("grpc-status", "16") + .body(BoxBody::empty()) + .unwrap(); + Ok(res) + } + } + }) + .clone() + .serve(addr, pb::server::EchoServer::new(server)) + .await?; + + Ok(()) +} diff --git a/tonic-examples/src/load_balance/client.rs b/tonic-examples/src/load_balance/client.rs new file mode 100644 index 0000000..3ab85d7 --- /dev/null +++ b/tonic-examples/src/load_balance/client.rs @@ -0,0 +1,29 @@ +pub mod pb { + include!(concat!(env!("OUT_DIR"), "/grpc.examples.echo.rs")); +} + +use pb::{client::EchoClient, EchoRequest}; +use tonic::transport::Channel; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let endpoints = ["http://[::1]:50051", "http://[::1]:50052"] + .into_iter() + .map(|a| Channel::from_static(a)); + + let channel = Channel::balance_list(endpoints); + + let mut client = EchoClient::new(channel); + + for _ in 0..12 { + let request = tonic::Request::new(EchoRequest { + message: "hello".into(), + }); + + let response = client.unary_echo(request).await?; + + println!("RESPONSE={:?}", response); + } + + Ok(()) +} diff --git a/tonic-examples/src/load_balance/server.rs b/tonic-examples/src/load_balance/server.rs new file mode 100644 index 0000000..b5f8209 --- /dev/null +++ b/tonic-examples/src/load_balance/server.rs @@ -0,0 +1,77 @@ +pub mod pb { + include!(concat!(env!("OUT_DIR"), "/grpc.examples.echo.rs")); +} + +use pb::{EchoRequest, EchoResponse}; +use std::{collections::VecDeque, net::SocketAddr}; +use tokio::sync::mpsc; +use tonic::{transport::Server, Request, Response, Status, Streaming}; + +type EchoResult = Result, Status>; +type Stream = VecDeque>; + +#[derive(Debug)] +pub struct EchoServer { + addr: SocketAddr, +} + +#[tonic::async_trait] +impl pb::server::Echo for EchoServer { + async fn unary_echo(&self, request: Request) -> EchoResult { + let message = format!("{} (from {})", request.into_inner().message, self.addr); + + Ok(Response::new(EchoResponse { message })) + } + + type ServerStreamingEchoStream = Stream; + + async fn server_streaming_echo( + &self, + _: Request, + ) -> EchoResult { + Err(Status::unimplemented("not implemented")) + } + + async fn client_streaming_echo( + &self, + _: Request>, + ) -> EchoResult { + Err(Status::unimplemented("not implemented")) + } + + type BidirectionalStreamingEchoStream = Stream; + + async fn bidirectional_streaming_echo( + &self, + _: Request>, + ) -> EchoResult { + Err(Status::unimplemented("not implemented")) + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let addrs = ["[::1]:50051", "[::1]:50052"]; + + let (tx, mut rx) = mpsc::unbounded_channel(); + + for addr in &addrs { + let addr = addr.parse()?; + let mut tx = tx.clone(); + + let server = EchoServer { addr }; + let serve = Server::builder().serve(addr, pb::server::EchoServer::new(server)); + + tokio::spawn(async move { + if let Err(e) = serve.await { + eprintln!("Error = {:?}", e); + } + + tx.try_send(()).unwrap(); + }); + } + + rx.recv().await; + + Ok(()) +} diff --git a/tonic/src/transport/channel.rs b/tonic/src/transport/channel.rs index ffa2e08..72ea7a4 100644 --- a/tonic/src/transport/channel.rs +++ b/tonic/src/transport/channel.rs @@ -14,6 +14,7 @@ use std::{ fmt, future::Future, pin::Pin, + sync::Arc, task::{Context, Poll}, }; use tower::{ @@ -35,6 +36,7 @@ const DEFAULT_BUFFER_SIZE: usize = 1024; #[derive(Clone)] pub struct Channel { svc: Buffer>, + interceptor_headers: Option>, } /// A future that resolves to an HTTP response. @@ -85,22 +87,35 @@ impl Channel { .and_then(|e| e.buffer_size) .unwrap_or(DEFAULT_BUFFER_SIZE); + let interceptor_headers = list + .iter() + .next() + .and_then(|e| e.interceptor_headers.clone()); + let discover = ServiceList::new(list); - Self::balance(discover, buffer_size) + Self::balance(discover, buffer_size, interceptor_headers) } pub(crate) fn connect(endpoint: Endpoint) -> Self { let buffer_size = endpoint.buffer_size.clone().unwrap_or(DEFAULT_BUFFER_SIZE); + let interceptor_headers = endpoint.interceptor_headers.clone(); let svc = Connection::new(endpoint); let svc = Buffer::new(Either::A(svc), buffer_size); - Channel { svc } + Channel { + svc, + interceptor_headers, + } } - pub(crate) fn balance(discover: D, buffer_size: usize) -> Self + pub(crate) fn balance( + discover: D, + buffer_size: usize, + interceptor_headers: Option>, + ) -> Self where D: Discover + Unpin + Send + 'static, D::Error: Into, @@ -111,7 +126,10 @@ impl Channel { let svc = BoxService::new(svc); let svc = Buffer::new(Either::B(svc), buffer_size); - Channel { svc } + Channel { + svc, + interceptor_headers, + } } } @@ -125,7 +143,11 @@ impl GrpcService for Channel { .map_err(|e| super::Error::from_source(super::ErrorKind::Client, e)) } - fn call(&mut self, request: Request) -> Self::Future { + fn call(&mut self, mut request: Request) -> Self::Future { + if let Some(interceptor) = self.interceptor_headers.clone() { + interceptor(request.headers_mut()); + } + let inner = GrpcService::call(&mut self.svc, request); ResponseFuture { inner } } diff --git a/tonic/src/transport/endpoint.rs b/tonic/src/transport/endpoint.rs index 3dd18ac..7d2a0f8 100644 --- a/tonic/src/transport/endpoint.rs +++ b/tonic/src/transport/endpoint.rs @@ -5,13 +5,15 @@ use bytes::Bytes; use http::uri::{InvalidUriBytes, Uri}; use std::{ convert::{TryFrom, TryInto}, + fmt, + sync::Arc, time::Duration, }; /// Channel builder. /// /// This struct is used to build and configure HTTP/2 channels. -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct Endpoint { pub(super) uri: Uri, pub(super) timeout: Option, @@ -20,6 +22,8 @@ pub struct Endpoint { #[cfg(feature = "tls")] pub(super) tls: Option, pub(super) buffer_size: Option, + pub(super) interceptor_headers: + Option>, } impl Endpoint { @@ -141,6 +145,15 @@ impl Endpoint { self } + /// Intercept outbound HTTP Request headers; + pub fn intercept_headers(&mut self, f: F) -> &mut Self + where + F: Fn(&mut http::HeaderMap) + Send + Sync + 'static, + { + self.interceptor_headers = Some(Arc::new(f)); + self + } + /// Create a channel from this config. pub fn channel(&self) -> Channel { Channel::connect(self.clone()) @@ -157,6 +170,7 @@ impl From for Endpoint { #[cfg(feature = "tls")] tls: None, buffer_size: None, + interceptor_headers: None, } } } @@ -195,3 +209,9 @@ impl std::fmt::Display for Never { } impl std::error::Error for Never {} + +impl fmt::Debug for Endpoint { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Endpoint").finish() + } +}