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