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