[Agent] fix windows upgrade failed #20102

deepflowio · Jan 12, 2023 · 167d70d · 167d70d
1 parent 4c0ba7b
commit 167d70d
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 3 deletions.
diff --git a/agent/Cargo.lock b/agent/Cargo.lock
diff --git a/agent/Cargo.toml b/agent/Cargo.toml
@@ -81,6 +81,7 @@ kube = { git = "https://github.com/deepflowys/kube", tag = "0.74.1" }
 
 [target.'cfg(target_os = "windows")'.dependencies]
 pcap = "0.10.1"
+winapi = { version = "0.3.9", features = ["errhandlingapi", "libloaderapi", "minwindef", "winerror" ] }
 windows = { version = "0.30", features = ["Win32_Foundation", "Win32_System_Diagnostics_ToolHelp", "Win32_NetworkManagement_IpHelper", "Win32_Networking_WinSock", "Win32_System_LibraryLoader", "Win32_System_Threading", "Win32_System_ProcessStatus"] }
 windows_recv_engine = { path = "plugins/windows_recv_engine" }
 

diff --git a/agent/src/dispatcher/local_mode_dispatcher.rs b/agent/src/dispatcher/local_mode_dispatcher.rs
@@ -52,6 +52,7 @@ use public::{
 
 pub(super) struct LocalModeDispatcher {
     pub(super) base: BaseDispatcher,
+    #[cfg(target_os = "linux")]
     pub(super) extractor: Arc<LibvirtXmlExtractor>,
 }
 
@@ -289,7 +290,10 @@ impl LocalModeDispatcher {
     }
 
     pub(super) fn listener(&self) -> LocalModeDispatcherListener {
-        LocalModeDispatcherListener::new(self.base.listener(), self.extractor.clone())
+        #[cfg(target_os = "linux")]
+        return LocalModeDispatcherListener::new(self.base.listener(), self.extractor.clone());
+        #[cfg(target_os = "windows")]
+        return LocalModeDispatcherListener::new(self.base.listener());
     }
 }
 
@@ -304,14 +308,19 @@ impl LocalModeDispatcher {
 #[derive(Clone)]
 pub struct LocalModeDispatcherListener {
     base: BaseDispatcherListener,
+    #[cfg(target_os = "linux")]
     extractor: Arc<LibvirtXmlExtractor>,
     rewriter: MacRewriter,
 }
 
 impl LocalModeDispatcherListener {
-    pub(super) fn new(base: BaseDispatcherListener, extractor: Arc<LibvirtXmlExtractor>) -> Self {
+    pub(super) fn new(
+        base: BaseDispatcherListener,
+        #[cfg(target_os = "linux")] extractor: Arc<LibvirtXmlExtractor>,
+    ) -> Self {
         Self {
             base,
+            #[cfg(target_os = "linux")]
             extractor,
             rewriter: MacRewriter::new(),
         }
@@ -387,6 +396,7 @@ impl LocalModeDispatcherListener {
         #[cfg(target_os = "linux")]
         let index_to_mac_map =
             Self::get_if_index_to_inner_mac_map(&self.base.platform_poller, &self.base.netns);
+        #[cfg(target_os = "linux")]
         let name_to_mac_map = self.get_if_name_to_mac_map(tap_mac_script);
 
         for iface in interfaces.iter() {
@@ -417,14 +427,18 @@ impl LocalModeDispatcherListener {
                     }
                     new_mac
                 }
+                #[cfg(target_os = "linux")]
                 IfMacSource::IfLibvirtXml => {
                     *name_to_mac_map.get(&iface.name).unwrap_or(&iface.mac_addr)
                 }
+                #[cfg(target_os = "windows")]
+                IfMacSource::IfLibvirtXml => MacAddr::ZERO,
             });
         }
         macs
     }
 
+    #[cfg(target_os = "linux")]
     fn get_if_name_to_mac_map(&self, tap_mac_script: &str) -> HashMap<String, MacAddr> {
         let mut result = HashMap::new();
         if let Some(entries) = self.extractor.get_entries() {

diff --git a/agent/src/dispatcher/mod.rs b/agent/src/dispatcher/mod.rs
@@ -828,10 +828,15 @@ impl DispatcherBuilder {
         );
         let mut dispatcher = match tap_mode {
             TapMode::Local => {
+                #[cfg(target_os = "linux")]
                 let extractor = self
                     .libvirt_xml_extractor
                     .ok_or(Error::ConfigIncomplete("no libvirt xml extractor".into()))?;
-                DispatcherFlavor::Local(LocalModeDispatcher { base, extractor })
+                DispatcherFlavor::Local(LocalModeDispatcher {
+                    base,
+                    #[cfg(target_os = "linux")]
+                    extractor,
+                })
             }
             TapMode::Mirror => DispatcherFlavor::Mirror(MirrorModeDispatcher {
                 base,

diff --git a/agent/src/utils/environment.rs b/agent/src/utils/environment.rs
@@ -26,6 +26,8 @@ use std::{
     process, thread,
     time::Duration,
 };
+#[cfg(target_os = "windows")]
+use std::{ffi::OsString, os::windows::ffi::OsStringExt, ptr};
 
 use bytesize::ByteSize;
 #[cfg(target_os = "linux")]
@@ -34,6 +36,20 @@ use log::{error, info, warn};
 #[cfg(target_os = "linux")]
 use nom::AsBytes;
 use sysinfo::{DiskExt, System, SystemExt};
+#[cfg(target_os = "windows")]
+use winapi::{
+    shared::{
+        minwindef::{DWORD, MAX_PATH},
+        winerror::ERROR_INSUFFICIENT_BUFFER,
+    },
+    um::{
+        errhandlingapi::GetLastError,
+        libloaderapi::{
+            GetModuleFileNameW, GetModuleHandleExW, GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
+            GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+        },
+    },
+};
 
 use crate::common::PROCESS_NAME;
 use crate::error::{Error, Result};
@@ -426,6 +442,7 @@ pub fn running_in_container() -> bool {
     env::var_os("IN_CONTAINER").is_some()
 }
 
+#[cfg(target_os = "linux")]
 pub fn get_executable_path() -> Result<PathBuf, io::Error> {
     let possible_paths = vec![
         "/proc/self/exe".to_owned(),
@@ -444,6 +461,24 @@ pub fn get_executable_path() -> Result<PathBuf, io::Error> {
     ))
 }
 
+#[cfg(target_os = "windows")]
+pub fn get_executable_path() -> Result<PathBuf, io::Error> {
+    let mut buf = Vec::with_capacity(MAX_PATH);
+    unsafe {
+        let ret = GetModuleFileNameW(ptr::null_mut(), buf.as_mut_ptr(), MAX_PATH as DWORD) as usize;
+        if ret > 0 && ret < MAX_PATH {
+            buf.set_len(ret);
+            let s = OsString::from_wide(&buf);
+            Ok(s.into())
+        } else {
+            Err(io::Error::new(
+                io::ErrorKind::NotFound,
+                "executable path not found",
+            ))
+        }
+    }
+}
+
 pub fn get_mac_by_name(src_interface: String) -> u32 {
     if src_interface.is_empty() {
         return 0;