1 use hyper::Body; 2 use std::{ 3 pin::Pin, 4 task::{Context, Poll}, 5 time::Duration, 6 }; 7 use tonic::{body::BoxBody, transport::Server, Request, Response, Status}; 8 use tower::{Layer, Service}; 9 10 use hello_world::greeter_server::{Greeter, GreeterServer}; 11 use hello_world::{HelloReply, HelloRequest}; 12 13 pub mod hello_world { 14 tonic::include_proto!("helloworld"); 15 } 16 17 #[derive(Default)] 18 pub struct MyGreeter {} 19 20 #[tonic::async_trait] 21 impl Greeter for MyGreeter { 22 async fn say_hello( 23 &self, 24 request: Request<HelloRequest>, 25 ) -> Result<Response<HelloReply>, Status> { 26 println!("Got a request from {:?}", request.remote_addr()); 27 28 let reply = hello_world::HelloReply { 29 message: format!("Hello {}!", request.into_inner().name), 30 }; 31 Ok(Response::new(reply)) 32 } 33 } 34 35 #[tokio::main] 36 async fn main() -> Result<(), Box<dyn std::error::Error>> { 37 let addr = "[::1]:50051".parse().unwrap(); 38 let greeter = MyGreeter::default(); 39 40 println!("GreeterServer listening on {}", addr); 41 42 let svc = GreeterServer::new(greeter); 43 44 // The stack of middleware that our service will be wrapped in 45 let layer = tower::ServiceBuilder::new() 46 // Apply middleware from tower 47 .timeout(Duration::from_secs(30)) 48 // Apply our own middleware 49 .layer(MyMiddlewareLayer::default()) 50 // Interceptors can be also be applied as middleware 51 .layer(tonic::service::interceptor(intercept)) 52 .into_inner(); 53 54 Server::builder() 55 // Wrap all services in the middleware stack 56 .layer(layer) 57 .add_service(svc) 58 .serve(addr) 59 .await?; 60 61 Ok(()) 62 } 63 64 // An interceptor function. 65 fn intercept(req: Request<()>) -> Result<Request<()>, Status> { 66 Ok(req) 67 } 68 69 #[derive(Debug, Clone, Default)] 70 struct MyMiddlewareLayer; 71 72 impl<S> Layer<S> for MyMiddlewareLayer { 73 type Service = MyMiddleware<S>; 74 75 fn layer(&self, service: S) -> Self::Service { 76 MyMiddleware { inner: service } 77 } 78 } 79 80 #[derive(Debug, Clone)] 81 struct MyMiddleware<S> { 82 inner: S, 83 } 84 85 type BoxFuture<'a, T> = Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>; 86 87 impl<S> Service<hyper::Request<Body>> for MyMiddleware<S> 88 where 89 S: Service<hyper::Request<Body>, Response = hyper::Response<BoxBody>> + Clone + Send + 'static, 90 S::Future: Send + 'static, 91 { 92 type Response = S::Response; 93 type Error = S::Error; 94 type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>; 95 96 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { 97 self.inner.poll_ready(cx) 98 } 99 100 fn call(&mut self, req: hyper::Request<Body>) -> Self::Future { 101 // This is necessary because tonic internally uses `tower::buffer::Buffer`. 102 // See https://github.com/tower-rs/tower/issues/547#issuecomment-767629149 103 // for details on why this is necessary 104 let clone = self.inner.clone(); 105 let mut inner = std::mem::replace(&mut self.inner, clone); 106 107 Box::pin(async move { 108 // Do extra async work here... 109 let response = inner.call(req).await?; 110 111 Ok(response) 112 }) 113 } 114 } 115