1 files changed, 51 insertions, 14 deletions

diff --git a/lib/rs/test/src/bin/kitchen_sink_client.rs b/lib/rs/test/src/bin/kitchen_sink_client.rs
index 74197de7f..b98afb814 100644
--- a/lib/rs/test/src/bin/kitchen_sink_client.rs
+++ b/lib/rs/test/src/bin/kitchen_sink_client.rs
@@ -16,8 +16,16 @@
 // under the License.
 
 use clap::{clap_app, value_t};
+use log::*;
 
 use std::convert::Into;
+use std::net::TcpStream;
+use std::net::ToSocketAddrs;
+
+#[cfg(unix)]
+use std::os::unix::net::UnixStream;
+#[cfg(unix)]
+use std::path::Path;
 
 use kitchen_sink::base_two::{TNapkinServiceSyncClient, TRamenServiceSyncClient};
 use kitchen_sink::midlayer::{MealServiceSyncClient, TMealServiceSyncClient};
@@ -30,9 +38,9 @@ use thrift::protocol::{
     TBinaryInputProtocol, TBinaryOutputProtocol, TCompactInputProtocol, TCompactOutputProtocol,
     TInputProtocol, TOutputProtocol,
 };
-use thrift::transport::{
-    ReadHalf, TFramedReadTransport, TFramedWriteTransport, TIoChannel, TTcpChannel, WriteHalf,
-};
+use thrift::transport::{TFramedReadTransport, TFramedWriteTransport, TIoChannel, TTcpChannel};
+
+type IoProtocol = (Box<dyn TInputProtocol>, Box<dyn TOutputProtocol>);
 
 fn main() {
     match run() {
@@ -51,6 +59,7 @@ fn run() -> thrift::Result<()> {
         (about: "Thrift Rust kitchen sink client")
         (@arg host: --host +takes_value "Host on which the Thrift test server is located")
         (@arg port: --port +takes_value "Port on which the Thrift test server is listening")
+        (@arg domain_socket: --("domain-socket") + takes_value "Unix Domain Socket on which the Thrift test server is listening")
         (@arg protocol: --protocol +takes_value "Thrift protocol implementation to use (\"binary\", \"compact\")")
         (@arg service: --service +takes_value "Service type to contact (\"part\", \"full\", \"recursive\")")
     )
@@ -58,10 +67,47 @@ fn run() -> thrift::Result<()> {
 
     let host = matches.value_of("host").unwrap_or("127.0.0.1");
     let port = value_t!(matches, "port", u16).unwrap_or(9090);
+    let domain_socket = matches.value_of("domain_socket");
     let protocol = matches.value_of("protocol").unwrap_or("compact");
     let service = matches.value_of("service").unwrap_or("part");
 
-    let (i_chan, o_chan) = tcp_channel(host, port)?;
+    let (i_prot, o_prot) = match domain_socket {
+        None => {
+            let listen_address = format!("{}:{}", host, port);
+            info!("Client binds to {} with {}", listen_address, protocol);
+            bind(listen_address, protocol)?
+        }
+        Some(domain_socket) => {
+            info!("Client binds to {} (UDS) with {}", domain_socket, protocol);
+            bind_uds(domain_socket, protocol)?
+        }
+    };
+
+    run_client(service, i_prot, o_prot)
+}
+
+fn bind<A: ToSocketAddrs>(listen_address: A, protocol: &str) -> Result<IoProtocol, thrift::Error> {
+    let stream = TcpStream::connect(listen_address)?;
+    let channel = TTcpChannel::with_stream(stream);
+
+    let (i_prot, o_prot) = build(channel, protocol)?;
+    Ok((i_prot, o_prot))
+}
+
+#[cfg(unix)]
+fn bind_uds<P: AsRef<Path>>(domain_socket: P, protocol: &str) -> Result<IoProtocol, thrift::Error> {
+    let stream = UnixStream::connect(domain_socket)?;
+
+    let (i_prot, o_prot) = build(stream, protocol)?;
+    Ok((i_prot, o_prot))
+}
+
+fn build<C: TIoChannel + 'static>(
+    channel: C,
+    protocol: &str,
+) -> thrift::Result<(Box<dyn TInputProtocol>, Box<dyn TOutputProtocol>)> {
+    let (i_chan, o_chan) = channel.split()?;
+
     let (i_tran, o_tran) = (
         TFramedReadTransport::new(i_chan),
         TFramedWriteTransport::new(o_chan),
@@ -79,7 +125,7 @@ fn run() -> thrift::Result<()> {
         unmatched => return Err(format!("unsupported protocol {}", unmatched).into()),
     };
 
-    run_client(service, i_prot, o_prot)
+    Ok((i_prot, o_prot))
 }
 
 fn run_client(
@@ -98,15 +144,6 @@ fn run_client(
     }
 }
 
-fn tcp_channel(
-    host: &str,
-    port: u16,
-) -> thrift::Result<(ReadHalf<TTcpChannel>, WriteHalf<TTcpChannel>)> {
-    let mut c = TTcpChannel::new();
-    c.open(&format!("{}:{}", host, port))?;
-    c.split()
-}
-
 fn exec_meal_client(
     i_prot: Box<dyn TInputProtocol>,
     o_prot: Box<dyn TOutputProtocol>,