1 use std::{ 2 pin::Pin, 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 type BoxFuture<'a, T> = Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>; 85 86 impl<S> Service<hyper::Request<BoxBody>> for MyMiddleware<S> 87 where 88 S: Service<hyper::Request<BoxBody>, Response = hyper::Response<BoxBody>> 89 + Clone 90 + Send 91 + 'static, 92 S::Future: Send + 'static, 93 { 94 type Response = S::Response; 95 type Error = S::Error; 96 type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>; 97 98 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { 99 self.inner.poll_ready(cx) 100 } 101 102 fn call(&mut self, req: hyper::Request<BoxBody>) -> Self::Future { 103 // This is necessary because tonic internally uses `tower::buffer::Buffer`. 104 // See https://github.com/tower-rs/tower/issues/547#issuecomment-767629149 105 // for details on why this is necessary 106 let clone = self.inner.clone(); 107 let mut inner = std::mem::replace(&mut self.inner, clone); 108 109 Box::pin(async move { 110 // Do extra async work here... 111 let response = inner.call(req).await?; 112 113 Ok(response) 114 }) 115 } 116 } 117