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 {
say_hello( &self, request: Request<HelloRequest>, ) -> Result<Response<HelloReply>, Status>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]
main() -> Result<(), Box<dyn std::error::Error>>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.
intercept(req: Request<()>) -> Result<Request<()>, Status>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
layer(&self, service: S) -> Self::Service71 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
poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>93 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
94 self.inner.poll_ready(cx)
95 }
96
call(&mut self, req: http::Request<ReqBody>) -> Self::Future97 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