figment_winreg/

lib.rs

//! figment provider which allows loading data from the Windows registry.
//!
//! Features:
//! - Supports nesting (into subkeys).
//! - Supports string expansions for `REG_EXPAND_SZ` values.
//! - Can be configured to silently ignore errors while iterating over registry
//!   keys/values.

use figment::value::{Empty, Value};
use figment::{
  value::{Dict, Map},
  Error, Metadata, Profile, Provider
};

#[allow(clippy::wildcard_imports)]
use winreg::{
  enums::*,
  types::FromRegValue,
  {RegKey, HKEY}
};

pub use winreg;

#[allow(clippy::wildcard_imports)]
use windows::{
  core::*, Win32::System::Environment::ExpandEnvironmentStringsA
};

#[cfg(not(windows))]
compile_error!("This crate is for windows only");


/// Different ways to handle non-fatal errors.
#[derive(Copy, Clone)]
pub enum FailStrategy {
  /// If a non-fatal error occurs, terminate the registry iteration and return
  /// an error to the the figment front-end.
  Bail,

  /// If a non-fatal error occurs, skip the operation and continue.
  Skip
}

impl Default for FailStrategy {
  fn default() -> Self {
    Self::Bail
  }
}

/// A provider that fetches its data from a Windows registry subkey.
///
/// # String Expansion
/// This provider supports automatically expanding `REG_EXPAND_SZ` values
/// (which is enabled by default), but it has a special caveat:  The
/// expansion is performed using `ExpandEnvironmentStringsA()`, which in
/// legacy terms means that it operates on "ANSI" strings (which is
/// nonsensical).  Modern Windows versions treat the `*A()` functions as UTF-8
/// strings, and Microsoft recommend doing this for new applications.
/// However, `*A()` == UTF-8 is apparently only true if the process is running
/// with the active codepage 65001, which can be accomplished by either
/// globally configuring Windows to default codepage 65001 or setting it in
/// the executable's manifest.
pub struct RegistryProvider {
  /// The profile to emit data to if nesting is disabled.
  pub profile: Profile,

  /// The root registry key.
  hkey: HKEY,

  /// The subkey path to load from.
  regkey: String,

  /// How to handle errors while iterating through the registry's subkeys and
  /// values.
  failstrategy: FailStrategy,

  /// If `true` `REG_EXPAND_SZ` types will be expanded using
  /// `ExpandEnvironmentStrings`.  If `false` `REG_EXPAND_SZ` will be treated
  /// like `REG_SZ`
  expand: bool
}

impl RegistryProvider {
  /// Create a Windows Registry figment [`Provider`].
  ///
  /// Defaults to bailing on error and expanding `REG_EXPAND_SZ` data.
  #[allow(clippy::needless_pass_by_value)]
  pub fn new(hkey: HKEY, regkey: impl ToString) -> Self {
    Self {
      profile: Profile::Default,
      regkey: regkey.to_string(),
      hkey,
      failstrategy: FailStrategy::Bail,
      expand: true
    }
  }

  /// Configure how to handle errors occurring while processing registry
  /// keys/values/data.  By default these will cause error
  /// (`FailStrategy::Bail`), but can be configured to `FailStrategy::Skip` to
  /// silently ignore errors.
  #[must_use]
  pub const fn fail_strategy(mut self, strategy: FailStrategy) -> Self {
    self.failstrategy = strategy;
    self
  }

  /// Configure whether to expand `REG_EXPAND_SZ` or not.
  ///
  /// If set to `true` (the default) `REG_EXPAND_SZ` values will be their data
  /// processed through `ExpandEnvironmentStrings()`.
  ///
  /// It set to `false` `REG_EXPAND_SZ` will be treated like `REG_SZ`.
  #[must_use]
  pub const fn expand(mut self, expand: bool) -> Self {
    self.expand = expand;
    self
  }
}


