talksik/tonic
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: Add gRPC interceptors (#232)

Browse Source 
This change introduces proper gRPC interceptors that are avilable
regardless of the transport used. Each codegen service now produces an
additional method called `with_interceptor` that accepts a
`Interceptor`.

All examples have been updated to use this new style and interop has a
custom `tower::Service` middleware to echo the headers. There is also a
new `interceptor` example that shows basic usage.

BREAKING CHANGE: removed `interceptor_fn` and `intercep_headers_fn` from `transport` in favor of using `tonic::Interceptor`.
This commit is contained in:
Lucio Franco
2020-01-13 21:26:08 -05:00
committed by GitHub
parent 0fa2bf1cea
commit eba7ec7b32
25 changed files with 456 additions and 273 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tonic/src/client/grpc.rs
+22 -1
View File
@@ -2,6 +2,7 @@ use crate::{
body::{Body, BoxBody},
client::GrpcService,
codec::{encode_client, Codec, Streaming},
interceptor::Interceptor,
Code, Request, Response, Status,
};
use futures_core::Stream;
@@ -28,12 +29,25 @@ use std::fmt;
/// [gRPC protocol definition]: https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests
pub struct Grpc<T> {
inner: T,
interceptor: Option<Interceptor>,
}
impl<T> Grpc<T> {
/// Creates a new gRPC client with the provided [`GrpcService`].
pub fn new(inner: T) -> Self {
Self { inner }
Self {
inner,
interceptor: None,
}
}
/// Creates a new gRPC client with the provided [`GrpcService`] and will apply
/// the provided interceptor on each request.
pub fn with_interceptor(inner: T, interceptor: impl Into<Interceptor>) -> Self {
Self {
inner,
interceptor: Some(interceptor.into()),
}
}
/// Check if the inner [`GrpcService`] is able to accept a new request.
@@ -134,6 +148,12 @@ impl<T> Grpc<T> {
M1: Send + Sync + 'static,
M2: Send + Sync + 'static,
{
let request = if let Some(interceptor) = &self.interceptor {
interceptor.call(request)?
} else {
request
};
let mut parts = Parts::default();
parts.path_and_query = Some(path);
@@ -192,6 +212,7 @@ impl<T: Clone> Clone for Grpc<T> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
interceptor: self.interceptor.clone(),
}
}
}
tonic/src/interceptor.rs
+54
View File
@@ -0,0 +1,54 @@
use crate::{Request, Status};
use std::{fmt, sync::Arc};
/// Represents a gRPC interceptor.
///
/// gRPC interceptors are similar to middleware but have much less
/// flexibility. This interceptor allows you to do two main things,
/// one is to add/remove/check items in the `MetadataMap` of each
/// request. Two, cancel a request with any `Status`.
///
/// An interceptor can be used on both the server and client side through
/// the `tonic-build` crate's generated structs.
///
/// These interceptors do not allow you to modify the `Message` of the request
/// but allow you to check for metadata. If you would like to apply middleware like
/// features to the body of the request, going through the `tower` abstraction is recommended.
#[derive(Clone)]
pub struct Interceptor {
f: Arc<dyn Fn(Request<()>) -> Result<Request<()>, Status> + Send + Sync + 'static>,
}
impl Interceptor {
/// Create a new `Interceptor` from the provided function.
pub fn new(
f: impl Fn(Request<()>) -> Result<Request<()>, Status> + Send + Sync + 'static,
) -> Self {
Interceptor { f: Arc::new(f) }
}
pub(crate) fn call<T>(&self, req: Request<T>) -> Result<Request<T>, Status> {
let (metadata, ext, message) = req.into_parts();
let temp_req = Request::from_parts(metadata, ext, ());
let (metadata, ext, _) = (self.f)(temp_req)?.into_parts();
Ok(Request::from_parts(metadata, ext, message))
}
}
impl<F> From<F> for Interceptor
where
F: Fn(Request<()>) -> Result<Request<()>, Status> + Send + Sync + 'static,
{
fn from(f: F) -> Self {
Interceptor::new(f)
}
}
impl fmt::Debug for Interceptor {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Interceptor").finish()
}
}
tonic/src/lib.rs
+2
View File
@@ -86,6 +86,7 @@ pub mod server;
#[cfg_attr(docsrs, doc(cfg(feature = "transport")))]
pub mod transport;
mod interceptor;
mod macros;
mod request;
mod response;
@@ -98,6 +99,7 @@ pub use async_trait::async_trait;
#[doc(inline)]
pub use codec::Streaming;
pub use interceptor::Interceptor;
pub use request::{IntoRequest, IntoStreamingRequest, Request};
pub use response::Response;
pub use status::{Code, Status};
tonic/src/request.rs
+12
View File
@@ -145,6 +145,18 @@ impl<T> Request<T> {
self.message
}
pub(crate) fn into_parts(self) -> (MetadataMap, Extensions, T) {
(self.metadata, self.extensions, self.message)
}
pub(crate) fn from_parts(metadata: MetadataMap, extensions: Extensions, message: T) -> Self {
Self {
metadata,
extensions,
message,
}
}
pub(crate) fn from_http_parts(parts: http::request::Parts, message: T) -> Self {
Request {
metadata: MetadataMap::from_headers(parts.headers),
tonic/src/server/grpc.rs
+56 -10
View File
@@ -1,6 +1,7 @@
use crate::{
body::BoxBody,
codec::{encode_server, Codec, Streaming},
interceptor::Interceptor,
server::{ClientStreamingService, ServerStreamingService, StreamingService, UnaryService},
Code, Request, Response, Status,
};
@@ -9,6 +10,16 @@ use futures_util::{future, stream, TryStreamExt};
use http_body::Body;
use std::fmt;
// A try! type macro for intercepting requests
macro_rules! t {
($expr : expr) => {
match $expr {
Ok(request) => request,
Err(res) => return res,
}
};
}
/// A gRPC Server handler.
///
/// This will wrap some inner [`Codec`] and provide utilities to handle
@@ -20,6 +31,7 @@ use std::fmt;
/// implements some [`Body`].
pub struct Grpc<T> {
codec: T,
interceptor: Option<Interceptor>,
}
impl<T> Grpc<T>
@@ -27,9 +39,21 @@ where
T: Codec,
T::Encode: Sync,
{
/// Creates a new gRPC client with the provided [`Codec`].
/// Creates a new gRPC server with the provided [`Codec`].
pub fn new(codec: T) -> Self {
Self { codec }
Self {
codec,
interceptor: None,
}
}
/// Creates a new gRPC server with the provided [`Codec`] and will apply the provided
/// interceptor on each inbound request.
pub fn with_interceptor(codec: T, interceptor: impl Into<Interceptor>) -> Self {
Self {
codec,
interceptor: Some(interceptor.into()),
}
}
/// Handle a single unary gRPC request.
@@ -53,6 +77,8 @@ where
}
};
let request = t!(self.intercept_request(request));
let response = service
.call(request)
.await
@@ -80,6 +106,8 @@ where
}
};
let request = t!(self.intercept_request(request));
let response = service.call(request).await;
self.map_response(response)
@@ -97,6 +125,7 @@ where
B::Error: Into<crate::Error> + Send + 'static,
{
let request = self.map_request_streaming(req);
let request = t!(self.intercept_request(request));
let response = service
.call(request)
.await
@@ -117,6 +146,7 @@ where
B::Error: Into<crate::Error> + Send,
{
let request = self.map_request_streaming(req);
let request = t!(self.intercept_request(request));
let response = service.call(request).await;
self.map_response(response)
}
@@ -180,18 +210,34 @@ where
http::Response::from_parts(parts, BoxBody::new(body))
}
Err(status) => {
let (mut parts, _body) = Response::new(()).into_http().into_parts();
Err(status) => Self::map_status(status),
}
}
parts.headers.insert(
http::header::CONTENT_TYPE,
http::header::HeaderValue::from_static("application/grpc"),
);
fn map_status(status: Status) -> http::Response<BoxBody> {
let (mut parts, _body) = Response::new(()).into_http().into_parts();
status.add_header(&mut parts.headers).unwrap();
parts.headers.insert(
http::header::CONTENT_TYPE,
http::header::HeaderValue::from_static("application/grpc"),
);
http::Response::from_parts(parts, BoxBody::empty())
status.add_header(&mut parts.headers).unwrap();
http::Response::from_parts(parts, BoxBody::empty())
}
fn intercept_request<A>(&self, req: Request<A>) -> Result<Request<A>, http::Response<BoxBody>> {
if let Some(interceptor) = &self.interceptor {
match interceptor.call(req) {
Ok(req) => Ok(req),
Err(status) => {
let res = Self::map_status(status);
return Err(res);
}
}
} else {
Ok(req)
}
}
}
tonic/src/transport/channel/endpoint.rs
-27
View File
@@ -10,7 +10,6 @@ use http::uri::{InvalidUri, Uri};
use std::{
convert::{TryFrom, TryInto},
fmt,
sync::Arc,
time::Duration,
};
use tower_make::MakeConnection;
@@ -27,8 +26,6 @@ pub struct Endpoint {
#[cfg(feature = "tls")]
pub(crate) tls: Option<TlsConnector>,
pub(crate) buffer_size: Option<usize>,
pub(crate) interceptor_headers:
Option<Arc<dyn Fn(&mut http::HeaderMap) + Send + Sync + 'static>>,
pub(crate) init_stream_window_size: Option<u32>,
pub(crate) init_connection_window_size: Option<u32>,
pub(crate) tcp_keepalive: Option<Duration>,
@@ -152,29 +149,6 @@ impl Endpoint {
}
}
/// Intercept outbound HTTP Request headers;
///
/// # Example
///
/// ```
/// # use tonic::transport::Endpoint;
/// # use std::time::Duration;
/// # let mut builder = Endpoint::from_static("https://example.com");
/// builder.intercept_headers(|headers| {
/// // Do something with headers
/// headers.insert("hello", "world".parse().unwrap());
/// });
/// ```
pub fn intercept_headers<F>(self, f: F) -> Self
where
F: Fn(&mut http::HeaderMap) + Send + Sync + 'static,
{
Endpoint {
interceptor_headers: Some(Arc::new(f)),
..self
}
}
/// Configures TLS for the endpoint.
#[cfg(feature = "tls")]
#[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
@@ -237,7 +211,6 @@ impl From<Uri> for Endpoint {
#[cfg(feature = "tls")]
tls: None,
buffer_size: None,
interceptor_headers: None,
init_stream_window_size: None,
init_connection_window_size: None,
tcp_keepalive: None,
tonic/src/transport/channel/mod.rs
+5 -27
View File
@@ -21,7 +21,6 @@ use std::{
fmt,
future::Future,
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
use tokio::io::{AsyncRead, AsyncWrite};
@@ -63,7 +62,6 @@ const DEFAULT_BUFFER_SIZE: usize = 1024;
#[derive(Clone)]
pub struct Channel {
svc: Buffer<Svc, Request<BoxBody>>,
interceptor_headers: Option<Arc<dyn Fn(&mut http::HeaderMap) + Send + Sync + 'static>>,
}
/// A future that resolves to an HTTP response.
@@ -114,14 +112,9 @@ 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, interceptor_headers)
Self::balance(discover, buffer_size)
}
pub(crate) async fn connect<C>(connector: C, endpoint: Endpoint) -> Result<Self, super::Error>
@@ -132,7 +125,6 @@ impl Channel {
C::Response: AsyncRead + AsyncWrite + HyperConnection + Unpin + Send + 'static,
{
let buffer_size = endpoint.buffer_size.clone().unwrap_or(DEFAULT_BUFFER_SIZE);
let interceptor_headers = endpoint.interceptor_headers.clone();
let svc = Connection::new(connector, endpoint)
.await
@@ -140,17 +132,10 @@ impl Channel {
let svc = Buffer::new(Either::A(svc), buffer_size);
Ok(Channel {
svc,
interceptor_headers,
})
Ok(Channel { svc })
}
pub(crate) fn balance<D>(
discover: D,
buffer_size: usize,
interceptor_headers: Option<Arc<dyn Fn(&mut http::HeaderMap) + Send + Sync + 'static>>,
) -> Self
pub(crate) fn balance<D>(discover: D, buffer_size: usize) -> Self
where
D: Discover<Service = Connection> + Unpin + Send + 'static,
D::Error: Into<crate::Error>,
@@ -161,10 +146,7 @@ impl Channel {
let svc = BoxService::new(svc);
let svc = Buffer::new(Either::B(svc), buffer_size);
Channel {
svc,
interceptor_headers,
}
Channel { svc }
}
}
@@ -177,11 +159,7 @@ impl GrpcService<BoxBody> for Channel {
GrpcService::poll_ready(&mut self.svc, cx).map_err(|e| super::Error::from_source(e))
}
fn call(&mut self, mut request: Request<BoxBody>) -> Self::Future {
if let Some(interceptor) = self.interceptor_headers.clone() {
interceptor(request.headers_mut());
}
fn call(&mut self, request: Request<BoxBody>) -> Self::Future {
let inner = GrpcService::call(&mut self.svc, request);
ResponseFuture { inner }
}
tonic/src/transport/mod.rs
+1 -6
View File
@@ -12,7 +12,6 @@
//! - Timeouts
//! - Concurrency Limits
//! - Rate limiting
//! - gRPC Interceptors
//!
//! # Examples
//!
@@ -77,10 +76,6 @@
//! .tls_config(ServerTlsConfig::with_rustls()
//! .identity(Identity::from_pem(&cert, &key)))
//! .concurrency_limit_per_connection(256)
//! .interceptor_fn(|svc, req| {
//! println!("Request: {:?}", req);
//! svc.call(req)
//! })
//! .add_service(my_svc)
//! .serve(addr)
//! .await?;
@@ -104,7 +99,7 @@ pub use self::error::Error;
#[doc(inline)]
pub use self::server::{Server, ServiceName};
pub use self::tls::{Certificate, Identity};
pub use hyper::Body;
pub use hyper::{Body, Uri};
#[cfg(feature = "tls")]
#[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
tonic/src/transport/server/mod.rs
+7 -54
View File
@@ -15,7 +15,7 @@ use super::service::TlsAcceptor;
use incoming::TcpIncoming;
use super::service::{layer_fn, Or, Routes, ServerIo, ServiceBuilderExt};
use super::service::{Or, Routes, ServerIo, ServiceBuilderExt};
use crate::{body::BoxBody, request::ConnectionInfo};
use futures_core::Stream;
use futures_util::{
@@ -35,15 +35,11 @@ use std::{
};
use tokio::io::{AsyncRead, AsyncWrite};
use tower::{
layer::{Layer, Stack},
limit::concurrency::ConcurrencyLimitLayer,
timeout::TimeoutLayer,
Service, ServiceBuilder,
limit::concurrency::ConcurrencyLimitLayer, timeout::TimeoutLayer, Service, ServiceBuilder,
};
use tracing_futures::{Instrument, Instrumented};
type BoxService = tower::util::BoxService<Request<Body>, Response<BoxBody>, crate::Error>;
type Interceptor = Arc<dyn Layer<BoxService, Service = BoxService> + Send + Sync + 'static>;
type TraceInterceptor = Arc<dyn Fn(&HeaderMap) -> tracing::Span + Send + Sync + 'static>;
/// A default batteries included `transport` server.
@@ -56,7 +52,6 @@ type TraceInterceptor = Arc<dyn Fn(&HeaderMap) -> tracing::Span + Send + Sync +
/// wanting to create a more complex and/or specific implementation.
#[derive(Default, Clone)]
pub struct Server {
interceptor: Option<Interceptor>,
trace_interceptor: Option<TraceInterceptor>,
concurrency_limit: Option<usize>,
timeout: Option<Duration>,
@@ -198,35 +193,6 @@ impl Server {
}
}
/// Intercept the execution of gRPC methods.
///
/// ```
/// # use tonic::transport::Server;
/// # use tower_service::Service;
/// # let mut builder = Server::builder();
/// builder.interceptor_fn(|svc, req| {
/// println!("request={:?}", req);
/// svc.call(req)
/// });
/// ```
pub fn interceptor_fn<F, Out>(self, f: F) -> Self
where
F: Fn(&mut BoxService, Request<Body>) -> Out + Send + Sync + 'static,
Out: Future<Output = Result<Response<BoxBody>, crate::Error>> + Send + 'static,
{
let f = Arc::new(f);
let interceptor = layer_fn(move |mut s| {
let f = f.clone();
tower::service_fn(move |req| f(&mut s, req))
});
let layer = Stack::new(interceptor, layer_fn(BoxService::new));
Server {
interceptor: Some(Arc::new(layer)),
..self
}
}
/// Intercept inbound headers and add a [`tracing::Span`] to each response future.
pub fn trace_fn<F>(self, f: F) -> Self
where
@@ -270,7 +236,6 @@ impl Server {
IE: Into<crate::Error>,
F: Future<Output = ()>,
{
let interceptor = self.interceptor.clone();
let span = self.trace_interceptor.clone();
let concurrency_limit = self.concurrency_limit;
let init_connection_window_size = self.init_connection_window_size;
@@ -283,7 +248,6 @@ impl Server {
let svc = MakeSvc {
inner: svc,
interceptor,
concurrency_limit,
timeout,
span,
@@ -480,7 +444,6 @@ impl<S> fmt::Debug for Svc<S> {
}
struct MakeSvc<S> {
interceptor: Option<Interceptor>,
concurrency_limit: Option<usize>,
timeout: Option<Duration>,
inner: S,
@@ -508,7 +471,6 @@ where
peer_certs: io.peer_certs().map(Arc::new),
};
let interceptor = self.interceptor.clone();
let svc = self.inner.clone();
let concurrency_limit = self.concurrency_limit;
let timeout = self.timeout.clone();
@@ -520,20 +482,11 @@ where
.optional_layer(timeout.map(TimeoutLayer::new))
.service(svc);
let svc = if let Some(interceptor) = interceptor {
let layered = interceptor.layer(BoxService::new(Svc {
inner: svc,
span,
conn_info,
}));
BoxService::new(layered)
} else {
BoxService::new(Svc {
inner: svc,
span,
conn_info,
})
};
let svc = BoxService::new(Svc {
inner: svc,
span,
conn_info,
});
Ok(svc)
})
tonic/src/transport/service/layer.rs
+2
View File
@@ -42,6 +42,8 @@ impl<L> ServiceBuilderExt<L> for ServiceBuilder<L> {
}
}
// TODO: figure out why this is causing a warning even though its used in optional_layer_fn
#[allow(dead_code)]
pub(crate) fn layer_fn<F>(f: F) -> LayerFn<F> {
LayerFn(f)
}
tonic/src/transport/service/mod.rs
+1 -1
View File
@@ -14,7 +14,7 @@ pub(crate) use self::connection::Connection;
pub(crate) use self::connector::connector;
pub(crate) use self::discover::ServiceList;
pub(crate) use self::io::ServerIo;
pub(crate) use self::layer::{layer_fn, ServiceBuilderExt};
pub(crate) use self::layer::ServiceBuilderExt;
pub(crate) use self::router::{Or, Routes};
#[cfg(feature = "tls")]
pub(crate) use self::tls::{TlsAcceptor, TlsConnector};