From 38caab24c5087551d6ec12d0002df798e5e4f5ac Mon Sep 17 00:00:00 2001
From: Priyansh Rathi <techiepriyansh@gmail.com>
Date: Mon, 11 Sep 2023 20:37:45 -0700
Subject: [PATCH] vsock: Don't allow duplicate CIDs

Currently, while updating the `cid_map`, it is not checked if the map
already contained a key with the same CID before. So, fail to create the
vsock thread if the CID is already in use.

Signed-off-by: Priyansh Rathi <techiepriyansh@gmail.com>
---
 crates/vsock/src/vhu_vsock.rs        |  2 ++
 crates/vsock/src/vhu_vsock_thread.rs | 44 +++++++++++++++++++++-------
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/crates/vsock/src/vhu_vsock.rs b/crates/vsock/src/vhu_vsock.rs
index 0d25b38..e67cc16 100644
--- a/crates/vsock/src/vhu_vsock.rs
+++ b/crates/vsock/src/vhu_vsock.rs
@@ -135,6 +135,8 @@ pub(crate) enum Error {
     EventFdCreate(std::io::Error),
     #[error("Raw vsock packets queue is empty")]
     EmptyRawPktsQueue,
+    #[error("CID already in use by another vsock device")]
+    CidAlreadyInUse,
 }
 
 impl std::convert::From<Error> for std::io::Error {
diff --git a/crates/vsock/src/vhu_vsock_thread.rs b/crates/vsock/src/vhu_vsock_thread.rs
index 0c2ab2f..b0fe774 100644
--- a/crates/vsock/src/vhu_vsock_thread.rs
+++ b/crates/vsock/src/vhu_vsock_thread.rs
@@ -111,14 +111,21 @@ impl VhostUserVsockThread {
             cid_map.clone(),
         );
 
-        cid_map.write().unwrap().insert(
-            guest_cid,
-            (
-                thread_backend.raw_pkts_queue.clone(),
-                groups_set,
-                sibling_event_fd.try_clone().unwrap(),
-            ),
-        );
+        {
+            let mut cid_map = cid_map.write().unwrap();
+            if cid_map.contains_key(&guest_cid) {
+                return Err(Error::CidAlreadyInUse);
+            }
+
+            cid_map.insert(
+                guest_cid,
+                (
+                    thread_backend.raw_pkts_queue.clone(),
+                    groups_set,
+                    sibling_event_fd.try_clone().unwrap(),
+                ),
+            );
+        }
 
         let thread = VhostUserVsockThread {
             mem: None,
@@ -824,9 +831,14 @@ mod tests {
             .join("test_vsock_thread_failures.vsock")
             .display()
             .to_string();
-        let mut t =
-            VhostUserVsockThread::new(vsock_socket_path, 3, CONN_TX_BUF_SIZE, groups, cid_map)
-                .unwrap();
+        let mut t = VhostUserVsockThread::new(
+            vsock_socket_path,
+            3,
+            CONN_TX_BUF_SIZE,
+            groups.clone(),
+            cid_map.clone(),
+        )
+        .unwrap();
         assert!(VhostUserVsockThread::epoll_register(-1, -1, epoll::Events::EPOLLIN).is_err());
         assert!(VhostUserVsockThread::epoll_modify(-1, -1, epoll::Events::EPOLLIN).is_err());
         assert!(VhostUserVsockThread::epoll_unregister(-1, -1).is_err());
@@ -848,6 +860,16 @@ mod tests {
         assert!(t.process_rx(&vring, false).is_err());
         assert!(t.process_rx(&vring, true).is_err());
 
+        // trying to use a CID that is already in use should fail
+        let vsock_socket_path2 = test_dir
+            .path()
+            .join("test_vsock_thread_failures2.vsock")
+            .display()
+            .to_string();
+        let t2 =
+            VhostUserVsockThread::new(vsock_socket_path2, 3, CONN_TX_BUF_SIZE, groups, cid_map);
+        assert!(t2.is_err());
+
         test_dir.close().unwrap();
     }
 }