Implements open_dir and write_data_root for Windows. (#1243)

obhq · Jan 12, 2025 · 9f1c61f · 9f1c61f
1 parent b84b0f0
commit 9f1c61f
Show file tree

Hide file tree

Showing 4 changed files with 140 additions and 11 deletions.
diff --git a/.vscode/tasks.json b/.vscode/tasks.json
@@ -29,6 +29,9 @@
       "osx": {
         "command": "${workspaceFolder}/dist/Obliteration.app/Contents/MacOS/Obliteration"
       },
+      "windows": {
+        "command": "${workspaceFolder}/dist/obliteration.exe"
+      },
       "args": [
         "--debug",
         "127.0.0.1:1234"

diff --git a/gui/Cargo.toml b/gui/Cargo.toml
@@ -75,11 +75,16 @@ windows-sys = { version = "0.59.0", features = [
     "Win32_Foundation",
     "Win32_Security",
     "Win32_System",
+    "Win32_System_Com",
     "Win32_System_Hypervisor",
     "Win32_System_Memory",
     "Win32_System_Registry",
     "Win32_System_SystemInformation",
     "Win32_UI",
+    "Win32_UI_Controls",
+    "Win32_UI_Controls_Dialogs",
+    "Win32_UI_Shell",
+    "Win32_UI_Shell_Common",
     "Win32_UI_WindowsAndMessaging",
 ] }
 

diff --git a/gui/src/setup/windows.rs b/gui/src/setup/windows.rs
@@ -3,10 +3,12 @@ use thiserror::Error;
 use windows_sys::w;
 use windows_sys::Win32::Foundation::{ERROR_FILE_NOT_FOUND, ERROR_SUCCESS};
 use windows_sys::Win32::System::Registry::{
-    RegCloseKey, RegCreateKeyExW, RegQueryValueExW, HKEY, HKEY_CURRENT_USER, KEY_ALL_ACCESS,
-    REG_OPTION_NON_VOLATILE, REG_SZ,
+    RegCloseKey, RegCreateKeyExW, RegQueryValueExW, RegSetValueExW, HKEY, HKEY_CURRENT_USER,
+    KEY_ALL_ACCESS, REG_OPTION_NON_VOLATILE, REG_SZ,
 };
 
+const FQVN: &str = "HKEY_CURRENT_USER\\Software\\OBHQ\\Obliteration\\DataRoot";
+
 pub fn read_data_root() -> Result<Option<String>, DataRootError> {
     // Open our registry key.
     let mut key = null_mut();
@@ -32,7 +34,6 @@ pub fn read_data_root() -> Result<Option<String>, DataRootError> {
     }
 
     // Get data size.
-    let fqvn = "HKEY_CURRENT_USER\\Software\\OBHQ\\Obliteration\\DataRoot";
     let key = Key(key);
     let name = w!("DataRoot");
     let mut ty = 0;
@@ -44,9 +45,9 @@ pub fn read_data_root() -> Result<Option<String>, DataRootError> {
     } else if e != ERROR_SUCCESS {
         let e = std::io::Error::from_raw_os_error(e.try_into().unwrap());
 
-        return Err(DataRootError::QueryRegKey(fqvn, e));
+        return Err(DataRootError::QueryRegKey(FQVN, e));
     } else if ty != REG_SZ {
-        return Err(DataRootError::InvalidRegValue(fqvn));
+        return Err(DataRootError::InvalidRegValue(FQVN));
     }
 
     // Read data.
@@ -65,9 +66,9 @@ pub fn read_data_root() -> Result<Option<String>, DataRootError> {
     if e != ERROR_SUCCESS {
         let e = std::io::Error::from_raw_os_error(e.try_into().unwrap());
 
-        return Err(DataRootError::QueryRegKey(fqvn, e));
+        return Err(DataRootError::QueryRegKey(FQVN, e));
     } else if ty != REG_SZ {
-        return Err(DataRootError::InvalidRegValue(fqvn));
+        return Err(DataRootError::InvalidRegValue(FQVN));
     }
 
     // Remove null-terminators if any.
@@ -79,12 +80,61 @@ pub fn read_data_root() -> Result<Option<String>, DataRootError> {
 
     // Convert to Rust string.
     String::from_utf16(&buf)
-        .map_err(|_| DataRootError::InvalidRegValue(fqvn))
+        .map_err(|_| DataRootError::InvalidRegValue(FQVN))
         .map(Some)
 }
 
 pub fn write_data_root(path: impl AsRef<str>) -> Result<(), DataRootError> {
-    todo!()
+    // Open our registry key.
+    let mut key = null_mut();
+    let key_path = "HKEY_CURRENT_USER\\Software\\OBHQ\\Obliteration";
+    let e = unsafe {
+        RegCreateKeyExW(
+            HKEY_CURRENT_USER,
+            w!("Software\\OBHQ\\Obliteration"),
+            0,
+            null(),
+            REG_OPTION_NON_VOLATILE,
+            KEY_ALL_ACCESS,
+            null(),
+            &mut key,
+            null_mut(),
+        )
+    };
+
+    if e != ERROR_SUCCESS {
+        let e = std::io::Error::from_raw_os_error(e.try_into().unwrap());
+        return Err(DataRootError::CreateRegKey(key_path, e));
+    }
+
+    // RAII to auto-close key.
+    let key = Key(key);
+
+    // Prepare the path in UTF-16 with a terminating 0.
+    let wide_path = path
+        .as_ref()
+        .encode_utf16()
+        .chain(std::iter::once(0))
+        .collect::<Vec<u16>>();
+
+    // Write the value.
+    let e = unsafe {
+        RegSetValueExW(
+            key.0,
+            w!("DataRoot"),
+            0,
+            REG_SZ,
+            wide_path.as_ptr().cast(),
+            (wide_path.len() * 2) as u32,
+        )
+    };
+
+    if e != ERROR_SUCCESS {
+        let e = std::io::Error::from_raw_os_error(e.try_into().unwrap());
+        return Err(DataRootError::WriteRegKey(FQVN, e));
+    }
+
+    Ok(())
 }
 
 /// RAII struct to close `HKEY` when dropped.
@@ -105,6 +155,9 @@ pub enum DataRootError {
     #[error("couldn't read {0}")]
     QueryRegKey(&'static str, #[source] std::io::Error),
 
+    #[error("couldn't write {0}")]
+    WriteRegKey(&'static str, #[source] std::io::Error),
+
     #[error("{0} has invalid value")]
     InvalidRegValue(&'static str),
 }
diff --git a/gui/src/ui/windows/dialogs.rs b/gui/src/ui/windows/dialogs.rs
@@ -1,15 +1,83 @@
 use crate::rt::WinitWindow;
 use crate::ui::FileType;
+use futures::channel::oneshot;
+use raw_window_handle::{HasWindowHandle, RawWindowHandle};
+use std::future::Future;
+use std::num::NonZero;
 use std::path::PathBuf;
+use windows_sys::Win32::System::Com::CoTaskMemFree;
+use windows_sys::Win32::UI::Controls::Dialogs::{GetOpenFileNameW, OPENFILENAMEW};
+use windows_sys::Win32::UI::Shell::{
+    SHBrowseForFolderW, SHGetPathFromIDListW, BIF_NEWDIALOGSTYLE, BIF_RETURNONLYFSDIRS, BROWSEINFOW,
+};
+
+fn get_hwnd<T: WinitWindow>(parent: &T) -> NonZero<isize> {
+    let parent = parent.handle();
+    let parent = parent.window_handle().unwrap();
+    let RawWindowHandle::Win32(win) = parent.as_ref() else {
+        unreachable!();
+    };
+
+    win.hwnd
+}
+
+fn spawn_dialog<F>(dialog_fn: F) -> impl Future<Output = Option<PathBuf>>
+where
+    F: Send + 'static + FnOnce() -> Option<PathBuf>,
+{
+    let (tx, rx) = oneshot::channel();
+
+    std::thread::spawn(move || {
+        let res = dialog_fn();
+        tx.send(res).unwrap();
+    });
+
+    async move { rx.await.unwrap_or(None) }
+}
 
 pub async fn open_file<T: WinitWindow>(
     parent: &T,
     title: impl AsRef<str>,
     ty: FileType,
 ) -> Option<PathBuf> {
-    todo!();
+    todo!()
 }
 
 pub async fn open_dir<T: WinitWindow>(parent: &T, title: impl AsRef<str>) -> Option<PathBuf> {
-    todo!()
+    let hwnd = get_hwnd(parent);
+
+    let title_wide: Vec<u16> = title
+        .as_ref()
+        .encode_utf16()
+        .chain(std::iter::once(0))
+        .collect();
+
+    spawn_dialog(move || unsafe {
+        let mut bi: BROWSEINFOW = std::mem::zeroed();
+        bi.hwndOwner = hwnd.get() as _;
+
+        bi.lpszTitle = title_wide.as_ptr();
+
+        bi.ulFlags = BIF_RETURNONLYFSDIRS | BIF_NEWDIALOGSTYLE;
+
+        let pidl = SHBrowseForFolderW(&mut bi);
+        if !pidl.is_null() {
+            let mut buffer = [0u16; 260];
+            if SHGetPathFromIDListW(pidl, buffer.as_mut_ptr()) != 0 {
+                let len = buffer.iter().position(|&c| c == 0).unwrap_or(buffer.len());
+                let path_str = String::from_utf16_lossy(&buffer[..len]);
+                let dir_path = PathBuf::from(path_str);
+
+                CoTaskMemFree(pidl as _);
+
+                Some(dir_path)
+            } else {
+                CoTaskMemFree(pidl as _);
+                None
+            }
+        } else {
+            None
+        }
+    })
+    .await
 }