impl Provider for RegistryProvider {
  /// Returns metadata with kind Windows Registry.
  ///
  /// If the root key is known, then it will be prefix it accordingly:
  /// `HKLM:Path\To\SubKey\Value`.
  ///
  /// If the root key is not known then it will simply ignore the root:
  /// `Path\To\SubKey\Value`.  Note that unknown root keys are unsupported.
  ///
  /// A profile will be inserted as the first part of the path:
  /// `HKLM:<Profile>\Path\To\SubKey\Value`
  fn metadata(&self) -> Metadata {
    let reg = self.regkey.clone();
    let root = rootname(self.hkey);

    Metadata::named("Windows Registry")
      .source(self.regkey.clone())
      .interpolater(move |profile, keys| {
        //println!("profile: {:?}\nkeys: {:?}", profile, keys);
        if profile.is_custom() {
          root.map_or_else(
            || format!(r"{profile}\{reg}\{}", keys.join(r"\")),
            |root| format!(r"{root}:{profile}\{reg}\{}", keys.join(r"\"))
          )
        } else {
          root.map_or_else(
            || format!(r"{reg}\{}", keys.join(r"\")),
            |root| format!(r"{root}:{reg}\{}", keys.join(r"\"))
          )
        }
      })
  }

  /// Load the requested data from the Windows registry.
  fn data(&self) -> std::result::Result<Map<Profile, Dict>, Error> {
    let hkey = match RegKey::predef(self.hkey).open_subkey(&self.regkey) {
      Ok(result) => result,
      Err(e) => match self.failstrategy {
        FailStrategy::Bail => {
          return Err(Error::from(format!("Unable to open subkey; {e}")));
        }
        FailStrategy::Skip => {
          return Ok(self.profile.collect(Dict::new()));
        }
      }
    };

    let map = self.traverse_tree(hkey, &self.regkey)?;

    Ok(self.profile.collect(map))
  }
}


impl RegistryProvider {
  #[allow(clippy::too_many_lines)]
  fn traverse_tree(
    &self,
    regkey: RegKey,
    subkey: &str
  ) -> std::result::Result<Dict, Error> {
    // we use a stack so that we can go down the registry hierarchy.
    // once we find a unexplored subkey we add the current key to the stack
    // and start exploring the subkey.
    let mut stack: Vec<(RegKey, String, Dict)> = Vec::new();

    // push the first entry onto the stack so we have a starting point
    stack.push((regkey, subkey.to_string(), Dict::new()));

    let result = 'outer: loop {
      // Note: unwrap() should be safe here since we can only en up here if
      // data has just been added.  Make sure this invariant is upheld if
      // making changes.
      let (reg, key, mut dic) = stack.pop().unwrap();

      // if we havent done this already
      // add all the subkeys to the dict with Empty values
      if dic.is_empty() {
        for ekey in reg.enum_keys() {
          let ekey = match ekey {
            Ok(ekey) => ekey,
            Err(e) => match self.failstrategy {
              FailStrategy::Bail => {
                return Err(Error::from(format!("Unable to read key {e}")));
              }
              FailStrategy::Skip => {
                continue;
              }
            }
          };
          dic.insert(ekey.clone(), Value::from(Empty::None));
        }
      }

      // look through all subkeys
      // if one is Empty add it too the top of the stack and restart the loop
      for (nkey, nval) in dic.clone() {
        if nval == Value::from(Empty::None) {
          let nreg = match reg.open_subkey(nkey.clone()) {
            Ok(v) => v,
            Err(e) => {
              return Err(Error::from(format!("Unable to open subkey. {e}")));
            }
          };
          let ndic = Dict::new();
          stack.push((reg, key, dic));
          stack.push((nreg, nkey.clone(), ndic));
          continue 'outer;
        }
      }

      // all subkeys are explored
      // grab any values in this key before we leave
      for eval in reg.enum_values() {
        let (vkey, vvalue) = match eval {
          Ok((vkey, vvalue)) => (vkey, vvalue),
          Err(e) => match self.failstrategy {
            FailStrategy::Bail => {
              return Err(Error::from(format!(
                "Unable to read key value {e}"
              )));
            }
            FailStrategy::Skip => {
              continue;
            }
          }
        };
        match vvalue.vtype {
          REG_SZ => {
            dic.insert(
              vkey,
              Value::from(match String::from_reg_value(&vvalue) {
                Ok(v) => v,
                Err(e) => match self.failstrategy {
                  FailStrategy::Bail => {
                    return Err(Error::from(format!(
                      "Unable to convert registry entry to a String value. \
                       {e}"
                    )));
                  }
                  FailStrategy::Skip => {
                    continue;
                  }
                }
              })
            );
          }
          // Strings with enviroment variables such as %USERPROFILE%
          // we take them in and expand them meaning we replace the variable
          // with its current value
          REG_EXPAND_SZ => {
            if !self.expand {
              // self.expand is false, so treat this as if it were a REG_SZ
              dic.insert(
                vkey,
                Value::from(match String::from_reg_value(&vvalue) {
                  Ok(v) => v,
                  Err(e) => match self.failstrategy {
                    FailStrategy::Bail => {
                      return Err(Error::from(format!(
                        "Unable to convert registry entry to a String value. \
                         {e}"
                      )));
                    }
                    FailStrategy::Skip => {
                      continue;
                    }
                  }
                })
              );

              // process next entry
              continue;
            }


            match w32_expand_string(
              vkey.as_ref(),
              &vvalue.to_string(),
              self.failstrategy
            )? {
              Some(s) => {
                dic.insert(vkey, Value::from(s));
              }
              None => {
                // Skip on error is enabled, and an error occurred.
                continue;
              }
            }
          }
          REG_MULTI_SZ => {
            dic.insert(
              vkey,
              Value::from(match Vec::<String>::from_reg_value(&vvalue) {
                Ok(v) => v,
                Err(e) => match self.failstrategy {
                  FailStrategy::Bail => {
                    return Err(Error::from(format!(
                      "Unable to convert registry entry to a Vec<String> \
                       value. {e}"
                    )))
                  }
                  FailStrategy::Skip => {
                    continue;
                  }
                }
              })
            );
          }
          REG_DWORD => {
            dic.insert(
              vkey,
              Value::from(match u32::from_reg_value(&vvalue) {
                Ok(v) => v,
                Err(e) => match self.failstrategy {
                  FailStrategy::Bail => {
                    return Err(Error::from(format!(
                      "Unable to convert registry entry to a u32 value. {e}"
                    )))
                  }
                  FailStrategy::Skip => {
                    continue;
                  }
                }
              })
            );
          }
          REG_QWORD => {
            dic.insert(
              vkey,
              Value::from(match u64::from_reg_value(&vvalue) {
                Ok(v) => v,
                Err(e) => match self.failstrategy {
                  FailStrategy::Bail => {
                    return Err(Error::from(format!(
                      "Unable to convert registry entry to a u64 value. {e}"
                    )))
                  }
                  FailStrategy::Skip => {
                    continue;
                  }
                }
              })
            );
          }
          REG_DWORD_BIG_ENDIAN => {
            let v = match u32::from_reg_value(&vvalue) {
              Ok(v) => v,
              Err(e) => match self.failstrategy {
                FailStrategy::Bail => {
                  return Err(Error::from(format!(
                    "Unable to convert registry entry to a u32 value. {e}"
                  )))
                }
                FailStrategy::Skip => {
                  continue;
                }
              }
            }
            .to_be();

            dic.insert(vkey, Value::from(v));
          }

          REG_NONE
          | REG_BINARY
          | REG_LINK
          | REG_RESOURCE_LIST
          | REG_FULL_RESOURCE_DESCRIPTOR
          | REG_RESOURCE_REQUIREMENTS_LIST => {
            // Ignored
            continue;
          }
        }
      }

      // Collapse this entry into the underlying stack's dictionary and go
      // back one level unless were at the end then return the current
      // dictionary
      if stack.is_empty() {
        // The stack is empty at this point which means the current node being
        // held is the root node and we've exhaused all the keys and values in
        // the registry.  Break out of loop and return the root node.
        break dic;
      }
      // safe to unwrap since we just checked and found stack to be
      // non-empty.
      let (lreg, lkey, mut ldic) = stack.pop().unwrap();
      ldic.insert(key, Value::from(dic));
      stack.push((lreg, lkey, ldic));
    };

    Ok(result)
  }
}


fn w32_expand_string(
  key: &str,
  src: &str,
  failstrategy: FailStrategy
) -> std::result::Result<Option<String>, Error> {
  // Use ExpandEnvironmentStrings() to expand the value data.
  //
  // Passing an empty slice as the destination as will cause this function to
  // calculate the size needed to store the expanded buffer.
  //
  // The return value of this function can tell us two things.
  // If the function succeeds in expanding the string it returns the length of
  // the expanded string.
  // If the buffer we send in is too small it returns the length the buffer
  // needs to be for it to succeed.
  //
  // So we start of by making it fail so we know how big of a buffer we need
  // to store the expanded string. So we send in an 0 length destination
  // buffer.
  //
  // SAFETY: The windows crate's impl IntoParam on the source string should
  //         ensure that the input string is null terminated as appropriate.
  //         We're providing a 0-length destination slice, which the windows
  //         crate wrapper will translate to a null pointer, so
  //         ExpandEnvironmentStrings() shouldn't actually try to write any
  //         data to the target buffer.
  let alloc_needed =
    unsafe { ExpandEnvironmentStringsA(PCSTR(src.as_ptr()), None) };

  if alloc_needed == 0 {
    match failstrategy {
      FailStrategy::Bail => {
        return Err(Error::from(format!(
          "Unable to get size for expand environment string buffer {key}"
        )));
      }
      FailStrategy::Skip => {
        // Just move on to next entry.
        return Ok(None);
      }
    }
  }

  // Allocate a buffer for storing the expanded string buffer.
  // alloc_needed allegedly includes the null terminator.
  let mut exp_buf = Vec::<u8>::with_capacity(alloc_needed as usize);

  // Send in the string we want to expand together with the new target buffer
  // and its size.
  //
  // This should expanded the string into the buffer and return how many
  // characters it stored (including terminating null character).
  //
  // SAFETY: We have allocated all the memory that gets modified and we only
  //         work within that buffer by passing it as a slice.
  let n = unsafe {
    exp_buf.set_len(alloc_needed as usize);

    ExpandEnvironmentStringsA(
      PCSTR(src.as_ptr()),
      Some(exp_buf.as_mut_slice())
    )
  };

  if n == 0 {
    match failstrategy {
      FailStrategy::Bail => {
        return Err(Error::from(format!(
          "Unable to expand environment string for value {key}"
        )));
      }
      FailStrategy::Skip => {
        // Just move on to next entry.
        return Ok(None);
      }
    }
  }

  // Set the length of the vector to the length of the returned
  // buffer, but remove one byte as it should be the terminating
  // null.
  unsafe { exp_buf.set_len((n - 1) as usize) };

  // Attempt to create a regular String from the Vec<u8>.
  let Ok(exp_str) = String::from_utf8(exp_buf) else {
    match failstrategy {
      FailStrategy::Bail => {
        return Err(Error::from(format!(
          "Unable to expand environment string for value {key}"
        )));
      }
      FailStrategy::Skip => {
        // Just move on to next entry.
        return Ok(None);
      }
    }
  };


  Ok(Some(exp_str))
}


const fn rootname(hkey: HKEY) -> Option<&'static str> {
  match hkey {
    HKEY_CLASSES_ROOT => Some("HKCR"),
    HKEY_CURRENT_USER => Some("HKCU"),
    HKEY_LOCAL_MACHINE => Some("HKLM"),
    HKEY_USERS => Some("HKEY_USERS"),
    HKEY_CURRENT_CONFIG => Some("HKCC"),
    _ => None
  }
}

// vim: set ft=rust et sw=2 ts=2 sts=2 cinoptions=2 tw=79 :