From 368d5eca4621e73b9b4d9e9296b9537529f59f43 Mon Sep 17 00:00:00 2001
From: Florian RICHER <florian.richer@protonmail.com>
Date: Thu, 26 Jan 2023 22:40:09 +0100
Subject: [PATCH] Add support of Unix socket

---
 Cargo.lock         |  2 ++
 Cargo.toml         |  2 ++
 src/client/main.rs | 26 +++++++++++++++++++++++---
 src/daemon/main.rs | 39 ++++++++++++++++++++++++++++++++-------
 4 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d7073ad..44ab4ce 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -122,8 +122,10 @@ version = "0.1.0"
 dependencies = [
  "prost",
  "tokio",
+ "tokio-stream",
  "tonic",
  "tonic-build",
+ "tower",
 ]
 
 [[package]]
diff --git a/Cargo.toml b/Cargo.toml
index 4ce03ed..1cc7021 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,6 +19,8 @@ path = "src/client/main.rs"
 
 [dependencies]
 tokio = { version = "1.0", features = ["rt-multi-thread", "macros"] } # Required for tonic
+tokio-stream = { version = "0.1", features = ["net"] } # Required for tonic with unix socket
+tower = "0.4" # Required for tonic with unix socket
 prost = "0.11" # Required for tonic
 tonic = "0.8"
 
diff --git a/src/client/main.rs b/src/client/main.rs
index 35cdb50..f9f80e3 100644
--- a/src/client/main.rs
+++ b/src/client/main.rs
@@ -1,13 +1,28 @@
-use hello_world::greeter_client::GreeterClient;
-use hello_world::HelloRequest;
+#![cfg_attr(not(unix), allow(unused_imports))]
 
 pub mod hello_world {
     tonic::include_proto!("helloworld");
 }
 
+use hello_world::{greeter_client::GreeterClient, HelloRequest};
+#[cfg(unix)]
+use tokio::net::UnixStream;
+use tonic::transport::{Endpoint, Uri};
+use tower::service_fn;
+
+#[cfg(unix)]
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let mut client = GreeterClient::connect("http://[::1]:50051").await?;
+    let channel = Endpoint::try_from("http://[::]:50051")?
+        .connect_with_connector(service_fn(|_: Uri| {
+            let path = "helloworld.sock";
+
+            // Connect to a Uds socket
+            UnixStream::connect(path)
+        }))
+        .await?;
+
+    let mut client = GreeterClient::new(channel);
 
     let request = tonic::Request::new(HelloRequest {
         name: "Tonic".into(),
@@ -18,4 +33,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     println!("RESPONSE={:?}", response);
 
     Ok(())
+}
+
+#[cfg(not(unix))]
+fn main() {
+    panic!("The `uds` example only works on unix systems!");
 }
\ No newline at end of file
diff --git a/src/daemon/main.rs b/src/daemon/main.rs
index 22af836..f8d00a1 100644
--- a/src/daemon/main.rs
+++ b/src/daemon/main.rs
@@ -1,12 +1,23 @@
-use tonic::{transport::Server, Request, Response, Status};
+#![cfg_attr(not(unix), allow(unused_imports))]
 
-use hello_world::greeter_server::{Greeter, GreeterServer};
-use hello_world::{HelloReply, HelloRequest};
+use std::path::Path;
+#[cfg(unix)]
+use tokio::net::UnixListener;
+#[cfg(unix)]
+use tokio_stream::wrappers::UnixListenerStream;
+#[cfg(unix)]
+use tonic::transport::server::UdsConnectInfo;
+use tonic::{transport::Server, Request, Response, Status};
 
 pub mod hello_world {
     tonic::include_proto!("helloworld");
 }
 
+use hello_world::{
+    greeter_server::{Greeter, GreeterServer},
+    HelloReply, HelloRequest,
+};
+
 #[derive(Default)]
 pub struct MyGreeter {}
 
@@ -16,7 +27,11 @@ impl Greeter for MyGreeter {
         &self,
         request: Request<HelloRequest>,
     ) -> Result<Response<HelloReply>, Status> {
-        println!("Got a request from {:?}", request.remote_addr());
+        #[cfg(unix)]
+        {
+            let conn_info = request.extensions().get::<UdsConnectInfo>().unwrap();
+            println!("Got a request {:?} with info {:?}", request, conn_info);
+        }
 
         let reply = hello_world::HelloReply {
             message: format!("Hello {}!", request.into_inner().name),
@@ -25,17 +40,27 @@ impl Greeter for MyGreeter {
     }
 }
 
+#[cfg(unix)]
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let addr = "[::1]:50051".parse().unwrap();
+    let path = "helloworld.sock";
+
+    std::fs::create_dir_all(Path::new(path).parent().unwrap())?;
+
     let greeter = MyGreeter::default();
 
-    println!("GreeterServer listening on {}", addr);
+    let uds = UnixListener::bind(path)?;
+    let uds_stream = UnixListenerStream::new(uds);
 
     Server::builder()
         .add_service(GreeterServer::new(greeter))
-        .serve(addr)
+        .serve_with_incoming(uds_stream)
         .await?;
 
     Ok(())
+}
+
+#[cfg(not(unix))]
+fn main() {
+    panic!("The `uds` example only works on unix systems!");
 }
\ No newline at end of file