From d3820f08e5feae6ab82539e4a22edf697bf93bc8 Mon Sep 17 00:00:00 2001
From: Liang Zhang <zhangliang5@huawei.com>
Date: Sun, 7 Feb 2021 17:03:44 +0800
Subject: [PATCH] socket: clean tmporary file after vm exit

Signed-off-by: Liang Zhang <zhangliang5@huawei.com>
---
 device_model/src/virtio/console.rs    |  9 ++---
 machine_manager/src/event_loop.rs     | 26 -------------
 machine_manager/src/lib.rs            |  1 +
 machine_manager/src/qmp/mod.rs        | 10 ++---
 machine_manager/src/signal_handler.rs |  9 +++--
 machine_manager/src/socket.rs         | 40 ++++++-------------
 machine_manager/src/temp_cleaner.rs   | 56 +++++++++++++++++++++++++++
 src/main.rs                           | 19 +++++----
 util/src/loop_context.rs              | 23 +----------
 9 files changed, 94 insertions(+), 99 deletions(-)
 create mode 100644 machine_manager/src/temp_cleaner.rs

diff --git a/device_model/src/virtio/console.rs b/device_model/src/virtio/console.rs
index 8305192..4cfbc63 100644
--- a/device_model/src/virtio/console.rs
+++ b/device_model/src/virtio/console.rs
@@ -18,7 +18,7 @@ use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::{Arc, Mutex};
 
 use address_space::AddressSpace;
-use machine_manager::{config::ConsoleConfig, event_loop::EventLoop};
+use machine_manager::{config::ConsoleConfig, event_loop::EventLoop, temp_cleaner::TempCleaner};
 use util::byte_code::ByteCode;
 use util::loop_context::{read_fd, EventNotifier, EventNotifierHelper, NotifierOperation};
 use util::num_ops::{read_u32, write_u32};
@@ -297,8 +297,6 @@ pub struct Console {
     driver_features: u64,
     /// UnixListener for virtio-console to communicate in host.
     listener: UnixListener,
-    /// Path to console socket file.
-    path: String,
 }
 
 impl Console {
@@ -312,6 +310,9 @@ impl Console {
         let listener = UnixListener::bind(path.as_str())
             .unwrap_or_else(|_| panic!("Failed to bind socket for console {}", path));
 
+        // add file to temporary pool, so it could be clean when vm exit.
+        TempCleaner::add_path(path.clone());
+
         limit_permission(path.as_str())
             .unwrap_or_else(|_| panic!("Failed to change file permission for console {}", path));
         Console {
@@ -319,7 +320,6 @@ impl Console {
             device_features: 0_u64,
             driver_features: 0_u64,
             listener,
-            path,
         }
     }
 }
@@ -329,7 +329,6 @@ impl VirtioDevice for Console {
     fn realize(&mut self) -> Result<()> {
         self.device_features = 1_u64 << VIRTIO_F_VERSION_1 | 1_u64 << VIRTIO_CONSOLE_F_SIZE;
 
-        EventLoop::add_files(self.path.clone());
         Ok(())
     }
 
diff --git a/machine_manager/src/event_loop.rs b/machine_manager/src/event_loop.rs
index 84188e5..03f55f7 100644
--- a/machine_manager/src/event_loop.rs
+++ b/machine_manager/src/event_loop.rs
@@ -138,30 +138,4 @@ impl EventLoop {
             }
         }
     }
-
-    /// Exit vm and clean environment.
-    ///
-    /// # Notes
-    ///
-    /// Set CURRENT_MAINLOOP to None to release MainLoopContext,
-    /// otherwise, it will not be released automatically.
-    pub fn clean() -> bool {
-        unsafe {
-            if GLOBAL_EVENT_LOOP.is_none() {
-                return true;
-            }
-        }
-        let ret = Self::get_ctx(None).remove_files();
-        unsafe {
-            if GLOBAL_EVENT_LOOP.is_some() {
-                GLOBAL_EVENT_LOOP = None;
-            }
-        }
-        ret
-    }
-
-    /// Add to be removed files to Eventloop.
-    pub fn add_files(path: String) {
-        Self::get_ctx(None).add_files(path);
-    }
 }
diff --git a/machine_manager/src/lib.rs b/machine_manager/src/lib.rs
index 7bc2536..8de5dae 100644
--- a/machine_manager/src/lib.rs
+++ b/machine_manager/src/lib.rs
@@ -34,6 +34,7 @@ pub mod machine;
 pub mod qmp;
 pub mod signal_handler;
 pub mod socket;
