]>
Commit | Line | Data |
---|---|---|
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 | // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license | |
7 | // <LICENSE-MIT or http://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::{OsString, OsStr}; | |
12 | use std::io; | |
13 | use std::ops::RangeFrom; | |
14 | use std::os::raw; | |
15 | use std::os::windows::prelude::*; | |
16 | ||
17 | pub struct RegistryKey(Repr); | |
18 | ||
19 | type HKEY = *mut u8; | |
20 | type DWORD = u32; | |
21 | type LPDWORD = *mut DWORD; | |
22 | type LPCWSTR = *const u16; | |
23 | type LPWSTR = *mut u16; | |
24 | type LONG = raw::c_long; | |
25 | type PHKEY = *mut HKEY; | |
26 | type PFILETIME = *mut u8; | |
27 | type LPBYTE = *mut u8; | |
28 | type REGSAM = u32; | |
29 | ||
30 | const ERROR_SUCCESS: DWORD = 0; | |
31 | const ERROR_NO_MORE_ITEMS: DWORD = 259; | |
32 | const HKEY_LOCAL_MACHINE: HKEY = 0x80000002 as HKEY; | |
33 | const REG_SZ: DWORD = 1; | |
34 | const KEY_READ: DWORD = 0x20019; | |
35 | const KEY_WOW64_32KEY: DWORD = 0x200; | |
36 | ||
37 | #[link(name = "advapi32")] | |
38 | extern "system" { | |
39 | fn RegOpenKeyExW(key: HKEY, | |
40 | lpSubKey: LPCWSTR, | |
41 | ulOptions: DWORD, | |
42 | samDesired: REGSAM, | |
43 | phkResult: PHKEY) | |
44 | -> LONG; | |
45 | fn RegEnumKeyExW(key: HKEY, | |
46 | dwIndex: DWORD, | |
47 | lpName: LPWSTR, | |
48 | lpcName: LPDWORD, | |
49 | lpReserved: LPDWORD, | |
50 | lpClass: LPWSTR, | |
51 | lpcClass: LPDWORD, | |
52 | lpftLastWriteTime: PFILETIME) | |
53 | -> LONG; | |
54 | fn RegQueryValueExW(hKey: HKEY, | |
55 | lpValueName: LPCWSTR, | |
56 | lpReserved: LPDWORD, | |
57 | lpType: LPDWORD, | |
58 | lpData: LPBYTE, | |
59 | lpcbData: LPDWORD) | |
60 | -> LONG; | |
61 | fn RegCloseKey(hKey: HKEY) -> LONG; | |
62 | } | |
63 | ||
64 | struct OwnedKey(HKEY); | |
65 | ||
66 | enum Repr { | |
67 | Const(HKEY), | |
68 | Owned(OwnedKey), | |
69 | } | |
70 | ||
71 | pub struct Iter<'a> { | |
72 | idx: RangeFrom<DWORD>, | |
73 | key: &'a RegistryKey, | |
74 | } | |
75 | ||
76 | unsafe impl Sync for Repr {} | |
77 | unsafe impl Send for Repr {} | |
78 | ||
79 | pub static LOCAL_MACHINE: RegistryKey = RegistryKey(Repr::Const(HKEY_LOCAL_MACHINE)); | |
80 | ||
81 | impl RegistryKey { | |
82 | fn raw(&self) -> HKEY { | |
83 | match self.0 { | |
84 | Repr::Const(val) => val, | |
85 | Repr::Owned(ref val) => val.0, | |
86 | } | |
87 | } | |
88 | ||
89 | pub fn open(&self, key: &OsStr) -> io::Result<RegistryKey> { | |
90 | let key = key.encode_wide().chain(Some(0)).collect::<Vec<_>>(); | |
91 | let mut ret = 0 as *mut _; | |
92 | let err = unsafe { | |
93 | RegOpenKeyExW(self.raw(), | |
94 | key.as_ptr(), | |
95 | 0, | |
96 | KEY_READ | KEY_WOW64_32KEY, | |
97 | &mut ret) | |
98 | }; | |
99 | if err == ERROR_SUCCESS as LONG { | |
100 | Ok(RegistryKey(Repr::Owned(OwnedKey(ret)))) | |
101 | } else { | |
102 | Err(io::Error::from_raw_os_error(err as i32)) | |
103 | } | |
104 | } | |
105 | ||
106 | pub fn iter(&self) -> Iter { | |
107 | Iter { | |
108 | idx: 0.., | |
109 | key: self, | |
110 | } | |
111 | } | |
112 | ||
113 | pub fn query_str(&self, name: &str) -> io::Result<OsString> { | |
114 | let name: &OsStr = name.as_ref(); | |
115 | let name = name.encode_wide().chain(Some(0)).collect::<Vec<_>>(); | |
116 | let mut len = 0; | |
117 | let mut kind = 0; | |
118 | unsafe { | |
119 | let err = RegQueryValueExW(self.raw(), | |
120 | name.as_ptr(), | |
121 | 0 as *mut _, | |
122 | &mut kind, | |
123 | 0 as *mut _, | |
124 | &mut len); | |
125 | if err != ERROR_SUCCESS as LONG { | |
126 | return Err(io::Error::from_raw_os_error(err as i32)); | |
127 | } | |
128 | if kind != REG_SZ { | |
129 | return Err(io::Error::new(io::ErrorKind::Other, "registry key wasn't a string")); | |
130 | } | |
131 | ||
132 | // The length here is the length in bytes, but we're using wide | |
133 | // characters so we need to be sure to halve it for the capacity | |
134 | // passed in. | |
135 | let mut v = Vec::with_capacity(len as usize / 2); | |
136 | let err = RegQueryValueExW(self.raw(), | |
137 | name.as_ptr(), | |
138 | 0 as *mut _, | |
139 | 0 as *mut _, | |
140 | v.as_mut_ptr() as *mut _, | |
141 | &mut len); | |
142 | if err != ERROR_SUCCESS as LONG { | |
143 | return Err(io::Error::from_raw_os_error(err as i32)); | |
144 | } | |
145 | v.set_len(len as usize / 2); | |
146 | ||
147 | // Some registry keys may have a terminating nul character, but | |
148 | // we're not interested in that, so chop it off if it's there. | |
149 | if v[v.len() - 1] == 0 { | |
150 | v.pop(); | |
151 | } | |
152 | Ok(OsString::from_wide(&v)) | |
153 | } | |
154 | } | |
155 | } | |
156 | ||
157 | impl Drop for OwnedKey { | |
158 | fn drop(&mut self) { | |
159 | unsafe { | |
160 | RegCloseKey(self.0); | |
161 | } | |
162 | } | |
163 | } | |
164 | ||
165 | impl<'a> Iterator for Iter<'a> { | |
166 | type Item = io::Result<OsString>; | |
167 | ||
168 | fn next(&mut self) -> Option<io::Result<OsString>> { | |
169 | self.idx.next().and_then(|i| unsafe { | |
170 | let mut v = Vec::with_capacity(256); | |
171 | let mut len = v.capacity() as DWORD; | |
172 | let ret = RegEnumKeyExW(self.key.raw(), | |
173 | i, | |
174 | v.as_mut_ptr(), | |
175 | &mut len, | |
176 | 0 as *mut _, | |
177 | 0 as *mut _, | |
178 | 0 as *mut _, | |
179 | 0 as *mut _); | |
180 | if ret == ERROR_NO_MORE_ITEMS as LONG { | |
181 | None | |
182 | } else if ret != ERROR_SUCCESS as LONG { | |
183 | Some(Err(io::Error::from_raw_os_error(ret as i32))) | |
184 | } else { | |
185 | v.set_len(len as usize); | |
186 | Some(Ok(OsString::from_wide(&v))) | |
187 | } | |
188 | }) | |
189 | } | |
190 | } |