xref: /tonic/examples/src/uds/server.rs (revision 380d81dd)
1 #![cfg_attr(not(unix), allow(unused_imports))]
2 
3 use futures::TryFutureExt;
4 use std::path::Path;
5 #[cfg(unix)]
6 use tokio::net::UnixListener;
7 use tonic::{transport::Server, Request, Response, Status};
8 
9 pub mod hello_world {
10     tonic::include_proto!("helloworld");
11 }
12 
13 use hello_world::{
14     greeter_server::{Greeter, GreeterServer},
15     HelloReply, HelloRequest,
16 };
17 
18 #[derive(Default)]
19 pub struct MyGreeter {}
20 
21 #[tonic::async_trait]
22 impl Greeter for MyGreeter {
23     async fn say_hello(
24         &self,
25         request: Request<HelloRequest>,
26     ) -> Result<Response<HelloReply>, Status> {
27         #[cfg(unix)]
28         {
29             let conn_info = request.extensions().get::<unix::UdsConnectInfo>().unwrap();
30             println!("Got a request {:?} with info {:?}", request, conn_info);
31         }
32 
33         let reply = hello_world::HelloReply {
34             message: format!("Hello {}!", request.into_inner().name),
35         };
36         Ok(Response::new(reply))
37     }
38 }
39 
40 #[cfg(unix)]
41 #[tokio::main]
42 async fn main() -> Result<(), Box<dyn std::error::Error>> {
43     let path = "/tmp/tonic/helloworld";
44 
45     tokio::fs::create_dir_all(Path::new(path).parent().unwrap()).await?;
46 
47     let greeter = MyGreeter::default();
48 
49     let incoming = {
50         let uds = UnixListener::bind(path)?;
51 
52         async_stream::stream! {
53             while let item = uds.accept().map_ok(|(st, _)| unix::UnixStream(st)).await {
54                 yield item;
55             }
56         }
57     };
58 
59     Server::builder()
60         .add_service(GreeterServer::new(greeter))
61         .serve_with_incoming(incoming)
62         .await?;
63 
64     Ok(())
65 }
66 
67 #[cfg(unix)]
68 mod unix {
69     use std::{
70         pin::Pin,
71         sync::Arc,
72         task::{Context, Poll},
73     };
74 
75     use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
76     use tonic::transport::server::Connected;
77 
78     #[derive(Debug)]
79     pub struct UnixStream(pub tokio::net::UnixStream);
80 
81     impl Connected for UnixStream {
82         type ConnectInfo = UdsConnectInfo;
83 
84         fn connect_info(&self) -> Self::ConnectInfo {
85             UdsConnectInfo {
86                 peer_addr: self.0.peer_addr().ok().map(Arc::new),
87                 peer_cred: self.0.peer_cred().ok(),
88             }
89         }
90     }
91 
92     #[derive(Clone, Debug)]
93     pub struct UdsConnectInfo {
94         pub peer_addr: Option<Arc<tokio::net::unix::SocketAddr>>,
95         pub peer_cred: Option<tokio::net::unix::UCred>,
96     }
97 
98     impl AsyncRead for UnixStream {
99         fn poll_read(
100             mut self: Pin<&mut Self>,
101             cx: &mut Context<'_>,
102             buf: &mut ReadBuf<'_>,
103         ) -> Poll<std::io::Result<()>> {
104             Pin::new(&mut self.0).poll_read(cx, buf)
105         }
106     }
107 
108     impl AsyncWrite for UnixStream {
109         fn poll_write(
110             mut self: Pin<&mut Self>,
111             cx: &mut Context<'_>,
112             buf: &[u8],
113         ) -> Poll<std::io::Result<usize>> {
114             Pin::new(&mut self.0).poll_write(cx, buf)
115         }
116 
117         fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
118             Pin::new(&mut self.0).poll_flush(cx)
119         }
120 
121         fn poll_shutdown(
122             mut self: Pin<&mut Self>,
123             cx: &mut Context<'_>,
124         ) -> Poll<std::io::Result<()>> {
125             Pin::new(&mut self.0).poll_shutdown(cx)
126         }
127     }
128 }
129 
130 #[cfg(not(unix))]
131 fn main() {
132     panic!("The `uds` example only works on unix systems!");
133 }
134