Stable, safe, fast wrapper for RegQueryValueEx

RegQueryValueEx.rs

/* common uses */
pub use std::ptr::{self, null, null_mut};

pub use widestring::*;
pub use winapi::shared::minwindef::*;
pub use winapi::shared::winerror::*;
pub use winapi::um::errhandlingapi::*;
pub use winapi::um::winbase::*;
pub use winapi::um::winnt::*;
pub use winapi::um::winreg::*;
/* end of common uses */

// Custom wrapper `WinError` for `GetLastError()` and `WinResult` type.
pub use crate::error::*;

const INITIAL_BUFFER_SIZE: usize = 8 * 1024;

/// Create new buffer and call `query_value_with_buffer`.
pub fn query_value(
    hkey: HKEY,
    value_name: &str,
    buffer_size_hint: Option<usize>
) -> WinResult<Vec<u8>>
{
    let mut buffer = Vec::new();
    query_value_with_buffer(
        hkey,
        value_name,
        buffer_size_hint,
        &mut buffer,
    )?;
    Ok(buffer)
}

/// Query registry value of potentially unknown size, reallocating larger buffer in a loop as needed.
/// Given buffer will be cleared and overridden with zeroes before usage.
/// # Panics
/// Will panic if value_name contains NULL characters.
pub fn query_value_with_buffer(
    hkey: HKEY,
    value_name: &str,
    mut buffer_size_hint: Option<usize>,
    buffer: &mut Vec<u8>
) -> WinResult<()>
{
    // prepare value name with trailing NULL char
    let wsz_value_name = U16CString::from_str(value_name).unwrap();

    // start with some non-zero size, even if explicit zero were provided, and gradually
    // increment it until value fits into buffer.
    // some queries may return undefined size and ERROR_MORE_DATA status when they don't know
    // the data size in advance.
    let mut buffer_size = match buffer_size_hint {
        Some(0) => {
            eprintln!("Explicit Some(0) size hint provided. Use None instead.");
            INITIAL_BUFFER_SIZE
        }
        Some(hint) => hint,
        None => {
            match try_get_size_hint(hkey, value_name, wsz_value_name.as_ref()) {
                Ok(Some(hint)) => hint as usize,
                Ok(None) => INITIAL_BUFFER_SIZE,
                // gracefully fallback to incremental buffer allocation, do not return error here.
                Err(why) => {
                    eprintln!("{}", why);
                    INITIAL_BUFFER_SIZE
                }
            }
        }
    };
    // From MSDN:
    // If hKey specifies HKEY_PERFORMANCE_DATA and the lpData buffer is not large enough to
    // contain all of the returned data, RegQueryValueEx returns ERROR_MORE_DATA and the value
    // returned through the lpcbData parameter is undefined.
    // [..]
    // You need to maintain a separate variable to keep track of the buffer size, because the
    // value returned by lpcbData is unpredictable.
    let mut buffer_size_out = buffer_size as DWORD;
    // buffer initialization
    buffer.clear();
    buffer.reserve(buffer_size);

    let mut error_code: DWORD = ERROR_SUCCESS;
    unsafe {
        error_code = RegQueryValueExW(
            hkey,
            wsz_value_name.as_ptr(),
            null_mut(),
            null_mut(),
            buffer.as_mut_ptr(),
            &mut buffer_size_out as LPDWORD,
        ) as DWORD;

        while error_code == ERROR_MORE_DATA {
            // initialize buffer size or double its value
            let increment = if buffer_size == 0 { INITIAL_BUFFER_SIZE } else { buffer_size };
            buffer_size += increment;
            buffer_size_out = buffer_size as DWORD;
            // buffer considers itself empty, so reversing for "additional" N items is the same as
            // reserving for total of N items.
            buffer.reserve(buffer_size);

            // exactly same call as above
            error_code = RegQueryValueExW(
                hkey,
                wsz_value_name.as_ptr(),
                null_mut(),
                null_mut(),
                buffer.as_mut_ptr(),
                &mut buffer_size_out as LPDWORD,
            ) as DWORD;
        }
    }

    if error_code != ERROR_SUCCESS {
        return Err(WinError::new_with_message(error_code).with_comment(format!("RegQueryValueExW with query: {}", value_name)));
    }

    // SAFETY: buffer_size_out is initialized to a valid value by a successful call to RegQueryValueExW
    unsafe { buffer.set_len(buffer_size_out as usize) };
    Ok(())
}

/// in cases when value size is known in advance, we may __try__ to use that as a size hint.
/// but it won't work with dynamic values, such as counters data; and it must not be trusted
/// as a final size, because anything could happen between two calls to RegQueryValueExW.
///
/// it certainly can NOT be used under these conditions:
/// - HKEY is HKEY_PERFORMANCE_DATA but
/// - value_name is not starting with neither "Counter" nor "Help".
fn try_get_size_hint(hkey: HKEY, value_name: &str, wsz_value_name: &U16CStr) -> WinResult<Option<DWORD>> {
    let can_not_use_reg_size_hint =
        (hkey == HKEY_PERFORMANCE_DATA)
            && (!value_name.starts_with("Counter")
            && !value_name.starts_with("Help"));

    if can_not_use_reg_size_hint {
        return Ok(None);
    }

    let mut reg_size_hint: DWORD = 0;
    // pass NULL data to figure out needed buffer size
    let error_code = unsafe {
        RegQueryValueExW(
            hkey,
            wsz_value_name.as_ptr(),
            null_mut(),
            null_mut(),
            null_mut(),
            &mut reg_size_hint as LPDWORD,
        ) as DWORD
    };

    if error_code != ERROR_SUCCESS {
        return Err(WinError::new_with_message(error_code).with_comment(format!(
            "Getting buffer size hint for registry value failed. \
            This should not happen for HKEY = {:p}, ValueName = {:?}", hkey, value_name)));
    }

    Ok(Some(reg_size_hint))
}

