Merge pull request #38 from Tim-Zhang/fix-async-vsock

Fix vsock in async mode

Author: Hui Zhu

Cargo.toml

@@ -21,6 +21,7 @@ async-trait = "0.1.31"
 
 tokio = { version = "0.2", features = ["rt-threaded", "sync", "uds", "stream", "macros", "io-util"] }
 futures = "0.3"
+tokio-vsock = "0.2.1"
 
 [build-dependencies]
 protobuf-codegen-pure = "2.14.0"

example/Cargo.lock

@@ -9,32 +9,37 @@ use std::os::unix::io::RawFd;
 use std::sync::Arc;
 
 use crate::asynchronous::stream::{receive, respond, respond_with_status};
-use crate::common;
-use crate::common::MESSAGE_TYPE_REQUEST;
+use crate::common::{self, Domain, MESSAGE_TYPE_REQUEST};
 use crate::error::{get_status, Error, Result};
 use crate::r#async::{MethodHandler, TtrpcContext};
 use crate::ttrpc::{Code, Request};
 use crate::MessageHeader;
 use futures::StreamExt as _;
+use std::marker::Unpin;
 use std::os::unix::io::FromRawFd;
+use std::os::unix::net::UnixListener as SysUnixListener;
 use tokio::{
     self,
     io::split,
     net::UnixListener,
     prelude::*,
+    stream::Stream,
     sync::mpsc::{channel, Receiver, Sender},
 };
+use tokio_vsock::VsockListener;
 
 pub struct Server {
     listeners: Vec<RawFd>,
     methods: Arc<HashMap<String, Box<dyn MethodHandler + Send + Sync>>>,
+    domain: Option<Domain>,
 }
 
 impl Default for Server {
     fn default() -> Self {
         Server {
             listeners: Vec::with_capacity(1),
             methods: Arc::new(HashMap::new()),
+            domain: None,
         }
     }
 }
@@ -51,9 +56,10 @@ impl Server {
             ));
         }
 
-        let fd = common::do_bind(host)?;
-        self.listeners.push(fd);
+        let (fd, domain) = common::do_bind(host)?;
+        self.domain = Some(domain);
 
+        self.listeners.push(fd);
         Ok(self)
     }
 
@@ -77,21 +83,42 @@ impl Server {
             return Err(Error::Others("ttrpc-rust not bind".to_string()));
         }
 
-        let listener = self.listeners[0];
-        common::do_listen(listener)?;
+        let listenfd = self.listeners[0];
+        common::do_listen(listenfd)?;
 
-        Ok(listener)
+        Ok(listenfd)
     }
 
     pub async fn start(&self) -> Result<()> {
-        let listener = self.listen()?;
-        let sys_unix_listener: std::os::unix::net::UnixListener;
-        unsafe {
-            sys_unix_listener = std::os::unix::net::UnixListener::from_raw_fd(listener);
+        let listenfd = self.listen()?;
+
+        match self.domain.as_ref().unwrap() {
+            Domain::Unix => {
+                let sys_unix_listener;
+                unsafe {
+                    sys_unix_listener = SysUnixListener::from_raw_fd(listenfd);
+                }
+                let mut unix_listener = UnixListener::from_std(sys_unix_listener).unwrap();
+                let incoming = unix_listener.incoming();
+
+                self.do_start(listenfd, incoming).await
+            }
+            Domain::Vsock => {
+                let incoming;
+                unsafe {
+                    incoming = VsockListener::from_raw_fd(listenfd).incoming();
+                }
+
+                self.do_start(listenfd, incoming).await
+            }
         }
-        let mut unix_listener = UnixListener::from_std(sys_unix_listener).unwrap();
-        let mut incoming = unix_listener.incoming();
+    }
 
+    pub async fn do_start<I, S>(&self, listenfd: RawFd, mut incoming: I) -> Result<()>
+    where
+        I: Stream<Item = std::io::Result<S>> + Unpin,
+        S: AsyncRead + AsyncWrite + Send + 'static,
+    {
         while let Some(result) = incoming.next().await {
             match result {
                 Ok(stream) => {
@@ -115,7 +142,7 @@ impl Server {
                             match receive(&mut reader).await {
                                 Ok(message) => {
                                     tokio::spawn(async move {
-                                        handle_request(tx, listener, methods, message).await;
+                                        handle_request(tx, listenfd, methods, message).await;
                                     });
                                 }
                                 Err(e) => {

src/asynchronous/server.rs

@@ -11,6 +11,12 @@ use nix::sys::socket::*;
 use std::os::unix::io::RawFd;
 use std::str::FromStr;
 
+#[derive(Debug)]
+pub enum Domain {
+    Unix,
+    Vsock,
+}
+
 #[derive(Default, Debug)]
 pub struct MessageHeader {
     pub length: u32,
@@ -30,18 +36,29 @@ pub fn do_listen(listener: RawFd) -> Result<()> {
     listen(listener, 10).map_err(|e| Error::Socket(e.to_string()))
 }
 
-pub fn do_bind(host: &str) -> Result<RawFd> {
+pub fn parse_host(host: &str) -> Result<(Domain, Vec<&str>)> {
     let hostv: Vec<&str> = host.trim().split("://").collect();
     if hostv.len() != 2 {
         return Err(Error::Others(format!("Host {} is not right", host)));
     }
-    let scheme = hostv[0].to_lowercase();
+
+    let domain = match &hostv[0].to_lowercase()[..] {
+        "unix" => Domain::Unix,
+        "vsock" => Domain::Vsock,
+        x => return Err(Error::Others(format!("Scheme {:?} is not supported", x))),
+    };
+
+    Ok((domain, hostv))
+}
+
+pub fn do_bind(host: &str) -> Result<(RawFd, Domain)> {
+    let (domain, hostv) = parse_host(host)?;
 
     let sockaddr: SockAddr;
     let fd: RawFd;
 
-    match scheme.as_str() {
-        "unix" => {
+    match domain {
+        Domain::Unix => {
             fd = socket(
                 AddressFamily::Unix,
                 SockType::Stream,
@@ -54,8 +71,7 @@ pub fn do_bind(host: &str) -> Result<RawFd> {
                 UnixAddr::new_abstract(sockaddr_h.as_bytes()).map_err(err_to_Others!(e, ""))?;
             sockaddr = SockAddr::Unix(sockaddr_u);
         }
-
-        "vsock" => {
+        Domain::Vsock => {
             let host_port_v: Vec<&str> = hostv[1].split(':').collect();
             if host_port_v.len() != 2 {
                 return Err(Error::Others(format!(
@@ -75,12 +91,11 @@ pub fn do_bind(host: &str) -> Result<RawFd> {
             .map_err(|e| Error::Socket(e.to_string()))?;
             sockaddr = SockAddr::new_vsock(cid, port);
         }
-        _ => return Err(Error::Others(format!("Scheme {} is not supported", scheme))),
     };
 
     bind(fd, &sockaddr).map_err(err_to_Others!(e, ""))?;
 
-    Ok(fd)
+    Ok((fd, domain))
 }
 
 macro_rules! cfg_sync {

src/common.rs

@@ -275,7 +275,7 @@ impl Server {
             ));
         }
 
-        let fd = common::do_bind(host)?;
+        let (fd, _) = common::do_bind(host)?;
         self.listeners.push(fd);
 
         Ok(self)