vendor/cc-1.0.79/src/registry.rs

   1 // Copyright 2015 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 use std::ffi::{OsStr, OsString};
  12 use std::io;
  13 use std::ops::RangeFrom;
  14 use std::os::raw;
  15 use std::os::windows::prelude::*;
  16
  17 /// Must never be `HKEY_PERFORMANCE_DATA`.
  18 pub(crate) struct RegistryKey(Repr);
  19
  20 type HKEY = *mut u8;
  21 type DWORD = u32;
  22 type LPDWORD = *mut DWORD;
  23 type LPCWSTR = *const u16;
  24 type LPWSTR = *mut u16;
  25 type LONG = raw::c_long;
  26 type PHKEY = *mut HKEY;
  27 type PFILETIME = *mut u8;
  28 type LPBYTE = *mut u8;
  29 type REGSAM = u32;
  30
  31 const ERROR_SUCCESS: DWORD = 0;
  32 const ERROR_NO_MORE_ITEMS: DWORD = 259;
  33 // Sign-extend into 64 bits if needed.
  34 const HKEY_LOCAL_MACHINE: HKEY = 0x80000002u32 as i32 as isize as HKEY;
  35 const REG_SZ: DWORD = 1;
  36 const KEY_READ: DWORD = 0x20019;
  37 const KEY_WOW64_32KEY: DWORD = 0x200;
  38
  39 #[link(name = "advapi32")]
  40 extern "system" {
  41     fn RegOpenKeyExW(
  42         key: HKEY,
  43         lpSubKey: LPCWSTR,
  44         ulOptions: DWORD,
  45         samDesired: REGSAM,
  46         phkResult: PHKEY,
  47     ) -> LONG;
  48     fn RegEnumKeyExW(
  49         key: HKEY,
  50         dwIndex: DWORD,
  51         lpName: LPWSTR,
  52         lpcName: LPDWORD,
  53         lpReserved: LPDWORD,
  54         lpClass: LPWSTR,
  55         lpcClass: LPDWORD,
  56         lpftLastWriteTime: PFILETIME,
  57     ) -> LONG;
  58     fn RegQueryValueExW(
  59         hKey: HKEY,
  60         lpValueName: LPCWSTR,
  61         lpReserved: LPDWORD,
  62         lpType: LPDWORD,
  63         lpData: LPBYTE,
  64         lpcbData: LPDWORD,
  65     ) -> LONG;
  66     fn RegCloseKey(hKey: HKEY) -> LONG;
  67 }
  68
  69 struct OwnedKey(HKEY);
  70
  71 /// Note: must not encode `HKEY_PERFORMANCE_DATA` or one of its subkeys.
  72 enum Repr {
  73     /// `HKEY_LOCAL_MACHINE`.
  74     LocalMachine,
  75     /// A subkey of `HKEY_LOCAL_MACHINE`.
  76     Owned(OwnedKey),
  77 }
  78
  79 pub struct Iter<'a> {
  80     idx: RangeFrom<DWORD>,
  81     key: &'a RegistryKey,
  82 }
  83
  84 unsafe impl Sync for Repr {}
  85 unsafe impl Send for Repr {}
  86
  87 pub(crate) const LOCAL_MACHINE: RegistryKey = RegistryKey(Repr::LocalMachine);
  88
  89 impl RegistryKey {
  90     fn raw(&self) -> HKEY {
  91         match self.0 {
  92             Repr::LocalMachine => HKEY_LOCAL_MACHINE,
  93             Repr::Owned(ref val) => val.0,
  94         }
  95     }
  96
  97     /// Open a sub-key of `self`.
  98     pub fn open(&self, key: &OsStr) -> io::Result<RegistryKey> {
  99         let key = key.encode_wide().chain(Some(0)).collect::<Vec<_>>();
 100         let mut ret = 0 as *mut _;
 101         let err = unsafe {
 102             RegOpenKeyExW(
 103                 self.raw(),
 104                 key.as_ptr(),
 105                 0,
 106                 KEY_READ | KEY_WOW64_32KEY,
 107                 &mut ret,
 108             )
 109         };
 110         if err == ERROR_SUCCESS as LONG {
 111             Ok(RegistryKey(Repr::Owned(OwnedKey(ret))))
 112         } else {
 113             Err(io::Error::from_raw_os_error(err as i32))
 114         }
 115     }
 116
 117     pub fn iter(&self) -> Iter {
 118         Iter {
 119             idx: 0..,
 120             key: self,
 121         }
 122     }
 123
 124     pub fn query_str(&self, name: &str) -> io::Result<OsString> {
 125         let name: &OsStr = name.as_ref();
 126         let name = name.encode_wide().chain(Some(0)).collect::<Vec<_>>();
 127         let mut len = 0;
 128         let mut kind = 0;
 129         unsafe {
 130             let err = RegQueryValueExW(
 131                 self.raw(),
 132                 name.as_ptr(),
 133                 0 as *mut _,
 134                 &mut kind,
 135                 0 as *mut _,
 136                 &mut len,
 137             );
 138             if err != ERROR_SUCCESS as LONG {
 139                 return Err(io::Error::from_raw_os_error(err as i32));
 140             }
 141             if kind != REG_SZ {
 142                 return Err(io::Error::new(
 143                     io::ErrorKind::Other,
 144                     "registry key wasn't a string",
 145                 ));
 146             }
 147
 148             // The length here is the length in bytes, but we're using wide
 149             // characters so we need to be sure to halve it for the length
 150             // passed in.
 151             assert!(len % 2 == 0, "impossible wide string size: {} bytes", len);
 152             let vlen = len as usize / 2;
 153             // Defensively initialized, see comment about
 154             // `HKEY_PERFORMANCE_DATA` below.
 155             let mut v = vec![0u16; vlen];
 156             let err = RegQueryValueExW(
 157                 self.raw(),
 158                 name.as_ptr(),
 159                 0 as *mut _,
 160                 0 as *mut _,
 161                 v.as_mut_ptr() as *mut _,
 162                 &mut len,
 163             );
 164             // We don't check for `ERROR_MORE_DATA` (which would if the value
 165             // grew between the first and second call to `RegQueryValueExW`),
 166             // both because it's extremely unlikely, and this is a bit more
 167             // defensive more defensive against weird types of registry keys.
 168             if err != ERROR_SUCCESS as LONG {
 169                 return Err(io::Error::from_raw_os_error(err as i32));
 170             }
 171             // The length is allowed to change, but should still be even, as
 172             // well as smaller.
 173             assert!(len % 2 == 0, "impossible wide string size: {} bytes", len);
 174             // If the length grew but returned a success code, it *probably*
 175             // indicates we're `HKEY_PERFORMANCE_DATA` or a subkey(?). We
 176             // consider this UB, since those keys write "undefined" or
 177             // "unpredictable" values to len, and need to use a completely
 178             // different loop structure. This should be impossible (and enforce
 179             // it in the API to the best of our ability), but to mitigate the
 180             // damage we do some smoke-checks on the len, and ensure `v` has
 181             // been fully initialized (rather than trusting the result of
 182             // `RegQueryValueExW`).
 183             let actual_len = len as usize / 2;
 184             assert!(actual_len <= v.len());
 185             v.truncate(actual_len);
 186             // Some registry keys may have a terminating nul character, but
 187             // we're not interested in that, so chop it off if it's there.
 188             if !v.is_empty() && v[v.len() - 1] == 0 {
 189                 v.pop();
 190             }
 191             return Ok(OsString::from_wide(&v));
 192         }
 193     }
 194 }
 195
 196 impl Drop for OwnedKey {
 197     fn drop(&mut self) {
 198         unsafe {
 199             RegCloseKey(self.0);
 200         }
 201     }
 202 }
 203
 204 impl<'a> Iterator for Iter<'a> {
 205     type Item = io::Result<OsString>;
 206
 207     fn next(&mut self) -> Option<io::Result<OsString>> {
 208         self.idx.next().and_then(|i| unsafe {
 209             let mut v = Vec::with_capacity(256);
 210             let mut len = v.capacity() as DWORD;
 211             let ret = RegEnumKeyExW(
 212                 self.key.raw(),
 213                 i,
 214                 v.as_mut_ptr(),
 215                 &mut len,
 216                 0 as *mut _,
 217                 0 as *mut _,
 218                 0 as *mut _,
 219                 0 as *mut _,
 220             );
 221             if ret == ERROR_NO_MORE_ITEMS as LONG {
 222                 None
 223             } else if ret != ERROR_SUCCESS as LONG {
 224                 Some(Err(io::Error::from_raw_os_error(ret as i32)))
 225             } else {
 226                 v.set_len(len as usize);
 227                 Some(Ok(OsString::from_wide(&v)))
 228             }
 229         })
 230     }
 231 }