1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155
use failure::Fail; #[derive(Debug, Fail)] pub enum RegistryError { #[fail(display = "Invalid registry value type")] InvalidValueType, #[fail(display = "{}", _0)] Io(#[fail(cause)] ::std::io::Error), #[fail(display = "{}", _0)] Serialization(super::serialization::SerializationError), } impl From<::std::io::Error> for RegistryError { fn from(error: ::std::io::Error) -> RegistryError { RegistryError::Io(error) } } impl From<super::serialization::SerializationError> for RegistryError { fn from(error: super::serialization::SerializationError) -> RegistryError { RegistryError::Serialization(error) } } use super::serialization; use super::types; use winreg::enums::*; use winreg::{RegKey, RegValue}; const KEY_PATH: &'static str = "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Internet Settings\\Connections"; const KEY_PATH_WOW6432: &'static str = "SOFTWARE\\WOW6432Node\\Microsoft\\Windows\\CurrentVersion\\Internet Settings\\Connections"; pub const DEFAULT_CONNECTION_NAME: &'static str = "DefaultConnectionSettings"; pub const WINHTTP_CONNECTION_NAME: &'static str = "WinHttpSettings"; #[derive(Debug, Clone)] pub enum Target { System, CurrentUser, } #[derive(Debug, Clone)] pub struct Location { pub target: Target, pub connection_name: String, } pub fn get_current_user_location() -> Location { Location { target: Target::CurrentUser, connection_name: String::from(DEFAULT_CONNECTION_NAME), } } pub fn get_winhttp_location() -> Location { Location { target: Target::System, connection_name: String::from(WINHTTP_CONNECTION_NAME), } } fn open_key(target: &Target, write: bool, wow6432: bool) -> Result<RegKey, RegistryError> { let root_key = match target { Target::System => RegKey::predef(HKEY_LOCAL_MACHINE), Target::CurrentUser => RegKey::predef(HKEY_CURRENT_USER), }; let access = if write { KEY_ALL_ACCESS } else { KEY_READ }; let key_path = if wow6432 { KEY_PATH_WOW6432 } else { KEY_PATH }; let key = root_key.open_subkey_with_flags(key_path, access)?; return Ok(key); } fn write_raw(location: &Location, bytes: &Vec<u8>, wow6432: bool) -> Result<(), RegistryError> { let value = RegValue { vtype: REG_BINARY, bytes: bytes.to_owned(), }; let key = open_key(&location.target, true, wow6432)?; key.set_raw_value(&location.connection_name, &value)?; return Ok(()); } pub fn write_full(location: &Location, config: &types::FullConfig) -> Result<(), RegistryError> { let mut bytes = Vec::new(); serialization::serialize(config, &mut bytes)?; match location.target { Target::System => write_raw(location, &bytes, true)?, Target::CurrentUser => {} } write_raw(location, &bytes, false)?; return Ok(()); } fn read_raw(location: &Location) -> Result<Vec<u8>, RegistryError> { let key = open_key(&location.target, false, false)?; let value = key.get_raw_value(&location.connection_name)?; match value.vtype { REG_BINARY => Ok(value.bytes), _ => Err(RegistryError::InvalidValueType), } } pub fn read_full(location: &Location) -> Result<types::FullConfig, RegistryError> { let bytes = read_raw(location)?; let conf = serialization::deserialize(&bytes[..])?; return Ok(conf); } pub fn get_next_counter(location: &Location) -> u32 { let full_result = read_full(location); match full_result { Ok(full) => full.counter + 1, _ => 0, } } pub fn read(location: &Location) -> Result<types::ProxyConfig, RegistryError> { return Ok(read_full(location)?.config); } pub fn write(location: &Location, config: types::ProxyConfig) -> Result<(), RegistryError> { let full_before = read_full(location)?; let full_after = types::FullConfig { version: super::IE7_VERSION, counter: full_before.counter + 1, config, }; write_full(location, &full_after)?; Ok(()) } pub fn update<F>(location: &Location, updater: F) -> Result<(), RegistryError> where F: FnOnce(types::ProxyConfig) -> types::ProxyConfig, { let full_before = read_full(location)?; let after = updater(full_before.config); let full_after = types::FullConfig { version: super::IE7_VERSION, counter: full_before.counter + 1, config: after, }; write_full(location, &full_after)?; Ok(()) }