winctx/
registry.rs

1use std::ffi::{OsStr, OsString};
2use std::io;
3use std::mem::MaybeUninit;
4use std::ptr;
5
6use windows_sys::Win32::Foundation::ERROR_SUCCESS;
7use windows_sys::Win32::System::Registry::{self as winreg, HKEY};
8
9use crate::convert::{FromWide, ToWide};
10
11/// An open registry key.
12///
13/// This is constructed using [`OpenRegistryKey`].
14pub struct RegistryKey(winreg::HKEY);
15
16unsafe impl Sync for RegistryKey {}
17unsafe impl Send for RegistryKey {}
18
19/// Helper to open a registry key with the ability to specify desired
20/// permissions.
21pub struct OpenRegistryKey {
22    key: HKEY,
23    desired: u32,
24}
25
26impl OpenRegistryKey {
27    /// Open the given key in the `HKEY_CURRENT_USER` registry.
28    pub fn current_user() -> Self {
29        Self {
30            key: winreg::HKEY_CURRENT_USER,
31            desired: winreg::KEY_READ | winreg::KEY_WOW64_64KEY,
32        }
33    }
34
35    /// Open the given key in the `HKEY_LOCAL_MACHINE` registry.
36    pub fn local_machine() -> Self {
37        Self {
38            key: winreg::HKEY_LOCAL_MACHINE,
39            desired: winreg::KEY_READ,
40        }
41    }
42
43    /// Enable the `KEY_SET_VALUE` desired access mode.
44    pub fn set_value(mut self) -> Self {
45        self.desired |= winreg::KEY_SET_VALUE;
46        self
47    }
48
49    /// Internal open implementation.
50    pub fn open<K>(self, key: K) -> io::Result<RegistryKey>
51    where
52        K: AsRef<OsStr>,
53    {
54        let key = key.to_wide_null();
55        self.open_inner(&key)
56    }
57
58    fn open_inner(&self, key: &[u16]) -> io::Result<RegistryKey> {
59        unsafe {
60            let mut hkey = MaybeUninit::uninit();
61
62            let status =
63                winreg::RegOpenKeyExW(self.key, key.as_ptr(), 0, self.desired, hkey.as_mut_ptr());
64
65            if status != ERROR_SUCCESS {
66                return Err(io::Error::from_raw_os_error(status as i32));
67            }
68
69            Ok(RegistryKey(hkey.assume_init()))
70        }
71    }
72}
73
74impl RegistryKey {
75    /// Open the given key in the HKEY_CURRENT_USER scope.
76    #[inline]
77    pub fn current_user<K>(key: K) -> io::Result<RegistryKey>
78    where
79        K: AsRef<OsStr>,
80    {
81        OpenRegistryKey::current_user().open(key)
82    }
83
84    /// Open the given key in the HKEY_LOCAL_MACHINE scope.
85    pub fn local_machine<K>(key: K) -> io::Result<RegistryKey>
86    where
87        K: AsRef<OsStr>,
88    {
89        OpenRegistryKey::local_machine().open(key)
90    }
91
92    /// Get the given value as a string.
93    pub fn get_string<N>(&self, name: N) -> io::Result<OsString>
94    where
95        N: AsRef<OsStr>,
96    {
97        let name = name.to_wide_null();
98        let bytes = self.get_wide(&name, winreg::RRF_RT_REG_SZ)?;
99        // Skip the terminating null.
100        Ok(OsString::from_wide(&bytes[..bytes.len().saturating_sub(1)]))
101    }
102
103    fn get_wide(&self, name: &[u16], flags: u32) -> io::Result<Vec<u16>> {
104        let mut len = 0;
105
106        unsafe {
107            let status = winreg::RegGetValueW(
108                self.0,
109                ptr::null_mut(),
110                name.as_ptr(),
111                flags,
112                ptr::null_mut(),
113                ptr::null_mut(),
114                &mut len,
115            );
116
117            if status != ERROR_SUCCESS {
118                return Err(io::Error::from_raw_os_error(status as i32));
119            }
120
121            debug_assert!(len % 2 == 0);
122            let mut value = vec![0u16; (len / 2) as usize];
123
124            let status = winreg::RegGetValueW(
125                self.0,
126                ptr::null_mut(),
127                name.as_ptr(),
128                flags,
129                ptr::null_mut(),
130                value.as_mut_ptr().cast(),
131                &mut len,
132            );
133
134            if status != ERROR_SUCCESS {
135                return Err(io::Error::from_raw_os_error(status as i32));
136            }
137
138            debug_assert!(len % 2 == 0);
139
140            // Length reported *including* wide null terminator.
141            value.truncate((len / 2) as usize);
142            Ok(value)
143        }
144    }
145
146    /// Set the given value.
147    pub fn set<N>(&self, name: N, value: impl AsRef<OsStr>) -> io::Result<()>
148    where
149        N: AsRef<OsStr>,
150    {
151        let name = name.to_wide_null();
152        let value = value.to_wide_null();
153        self.set_inner(&name, &value)
154    }
155
156    fn set_inner(&self, name: &[u16], value: &[u16]) -> io::Result<()> {
157        let value_len = value
158            .len()
159            .checked_mul(2)
160            .and_then(|n| u32::try_from(n).ok())
161            .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "Value size overflow"))?;
162
163        let status = unsafe {
164            winreg::RegSetValueExW(
165                self.0,
166                name.as_ptr(),
167                0,
168                winreg::REG_SZ,
169                value.as_ptr().cast(),
170                value_len,
171            )
172        };
173
174        if status != 0 {
175            return Err(io::Error::last_os_error());
176        }
177
178        Ok(())
179    }
180
181    /// Delete the given registry key.
182    pub fn delete<N>(&self, name: N) -> io::Result<()>
183    where
184        N: AsRef<OsStr>,
185    {
186        let name = name.to_wide_null();
187        self.delete_inner(&name)
188    }
189
190    fn delete_inner(&self, name: &[u16]) -> io::Result<()> {
191        let status = unsafe { winreg::RegDeleteKeyValueW(self.0, ptr::null_mut(), name.as_ptr()) };
192
193        if status != 0 {
194            return Err(io::Error::last_os_error());
195        }
196
197        Ok(())
198    }
199}
200
201impl Drop for RegistryKey {
202    fn drop(&mut self) {
203        unsafe {
204            winreg::RegCloseKey(self.0);
205        }
206    }
207}