| // Copyright 2015 The Rust Project Developers. See the COPYRIGHT |
| // file at the top-level directory of this distribution and at |
| // http://rust-lang.org/COPYRIGHT. |
| // |
| // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or |
| // https://www.apache.org/licenses/LICENSE-2.0> or the MIT license |
| // <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your |
| // option. This file may not be copied, modified, or distributed |
| // except according to those terms. |
| |
| use std::ffi::{OsStr, OsString}; |
| use std::io; |
| use std::ops::RangeFrom; |
| use std::os::raw; |
| use std::os::windows::prelude::*; |
| |
| /// Must never be `HKEY_PERFORMANCE_DATA`. |
| pub(crate) struct RegistryKey(Repr); |
| |
| type HKEY = *mut u8; |
| type DWORD = u32; |
| type LPDWORD = *mut DWORD; |
| type LPCWSTR = *const u16; |
| type LPWSTR = *mut u16; |
| type LONG = raw::c_long; |
| type PHKEY = *mut HKEY; |
| type PFILETIME = *mut u8; |
| type LPBYTE = *mut u8; |
| type REGSAM = u32; |
| |
| const ERROR_SUCCESS: DWORD = 0; |
| const ERROR_NO_MORE_ITEMS: DWORD = 259; |
| // Sign-extend into 64 bits if needed. |
| const HKEY_LOCAL_MACHINE: HKEY = 0x80000002u32 as i32 as isize as HKEY; |
| const REG_SZ: DWORD = 1; |
| const KEY_READ: DWORD = 0x20019; |
| const KEY_WOW64_32KEY: DWORD = 0x200; |
| |
| #[link(name = "advapi32")] |
| extern "system" { |
| fn RegOpenKeyExW( |
| key: HKEY, |
| lpSubKey: LPCWSTR, |
| ulOptions: DWORD, |
| samDesired: REGSAM, |
| phkResult: PHKEY, |
| ) -> LONG; |
| fn RegEnumKeyExW( |
| key: HKEY, |
| dwIndex: DWORD, |
| lpName: LPWSTR, |
| lpcName: LPDWORD, |
| lpReserved: LPDWORD, |
| lpClass: LPWSTR, |
| lpcClass: LPDWORD, |
| lpftLastWriteTime: PFILETIME, |
| ) -> LONG; |
| fn RegQueryValueExW( |
| hKey: HKEY, |
| lpValueName: LPCWSTR, |
| lpReserved: LPDWORD, |
| lpType: LPDWORD, |
| lpData: LPBYTE, |
| lpcbData: LPDWORD, |
| ) -> LONG; |
| fn RegCloseKey(hKey: HKEY) -> LONG; |
| } |
| |
| struct OwnedKey(HKEY); |
| |
| /// Note: must not encode `HKEY_PERFORMANCE_DATA` or one of its subkeys. |
| enum Repr { |
| /// `HKEY_LOCAL_MACHINE`. |
| LocalMachine, |
| /// A subkey of `HKEY_LOCAL_MACHINE`. |
| Owned(OwnedKey), |
| } |
| |
| pub struct Iter<'a> { |
| idx: RangeFrom<DWORD>, |
| key: &'a RegistryKey, |
| } |
| |
| unsafe impl Sync for Repr {} |
| unsafe impl Send for Repr {} |
| |
| pub(crate) const LOCAL_MACHINE: RegistryKey = RegistryKey(Repr::LocalMachine); |
| |
| impl RegistryKey { |
| fn raw(&self) -> HKEY { |
| match self.0 { |
| Repr::LocalMachine => HKEY_LOCAL_MACHINE, |
| Repr::Owned(ref val) => val.0, |
| } |
| } |
| |
| /// Open a sub-key of `self`. |
| pub fn open(&self, key: &OsStr) -> io::Result<RegistryKey> { |
| let key = key.encode_wide().chain(Some(0)).collect::<Vec<_>>(); |
| let mut ret = 0 as *mut _; |
| let err = unsafe { |
| RegOpenKeyExW( |
| self.raw(), |
| key.as_ptr(), |
| 0, |
| KEY_READ | KEY_WOW64_32KEY, |
| &mut ret, |
| ) |
| }; |
| if err == ERROR_SUCCESS as LONG { |
| Ok(RegistryKey(Repr::Owned(OwnedKey(ret)))) |
| } else { |
| Err(io::Error::from_raw_os_error(err as i32)) |
| } |
| } |
| |
| pub fn iter(&self) -> Iter { |
| Iter { |
| idx: 0.., |
| key: self, |
| } |
| } |
| |
| pub fn query_str(&self, name: &str) -> io::Result<OsString> { |
| let name: &OsStr = name.as_ref(); |
| let name = name.encode_wide().chain(Some(0)).collect::<Vec<_>>(); |
| let mut len = 0; |
| let mut kind = 0; |
| unsafe { |
| let err = RegQueryValueExW( |
| self.raw(), |
| name.as_ptr(), |
| 0 as *mut _, |
| &mut kind, |
| 0 as *mut _, |
| &mut len, |
| ); |
| if err != ERROR_SUCCESS as LONG { |
| return Err(io::Error::from_raw_os_error(err as i32)); |
| } |
| if kind != REG_SZ { |
| return Err(io::Error::new( |
| io::ErrorKind::Other, |
| "registry key wasn't a string", |
| )); |
| } |
| |
| // The length here is the length in bytes, but we're using wide |
| // characters so we need to be sure to halve it for the length |
| // passed in. |
| assert!(len % 2 == 0, "impossible wide string size: {} bytes", len); |
| let vlen = len as usize / 2; |
| // Defensively initialized, see comment about |
| // `HKEY_PERFORMANCE_DATA` below. |
| let mut v = vec![0u16; vlen]; |
| let err = RegQueryValueExW( |
| self.raw(), |
| name.as_ptr(), |
| 0 as *mut _, |
| 0 as *mut _, |
| v.as_mut_ptr() as *mut _, |
| &mut len, |
| ); |
| // We don't check for `ERROR_MORE_DATA` (which would if the value |
| // grew between the first and second call to `RegQueryValueExW`), |
| // both because it's extremely unlikely, and this is a bit more |
| // defensive more defensive against weird types of registry keys. |
| if err != ERROR_SUCCESS as LONG { |
| return Err(io::Error::from_raw_os_error(err as i32)); |
| } |
| // The length is allowed to change, but should still be even, as |
| // well as smaller. |
| assert!(len % 2 == 0, "impossible wide string size: {} bytes", len); |
| // If the length grew but returned a success code, it *probably* |
| // indicates we're `HKEY_PERFORMANCE_DATA` or a subkey(?). We |
| // consider this UB, since those keys write "undefined" or |
| // "unpredictable" values to len, and need to use a completely |
| // different loop structure. This should be impossible (and enforce |
| // it in the API to the best of our ability), but to mitigate the |
| // damage we do some smoke-checks on the len, and ensure `v` has |
| // been fully initialized (rather than trusting the result of |
| // `RegQueryValueExW`). |
| let actual_len = len as usize / 2; |
| assert!(actual_len <= v.len()); |
| v.truncate(actual_len); |
| // Some registry keys may have a terminating nul character, but |
| // we're not interested in that, so chop it off if it's there. |
| if !v.is_empty() && v[v.len() - 1] == 0 { |
| v.pop(); |
| } |
| return Ok(OsString::from_wide(&v)); |
| } |
| } |
| } |
| |
| impl Drop for OwnedKey { |
| fn drop(&mut self) { |
| unsafe { |
| RegCloseKey(self.0); |
| } |
| } |
| } |
| |
| impl<'a> Iterator for Iter<'a> { |
| type Item = io::Result<OsString>; |
| |
| fn next(&mut self) -> Option<io::Result<OsString>> { |
| self.idx.next().and_then(|i| unsafe { |
| let mut v = Vec::with_capacity(256); |
| let mut len = v.capacity() as DWORD; |
| let ret = RegEnumKeyExW( |
| self.key.raw(), |
| i, |
| v.as_mut_ptr(), |
| &mut len, |
| 0 as *mut _, |
| 0 as *mut _, |
| 0 as *mut _, |
| 0 as *mut _, |
| ); |
| if ret == ERROR_NO_MORE_ITEMS as LONG { |
| None |
| } else if ret != ERROR_SUCCESS as LONG { |
| Some(Err(io::Error::from_raw_os_error(ret as i32))) |
| } else { |
| v.set_len(len as usize); |
| Some(Ok(OsString::from_wide(&v))) |
| } |
| }) |
| } |
| } |