From 7e22a3a81323e59c66abe8a0001d69c363ef31a5 Mon Sep 17 00:00:00 2001
From: xuxuepeng <xuxuepeng1@huawei.com>
Date: Fri, 1 Sep 2023 19:15:53 +0800
Subject: [PATCH 4/4] Add vsock support

Signed-off-by: xuxuepeng <xuxuepeng1@huawei.com>
---
 src/client/client.rs | 45 ++++++++++++++++++++++++++++++++++++++------
 1 file changed, 39 insertions(+), 6 deletions(-)

diff --git a/src/client/client.rs b/src/client/client.rs
index 3f3e94a..2151e23 100644
--- a/src/client/client.rs
+++ b/src/client/client.rs
@@ -75,7 +75,35 @@ fn unix_sock(r#abstract: bool, socket_path: &str) -> Result<SockAddr> {
     Ok(sockaddr)
 }
 
-fn connect_to_socket(abs: bool, address: &str) -> Result<RawFd> {
+fn virtio_vsock(address: &str) -> Result<SockAddr> {
+    let (cid, port) = {
+        let vec: Vec<String> = address.split(":").map(String::from).collect();
+        if vec.len() != 2 {
+            let err_msg = format!("vsock address {address} is invalid");
+            return Err(other!(err_msg));
+        }
+        let cid = vec[0].parse::<u32>().map_err(other_error!(e, "failed to parse cid: "))?;
+        let port = vec[1].parse::<u32>().map_err(other_error!(e, "failed to parse port: "))?;
+        (cid, port)
+    };
+    let sockaddr = SockAddr::Vsock(VsockAddr::new(cid, port));
+    Ok(sockaddr)
+}
+
+fn connect_to_vsock(address: &str) -> Result<RawFd> {
+    let fd = socket(
+        AddressFamily::Vsock,
+        SockType::Stream,
+        SockFlag::empty(),
+        None,
+    )
+    .map_err(other_error!(e, "failed to create socket fd: "))?;
+    let sockaddr = virtio_vsock(address)?;
+    connect(fd, &sockaddr).map_err(other_error!(e, "failed to connect vsock: "))?;
+    Ok(fd)
+}
+
+fn connect_to_unix_socket(abs: bool, address: &str) -> Result<RawFd> {
     let fd = socket(
         AddressFamily::Unix,
         SockType::Stream,
@@ -90,14 +118,19 @@ fn connect_to_socket(abs: bool, address: &str) -> Result<RawFd> {
 }
 
 pub fn new_conn(container_id: &String, addr: &String) -> Result<()> {
-    let address = if addr.starts_with("unix://") {
-        addr.strip_prefix("unix://").unwrap()
+    let fd = if addr.starts_with("vsock://") {
+        let address = addr.strip_prefix("vsock://").unwrap();
+        connect_to_vsock(address)?
     } else {
-        addr
+        let address = if addr.starts_with("unix://") {
+            addr.strip_prefix("unix://").unwrap()
+        } else {
+            addr
+        };
+        let path = Path::new(&MAIN_SEPARATOR.to_string()).join(address);
+        connect_to_unix_socket(!addr.starts_with("unix://"), &path.to_string_lossy())?
     };
 
-    let path = Path::new(&MAIN_SEPARATOR.to_string()).join(address);
-    let fd = connect_to_socket(!addr.starts_with("unix://"), &path.to_string_lossy())?;
     TTRPC_CLIENTS.lock().unwrap().insert(
         container_id.clone(),
         Store {
-- 
2.40.1