src/error.rs:

use std::error::Error;
use std::fmt;

/* common uses */
pub use std::ptr::{self, null, null_mut};

pub use widestring::*;
pub use winapi::shared::minwindef::*;
pub use winapi::shared::winerror::*;
pub use winapi::um::errhandlingapi::*;
pub use winapi::um::winbase::*;
pub use winapi::um::winnt::*;
pub use winapi::um::winreg::*;
/* end of common uses */

#[derive(Debug, Clone)]
pub struct WinError {
    error_code: DWORD,
    comment: Option<String>,
    message: Option<String>,
    source: Option<Box<WinError>>,
}

impl WinError {
    fn _new(error_code: DWORD, comment: Option<String>, message: Option<String>, source: Option<Box<WinError>>) -> Self {
        WinError {
            error_code,
            comment,
            message,
            source,
        }
    }

    pub fn new(error_code: DWORD) -> Self {
        Self::_new(error_code, None, None, None)
    }

    pub fn new_with_message(error_code: DWORD) -> Self {
        Self::_new(error_code, None, None, None).with_message()
    }

    /// Call `GetLastError()` but do not attempt to get formatted message from system.
    pub fn get() -> Self {
        let error_code = Self::get_last_error();
        Self::new(error_code)
    }

    pub fn with_comment<S: Into<String>>(&self, comment: S) -> Self {
        let mut clone = self.clone();
        clone.comment = Some(comment.into());
        clone
    }

    /// If formatted message is not initialized, get one via `FormatMessage(...)` and return new error instance.
    pub fn with_message(&self) -> Self {
        match self.message.as_ref() {
            Some(_) => self.clone(),
            None => self.clone_with_message()
        }
    }

    pub fn with_source(&self, source: Self) -> Self {
        let mut clone = self.clone();
        clone.source = Some(Box::new(source));
        clone
    }

    /// Call `GetLastError()` & `FormatMessage(...)` at once.
    pub fn get_with_message() -> Self {
        Self::get().with_message()
    }

    /// Call `FormatMessage(...)` for given error code.
    /// If `FormatMessage(...)` fails, create new error for its status code wrapping original one as a source.
    fn clone_with_message(&self) -> Self {
        match Self::format_message_from_error_code(self.error_code) {
            Ok(message) => {
                let mut clone = self.clone();
                clone.message = Some(message);
                clone
            }
            Err(format_error) => {
                let message = Self::get_format_message_error(self.error_code);
                Self::_new(format_error, None, Some(message), Some(Box::new(self.clone())))
            }
        }
    }

    // getter
    fn error_code(&self) -> DWORD {
        self.error_code
    }

    #[inline(always)]
    fn get_last_error() -> DWORD {
        unsafe { GetLastError() }
    }

    fn format_message_from_error_code(error_code: DWORD) -> Result<String, DWORD> {
        unsafe {
            let mut buffer: LPWSTR = null_mut();
            // If the function succeeds, the return value is the number of TCHARs stored in the output buffer, excluding the terminating null character.
            let len = FormatMessageW(
                FORMAT_MESSAGE_IGNORE_INSERTS
                    | FORMAT_MESSAGE_FROM_SYSTEM
                    | FORMAT_MESSAGE_ALLOCATE_BUFFER, // dwFlags
                null(), // lpSource
                error_code, // dwMessageId
                MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT) as DWORD, // dwLanguageId
                ::std::mem::transmute(&mut buffer as *mut LPWSTR),  // lpBuffer
                0, // nSize
                null_mut(), // va_args
            );

            // If the function fails, the return value is zero. To get extended error information, call GetLastError.
            if len == 0 {
                return Err(Self::get_last_error());
            }

            let message_u16 = U16Str::from_ptr(buffer, len as usize);
            let message_string = message_u16.to_string_lossy();

            if buffer != null_mut() {
                LocalFree(buffer as HLOCAL);
            }

            Ok(message_string)
        }
    }

    fn get_format_message_error(original_error_code: DWORD) -> String {
        format!("FormatMessageW failed while formatting error 0x{:08X}", original_error_code)
    }

    const UNKNOWN_ERROR: &'static str = "UNKNOWN ERROR CODE";
}

impl fmt::Display for WinError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let Some(comment) = self.comment.as_ref() {
            write!(f, "{}; ", comment)?;
        }
        let message = match self.message.as_ref() {
            Some(msg) => msg.as_str(),
            None => Self::UNKNOWN_ERROR,
        };
        write!(f, "Error Code 0x{:08X}: {}", self.error_code, message)?;

        if let Some(source) = self.source.as_ref() {
            write!(f, "; Caused by: {}", source)?;
        }
        Ok(())
    }
}

impl Error for WinError {
    fn description(&self) -> &str {
        unimplemented!()
    }
}

/// Rust + Windows extension for error handling
pub type WinResult<T> = Result<T, WinError>;

Cargo.toml:

[package]
name = "windows-rust-counters"
version = "0.1.0"
authors = ["..."]
edition = "2018"

[dependencies]
widestring = "0.4"

[target.'cfg(windows)'.dependencies.winapi]
version = "0.3"
features = [
    "errhandlingapi",
    "minwindef",
    "ntdef",
    "pdh",
    "processenv",
    "sysinfoapi",
    "vadefs",
    "winbase",
    "winnls",
    "winnt",
    "winreg",
    "winsvc",
]