+pub mod temp_cleaner;
 
 pub mod errors {
     error_chain! {
diff --git a/machine_manager/src/qmp/mod.rs b/machine_manager/src/qmp/mod.rs
index af565d8..5828fb8 100644
--- a/machine_manager/src/qmp/mod.rs
+++ b/machine_manager/src/qmp/mod.rs
@@ -46,7 +46,7 @@ use vmm_sys_util::terminal::Terminal;
 
 use crate::machine::MachineExternalInterface;
 use crate::socket::SocketRWHandler;
-use crate::{errors::Result, event_loop::EventLoop};
+use crate::{errors::Result, temp_cleaner::TempCleaner};
 use qmp_schema as schema;
 use schema::QmpCommand;
 
@@ -345,9 +345,7 @@ pub fn handle_qmp(stream_fd: RawFd, controller: &Arc<dyn MachineExternalInterfac
                     reason: "host-qmp-quit".to_string(),
                 };
                 event!(SHUTDOWN; shutdown_msg);
-                if !EventLoop::clean() {
-                    error!("Clean environment failed!");
-                }
+                TempCleaner::clean();
 
                 std::io::stdin()
                     .lock()
@@ -619,7 +617,7 @@ mod tests {
         let (listener, mut client, server) = prepare_unix_socket_environment("06");
 
         // Use event! macro to send event msg to client
-        let socket = Socket::from_unix_listener(listener, None, "test_06.sock".to_string());
+        let socket = Socket::from_unix_listener(listener, None);
         socket.bind_unix_stream(server);
         QmpChannel::bind_writer(SocketRWHandler::new(socket.get_stream_fd()));
 
@@ -669,7 +667,7 @@ mod tests {
         let (listener, mut client, server) = prepare_unix_socket_environment("07");
 
         // Use event! macro to send event msg to client
-        let socket = Socket::from_unix_listener(listener, None, "test_07.sock".to_string());
+        let socket = Socket::from_unix_listener(listener, None);
         socket.bind_unix_stream(server);
 
         // 1.send greeting response
diff --git a/machine_manager/src/signal_handler.rs b/machine_manager/src/signal_handler.rs
index e3bd8bb..79af118 100644
--- a/machine_manager/src/signal_handler.rs
+++ b/machine_manager/src/signal_handler.rs
@@ -10,7 +10,7 @@
 // NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 // See the Mulan PSL v2 for more details.
 extern crate vmm_sys_util;
-use crate::event_loop::EventLoop;
+use crate::temp_cleaner::TempCleaner;
 use std::io::Write;
 
 use libc::{c_int, c_void, siginfo_t};
@@ -22,9 +22,10 @@ const SYSTEMCALL_OFFSET: isize = 6;
 
 fn basic_clean() {
     info!("Removing environment and exiting...");
-    if !EventLoop::clean() {
-        error!("Clean environment failed!");
-    }
+
+    // clean temporary file
+    TempCleaner::clean();
+
     std::io::stdin()
         .lock()
         .set_canon_mode()
diff --git a/machine_manager/src/socket.rs b/machine_manager/src/socket.rs
index c7fb402..526923d 100644
--- a/machine_manager/src/socket.rs
+++ b/machine_manager/src/socket.rs
@@ -11,11 +11,10 @@
 // See the Mulan PSL v2 for more details.
 
 use serde::Deserialize;
-use std::io::{Read, Write};
+use std::io::{Error, ErrorKind, Read, Write};
 use std::os::unix::io::{AsRawFd, RawFd};
 use std::os::unix::net::{UnixListener, UnixStream};
 use std::sync::{Arc, Mutex, RwLock};
-use std::{fs, io};
 
 use util::loop_context::{EventNotifier, EventNotifierHelper, NotifierOperation};
 use vmm_sys_util::epoll::EventSet;
@@ -42,7 +41,7 @@ const MAX_SOCKET_MSG_LENGTH: usize = 8192;
 ///
 /// fn main() -> std::io::Result<()> {
 ///     let listener = UnixListener::bind("/path/to/my/socket")?;
-///     let socket = Socket::from_unix_listener(listener, None, "/path/to/my/socket".to_string());
+///     let socket = Socket::from_unix_listener(listener, None);
 ///     assert!(!socket.is_connected());
 ///
 ///     let client_stream = UnixStream::connect("/path/to/my/socket")?;
@@ -61,21 +60,6 @@ pub struct Socket {
     stream: RwLock<Option<SocketStream>>,
     /// Perform socket command
     performer: Option<Arc<dyn MachineExternalInterface>>,
-    /// path of socket file
-    path: String,
-}
-
-impl Drop for Socket {
-    fn drop(&mut self) {
-        if let Err(ref e) = fs::remove_file(self.path.clone()) {
-            error!(
-                "Failed to remove api socket file: {} , {} , please remove it manually",
-                &self.path, e
-            );
-        } else {
-            info!("Delete api Socket file: {} successfully", &self.path);
-        }
-    }
 }
 
 impl Socket {
@@ -88,14 +72,12 @@ impl Socket {
     pub fn from_unix_listener(
         listener: UnixListener,
         performer: Option<Arc<dyn MachineExternalInterface>>,
-        path: String,
     ) -> Self {
         Socket {
             sock_type: SocketType::Unix,
             listener,
             stream: RwLock::new(None),
             performer,
-            path,
         }
     }
 
@@ -413,8 +395,8 @@ impl SocketRWHandler {
             let ret = unsafe { recvmsg(self.socket_fd, &mut mhdr, MSG_DONTWAIT) };
 
             if ret == -1 {
-                let sock_err = io::Error::last_os_error();
-                if sock_err.kind() == io::ErrorKind::WouldBlock {
+                let sock_err = Error::last_os_error();
+                if sock_err.kind() == ErrorKind::WouldBlock {
                     break;
                 } else {
                     return Err(sock_err);
@@ -456,7 +438,7 @@ impl SocketRWHandler {
             if let Some(pos) = self.pos.checked_add(1) {
                 self.pos = pos;
             } else {
-                return Err(io::ErrorKind::InvalidInput.into());
+                return Err(ErrorKind::InvalidInput.into());
             }
         }
         Ok(())
@@ -496,8 +478,8 @@ impl SocketRWHandler {
         mhdr.msg_flags = 0;
 
         if unsafe { sendmsg(self.socket_fd, &mhdr, 0) } == -1 {
-            Err(io::Error::new(
-                io::ErrorKind::BrokenPipe,
+            Err(Error::new(
+                ErrorKind::BrokenPipe,
                 "The socket pipe is broken!",
             ))
         } else {
@@ -529,7 +511,7 @@ impl Write for SocketRWHandler {
         if let Some(pos) = self.pos.checked_add(buf.len()) {
             self.pos = pos;
         } else {
-            return Err(io::ErrorKind::InvalidInput.into());
+            return Err(ErrorKind::InvalidInput.into());
         }
 
         self.write_fd(buf.len())?;
@@ -628,8 +610,8 @@ impl SocketHandler {
                 let _ = self.stream.write(&[b'\n'])?;
                 Ok(())
             }
-            Err(_) => Err(io::Error::new(
-                io::ErrorKind::BrokenPipe,
+            Err(_) => Err(Error::new(
+                ErrorKind::BrokenPipe,
                 "The socket pipe is broken!",
             )),
         }
@@ -822,7 +804,7 @@ mod tests {
     fn test_socket_lifecycle() {
         // Pre test. Environment Preparation
         let (listener, _, server) = prepare_unix_socket_environment("04");
-        let socket = Socket::from_unix_listener(listener, None, "test_04.sock".to_string());
+        let socket = Socket::from_unix_listener(listener, None);
 
         // life cycle test
         // 1.Unconnected
diff --git a/machine_manager/src/temp_cleaner.rs b/machine_manager/src/temp_cleaner.rs
new file mode 100644
index 0000000..47c2b98
--- /dev/null
+++ b/machine_manager/src/temp_cleaner.rs
@@ -0,0 +1,56 @@
+// Copyright (c) 2020 Huawei Technologies Co.,Ltd. All rights reserved.
+//
+// StratoVirt is licensed under Mulan PSL v2.
+// You can use this software according to the terms and conditions of the Mulan
+// PSL v2.
+// You may obtain a copy of Mulan PSL v2 at:
+//         http://license.coscl.org.cn/MulanPSL2
+// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
+// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
+// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+// See the Mulan PSL v2 for more details.
+
+use std::fs;
+
+static mut GLOBAL_TEMP_CLEANER: Option<TempCleaner> = None;
+
+/// This structure used to keep temporary file which was created by program, and would be deleted
+/// when Vm exit.
+pub struct TempCleaner {
+    /// Path of files that should be removed after exiting the vm.
+    paths: Vec<String>,
+}
+
+impl TempCleaner {
+    pub fn object_init() {
+        unsafe {
+            if GLOBAL_TEMP_CLEANER.is_none() {
+                GLOBAL_TEMP_CLEANER = Some(TempCleaner { paths: Vec::new() });
+            }
+        }
+    }
+
+    /// Add to be removed file path
+    pub fn add_path(path: String) {
+        unsafe {
+            if let Some(tmp) = GLOBAL_TEMP_CLEANER.as_mut() {
+                tmp.paths.push(path);
+            }
+        }
+    }
+
+    /// Clean the temporary files
+    pub fn clean() {
+        unsafe {
+            if let Some(tmp) = GLOBAL_TEMP_CLEANER.as_mut() {
+                while let Some(path) = tmp.paths.pop() {
+                    if let Err(ref e) = fs::remove_file(&path) {
+                        error!("Failed to delete console / socket file:{} :{}", &path, e);
+                    } else {
+                        info!("Delete file: {} successfully.", &path);
+                    }
+                }
+            }
+        }
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index 78abf93..9cfecbe 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -28,6 +28,7 @@ use machine_manager::socket::Socket;
 use machine_manager::{
     cmdline::{check_api_channel, create_args_parser, create_vmconfig},
     event_loop::EventLoop,
+    temp_cleaner::TempCleaner,
 };
 use machine_manager::{config::VmConfig, signal_handler::register_kill_signal};
 use util::loop_context::EventNotifierHelper;
@@ -81,9 +82,9 @@ fn run() -> Result<()> {
         } else {
             error!("Panic at [{}: {}].", panic_file, panic_line);
         }
-        if !EventLoop::clean() {
-            error!("Clean environment failed!");
-        }
+
+        // clean temporary file
+        TempCleaner::clean();
     }));
 
     let vm_config: VmConfig = create_vmconfig(&cmd_args)?;
@@ -109,9 +110,8 @@ fn run() -> Result<()> {
         }
     }
 
-    if !EventLoop::clean() {
-        error!("Clean environment failed!");
-    }
+    // clean temporary file
+    TempCleaner::clean();
 
     Ok(())
 }
@@ -129,6 +129,7 @@ fn real_main(cmd_args: &arg_parser::ArgMatches, vm_config: VmConfig) -> Result<(
             .chain_err(|| "Failed to set terminal to raw mode.")?;
     }
 
+    TempCleaner::object_init();
     QmpChannel::object_init();
     EventLoop::object_init(&vm_config.iothreads)?;
     register_kill_signal();
@@ -140,9 +141,13 @@ fn real_main(cmd_args: &arg_parser::ArgMatches, vm_config: VmConfig) -> Result<(
         let (api_path, _) = check_api_channel(&cmd_args)?;
         let listener = UnixListener::bind(&api_path)
             .chain_err(|| format!("Failed to bind api socket {}", &api_path))?;
+
+        // add file to temporary pool, so it could be clean when vm exit.
+        TempCleaner::add_path(api_path.clone());
+
         limit_permission(&api_path)
             .chain_err(|| format!("Failed to limit permission for api socket {}", &api_path))?;
-        Socket::from_unix_listener(listener, Some(vm.clone()), api_path)
+        Socket::from_unix_listener(listener, Some(vm.clone()))
     };
 
     EventLoop::update_event(
diff --git a/util/src/loop_context.rs b/util/src/loop_context.rs
index eff7b0f..0f637aa 100644
--- a/util/src/loop_context.rs
+++ b/util/src/loop_context.rs
@@ -12,9 +12,9 @@
 
 extern crate vmm_sys_util;
 
+use std::collections::BTreeMap;
 use std::os::unix::io::RawFd;
 use std::sync::{Arc, Mutex, RwLock};
-use std::{collections::BTreeMap, fs};
 
 use libc::{c_void, read};
 use vmm_sys_util::epoll::{ControlOperation, Epoll, EpollEvent, EventSet};
@@ -114,8 +114,6 @@ pub struct EventLoopContext {
     gc: Arc<RwLock<Vec<Box<EventNotifier>>>>,
     /// Temp events vector, store wait returned events.
     ready_events: Vec<EpollEvent>,
-    /// Path of files that should be removed after exiting the vm.
-    files: Vec<String>,
 }
 
 unsafe impl Sync for EventLoopContext {}
@@ -130,7 +128,6 @@ impl EventLoopContext {
             events: Arc::new(RwLock::new(BTreeMap::new())),
             gc: Arc::new(RwLock::new(Vec::new())),
             ready_events: vec![EpollEvent::default(); READY_EVENT_MAX],
-            files: Vec::new(),
         }
     }
 
@@ -296,24 +293,6 @@ impl EventLoopContext {
 
         Ok(true)
     }
-
-    /// Add path of files.
-    pub fn add_files(&mut self, path: String) {
-        self.files.push(path);
-    }
-
-    /// Remove unused files.
-    pub fn remove_files(&mut self) -> bool {
-        while let Some(file) = self.files.pop() {
-            if let Err(ref e) = fs::remove_file(&file) {
-                error!("Failed to delete console socket file:{} :{}", &file, e);
-                return false;
-            } else {
-                info!("Delete file: {} successfully.", &file);
-            }
-        }
-        true
-    }
 }
 
 impl Default for EventLoopContext {
-- 
2.27.0

