blob: cae32219c7fb41fe5768ebec2c5e88242ef26ce5 [file] [log] [blame]
// 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)))
}
})
}
}