diff --git a/Cargo.toml b/Cargo.toml index 5bca66e..1587538 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,7 @@ exclude.workspace = true readme.workspace = true [dependencies] -bottles-core = { workspace = true } +bottles-core = { path = "../next-core" } tonic = { workspace = true } tokio = { workspace = true } @@ -27,3 +27,6 @@ features = [ "Win32_UI_WindowsAndMessaging", "Win32_UI_Shell", ] + +[dependencies.windows-registry] +version = "0.5.1" diff --git a/README.md b/README.md index 3611d46..e0d9002 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,15 @@ # next-winebridge The source code of the Next version of WineBridge + +## Testing +Tests for winebridge are meant to be run inside wine. Generate the test executables using the following command: +```bash +cargo test --no-run --lib +``` + +This will generate the test executables in the `target/$TARGET/debug/deps` which can be run using wine as follows: +```bash +wine target/$TARGET/debug/deps/next-winebridge-*.exe --test-threads=1 +``` + +The `--test-threads=1` flag is required becase parallel operations on Wine registry can interfere with each other and cause tests to fail. This issue is not specific to this project. diff --git a/src/lib.rs b/src/lib.rs index 9a01875..835882c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,11 +1,14 @@ mod processes; +mod registry; use bottles_core::proto::{self, wine_bridge_server::WineBridge}; use processes::{manager::ProcessManager, process::ProcessIdentifier}; +use registry::manager::{to_proto_reg_val, to_reg_data, KeyExtension, RegistryManager}; use windows::{ core::{s, PCSTR}, Win32::UI::WindowsAndMessaging::{MessageBoxA, MB_OK}, }; +use windows_registry::Key; #[derive(Debug, Default)] pub struct WineBridgeService; @@ -87,4 +90,135 @@ impl WineBridge for WineBridgeService { Ok(tonic::Response::new(proto::KillProcessResponse {})) } + + async fn create_registry_key( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let input = request.get_ref(); + let hive = input + .hive + .parse() + .map_err(|e| tonic::Status::invalid_argument(format!("Invalid hive: {:?}", e)))?; + let subkey = std::path::Path::new(input.subkey.as_str()); + + RegistryManager + .create_key(hive, subkey) + .map(|_| tonic::Response::new(proto::MessageResponse { success: true })) + .map_err(|e| tonic::Status::internal(format!("Failed to create registry key: {:?}", e))) + } + + async fn delete_registry_key( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let input = request.get_ref(); + let hive = input + .hive + .parse() + .map_err(|e| tonic::Status::invalid_argument(format!("Invalid hive: {:?}", e)))?; + let subkey = std::path::Path::new(input.subkey.as_str()); + + RegistryManager + .delete_key(hive, subkey) + .map(|_| tonic::Response::new(proto::MessageResponse { success: true })) + .map_err(|e| tonic::Status::internal(format!("Failed to create registry key: {:?}", e))) + } + + async fn get_registry_key( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let input = request.get_ref(); + let hive = input + .hive + .parse() + .map_err(|e| tonic::Status::invalid_argument(format!("Invalid hive: {:?}", e)))?; + let subkey = std::path::Path::new(input.subkey.as_str()); + + let key = RegistryManager + .key(hive, subkey) + .map_err(|e| { + tonic::Status::internal(format!("Failed to create registry key: {:?}", e)) + })? + .as_registry_key(hive, subkey); + + Ok(tonic::Response::new(key)) + } + + async fn get_registry_key_value( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let input = request.get_ref(); + let name = input.name.as_str(); + let hive = input + .hive + .parse() + .map_err(|e| tonic::Status::invalid_argument(format!("Invalid hive: {:?}", e)))?; + let subkey = std::path::Path::new(input.subkey.as_str()); + + let key = RegistryManager.key(hive, subkey).map_err(|e| { + tonic::Status::internal(format!("Failed to create registry key: {:?}", e)) + })?; + + let value = key.value(name).map_err(|e| { + tonic::Status::internal(format!("Failed to get registry value: {:?}", e)) + })?; + + Ok(tonic::Response::new(to_proto_reg_val(value))) + } + + async fn set_registry_key_value( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let input = request.get_ref(); + + let (name, key) = input + .key + .as_ref() + .map(|k| { + let hive = k.hive.parse().unwrap(); + let subkey = std::path::Path::new(k.subkey.as_str()); + (k.name.clone(), RegistryManager.key(hive, subkey).unwrap()) + }) + .ok_or_else(|| tonic::Status::invalid_argument("Key is required"))?; + + let value = input + .value + .as_ref() + .ok_or_else(|| tonic::Status::invalid_argument("Value is required"))?; + + key.create_value( + name.as_str(), + to_reg_data(value.r#type(), value.data.clone()), + ) + .map(|_| tonic::Response::new(proto::MessageResponse { success: true })) + .map_err(|e| tonic::Status::internal(format!("Failed to set registry value: {:?}", e))) + } + + async fn delete_registry_key_value( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let input = request.get_ref(); + let name = input.name.as_str(); + let hive = input + .hive + .parse() + .map_err(|e| tonic::Status::invalid_argument(format!("Invalid hive: {:?}", e)))?; + let subkey = std::path::Path::new(input.subkey.as_str()); + + RegistryManager + .key(hive, subkey) + .map_err(|e| { + tonic::Status::internal(format!("Failed to create registry key: {:?}", e)) + })? + .remove_value(name) + .map(|_| tonic::Response::new(proto::MessageResponse { success: true })) + .map_err(|e| { + tonic::Status::internal(format!("Failed to delete registry value: {:?}", e)) + }) + } } diff --git a/src/registry/manager.rs b/src/registry/manager.rs new file mode 100644 index 0000000..03bb0cf --- /dev/null +++ b/src/registry/manager.rs @@ -0,0 +1,304 @@ +use bottles_core::proto; +use std::{ops::Deref, path::Path, str::FromStr}; +use windows_registry::*; + +#[derive(Debug, Eq, PartialEq)] +pub enum Data { + DWord(u32), + QWord(u64), + String(String), + ExpandString(String), + MultiString(Vec), + Bytes(Vec), + Other, +} + +impl Hive { + pub fn inner(&self) -> &Key { + match self { + Hive::ClassesRoot => CLASSES_ROOT, + Hive::CurrentConfig => CURRENT_CONFIG, + Hive::CurrentUser => CURRENT_USER, + Hive::LocalMachine => LOCAL_MACHINE, + Hive::Users => USERS, + } + } + + pub fn to_string(&self) -> String { + match self { + Hive::ClassesRoot => "HKCR".to_string(), + Hive::CurrentConfig => "HKCC".to_string(), + Hive::CurrentUser => "HKCU".to_string(), + Hive::LocalMachine => "HKLM".to_string(), + Hive::Users => "HKU".to_string(), + } + } +} + +#[derive(Debug, Copy, Clone)] +pub enum Hive { + ClassesRoot, + CurrentConfig, + CurrentUser, + LocalMachine, + Users, +} + +impl FromStr for Hive { + type Err = &'static str; + + fn from_str(s: &str) -> std::result::Result { + match s.to_ascii_uppercase().as_str() { + "CLASSESROOT" | "HKCR" => Ok(Hive::ClassesRoot), + "CURRENTCONFIG" | "HKCC" => Ok(Hive::CurrentConfig), + "CURRENTUSER" | "HKCU" => Ok(Hive::CurrentUser), + "LOCALMACHINE" | "HKLM" => Ok(Hive::LocalMachine), + "USERS" | "HKU" => Ok(Hive::Users), + _ => Err("invalid registry hive"), + } + } +} + +pub trait KeyExtension { + fn get(hive: Hive, subkey: &Path) -> Result { + hive.inner().open(subkey.display().to_string()) + } + + fn new(hive: Hive, subkey: &Path) -> Result { + hive.inner().create(subkey.display().to_string()) + } + + fn delete(hive: Hive, subkey: &Path) -> Result<()> { + hive.inner().remove_tree(subkey.display().to_string()) + } + + fn value(&self, name: &str) -> Result; + fn values(&self) -> Result>; + fn create_value(&self, name: &str, data: Data) -> Result<()>; + fn rename_value(&self, old_name: &str, new_name: &str) -> Result<()>; + + fn as_registry_key(&self, hive: Hive, subkey: &Path) -> proto::RegistryKey; +} + +pub fn to_reg_data(ty: proto::RegistryValueType, data: Vec) -> Data { + match ty { + proto::RegistryValueType::RegBinary => Data::Bytes(data), + proto::RegistryValueType::RegDword => { + let val = u32::from_le_bytes(data.try_into().unwrap()); + Data::DWord(val) + } + proto::RegistryValueType::RegQword => { + let val = u64::from_le_bytes(data.try_into().unwrap()); + Data::QWord(val) + } + proto::RegistryValueType::RegSz => { + let val = String::from_utf8(data).unwrap(); + Data::String(val) + } + proto::RegistryValueType::RegExpandSz => { + let val = String::from_utf8(data).unwrap(); + Data::ExpandString(val) + } + proto::RegistryValueType::RegMultiSz => { + let val = String::from_utf8(data).unwrap(); + let strings: Vec = val.split('\0').map(|s| s.to_string()).collect(); + Data::MultiString(strings) + } + proto::RegistryValueType::RegNone => Data::Other, + } +} + +pub fn to_proto_reg_val(value: Value) -> proto::RegistryValue { + let ty = match value.ty() { + Type::Bytes => proto::RegistryValueType::RegBinary, + Type::U32 => proto::RegistryValueType::RegDword, + Type::U64 => proto::RegistryValueType::RegQword, + Type::String => proto::RegistryValueType::RegSz, + Type::ExpandString => proto::RegistryValueType::RegExpandSz, + Type::MultiString => proto::RegistryValueType::RegMultiSz, + Type::Other(_) => proto::RegistryValueType::RegNone, + }; + + let val = value.deref(); + proto::RegistryValue { + r#type: ty as i32, + data: val.to_vec(), + } +} + +impl KeyExtension for windows_registry::Key { + fn as_registry_key(&self, hive: Hive, subkey: &Path) -> proto::RegistryKey { + let values: Vec = self + .values() + .unwrap() + .map(|(name, value)| proto::RegistryKeyValue { + name, + value: Some(to_proto_reg_val(value)), + }) + .collect(); + + proto::RegistryKey { + hive: hive.to_string(), + subkey: subkey.display().to_string(), + values, + } + } + + fn value(&self, name: &str) -> Result { + let value = self.get_value(name)?; + Ok(value) + } + + fn values(&self) -> Result> { + let mut values = Vec::new(); + for value in self.values()? { + values.push(value); + } + Ok(values) + } + + fn create_value(&self, name: &str, data: Data) -> Result<()> { + match data { + Data::Bytes(val) => self.set_bytes(name, Type::Bytes, &val), + Data::DWord(val) => self.set_u32(name, val), + Data::QWord(val) => self.set_u64(name, val), + Data::String(val) => self.set_string(name, &val), + Data::ExpandString(val) => self.set_expand_string(name, &val), + Data::MultiString(val) => { + let d: Vec<&str> = val.iter().map(|s| s.as_str()).collect(); + self.set_multi_string(name, &d) + } + _ => unreachable!(), + } + } + + fn rename_value(&self, old_name: &str, new_name: &str) -> Result<()> { + let value = self.get_value(old_name)?; + self.remove_value(old_name)?; + self.set_value(new_name, &value) + } +} + +pub struct RegistryManager; + +impl RegistryManager { + pub fn value(&self, hive: Hive, subkey: &Path, name: &str) -> Result { + let key = hive.inner().open(subkey.display().to_string())?; + + key.value(name) + } + + pub fn key(&self, hive: Hive, subkey: &Path) -> Result { + hive.inner().open(subkey.display().to_string()) + } + + pub fn create_key(&self, hive: Hive, subkey: &Path) -> Result { + Key::new(hive, subkey) + } + + pub fn delete_key(&self, hive: Hive, subkey: &Path) -> Result<()> { + Key::delete(hive, subkey) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + fn test_subkey() -> PathBuf { + PathBuf::from("Software\\WineBridgeTest") + } + + #[test] + fn test_create_key() { + let hive = Hive::CurrentUser; + let subkey = test_subkey(); + assert!(Key::new(hive, &subkey).is_ok(), "Failed to create key"); + + // Check if the key exists + let key = hive.inner().open(&subkey.display().to_string()); + assert!(key.is_ok(), "Failed to open key"); + + // Clean up + hive.inner() + .remove_tree(&subkey.display().to_string()) + .expect("Failed to delete test key"); + } + + #[test] + fn test_get_key() { + let hive = Hive::CurrentUser; + let subkey = test_subkey(); + Key::new(hive, &subkey).expect("Failed to create key"); + + // Get the key + let key = Key::get(hive, &subkey); + assert!(key.is_ok(), "Failed to open key"); + + // Clean up + hive.inner() + .remove_tree(&subkey.display().to_string()) + .expect("Failed to delete test key"); + } + + #[test] + fn test_delete_key() { + let hive = Hive::CurrentUser; + let subkey = test_subkey(); + Key::new(hive, &subkey).expect("Failed to create key"); + + // Delete the key + Key::delete(hive, &subkey).expect("Failed to delete key"); + + // Check if the key is deleted + let key = hive.inner().open(&subkey.display().to_string()); + assert!(key.is_err(), "Key still exists after deletion"); + } + + #[test] + fn test_create_value() { + let hive = Hive::CurrentUser; + let subkey = test_subkey(); + + let key = Key::new(hive, &subkey).expect("Failed to create key"); + + // Set values + key.create_value("TestDWord", Data::DWord(42)) + .expect("Failed to set DWord"); + key.create_value("TestString", Data::String("hello".to_string())) + .expect("Failed to set String"); + + // Get values + let dword = key.get_u32("TestDWord").expect("Failed to get DWord"); + assert_eq!(dword, 42); + + let string = key.get_string("TestString").expect("Failed to get String"); + assert_eq!(string, "hello"); + + key.remove_value("TestDWord") + .expect("Failed to remove DWord"); + key.remove_value("TestString") + .expect("Failed to remove String"); + } + + #[test] + fn test_rename_value() { + let hive = Hive::CurrentUser; + let subkey = test_subkey(); + let key = Key::new(hive, &subkey).expect("Failed to open key"); + + key.create_value("FromDWord", Data::DWord(42)) + .expect("Failed to set DWord"); + + // Rename value + key.rename_value("FromDWord", "ToDWord") + .expect("Failed to rename value"); + let renamed = key.get_u32("ToDWord").expect("Failed to get renamed value"); + assert_eq!(renamed, 42); + + // Delete value + key.remove_value("ToDWord").expect("Failed to delete value"); + assert!(key.value("ToDWord").is_err()); + } +} diff --git a/src/registry/mod.rs b/src/registry/mod.rs new file mode 100644 index 0000000..ff8de9e --- /dev/null +++ b/src/registry/mod.rs @@ -0,0 +1 @@ +pub mod manager;