From c44a6812b32abbbf674fa877e2955f609cb842a7 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 23 Sep 2024 20:48:17 +0800 Subject: [PATCH] feat: Add Env in context Signed-off-by: Xuanwo --- crates/reqsign/Cargo.toml | 3 + crates/reqsign/src/context.rs | 11 +++ crates/reqsign/src/env.rs | 135 ++++++++++++++++++++++++++++++++++ crates/reqsign/src/lib.rs | 2 + 4 files changed, 151 insertions(+) create mode 100644 crates/reqsign/src/env.rs diff --git a/crates/reqsign/Cargo.toml b/crates/reqsign/Cargo.toml index c9ea6a1..654f896 100644 --- a/crates/reqsign/Cargo.toml +++ b/crates/reqsign/Cargo.toml @@ -26,3 +26,6 @@ sha2.workspace = true [target.'cfg(not(target_arch = "wasm32"))'.dependencies] home.workspace = true + +[target.'cfg(target_os = "windows")'.dependencies] +windows-sys = "0.59.0" \ No newline at end of file diff --git a/crates/reqsign/src/context.rs b/crates/reqsign/src/context.rs index 4fcfb6f..e327614 100644 --- a/crates/reqsign/src/context.rs +++ b/crates/reqsign/src/context.rs @@ -1,3 +1,4 @@ +use crate::env::{Env, OsEnv}; use crate::{FileRead, HttpSend}; use anyhow::Result; use bytes::Bytes; @@ -8,17 +9,27 @@ use std::sync::Arc; pub struct Context { fs: Arc, http: Arc, + env: Arc, } impl Context { /// Create a new context. + #[inline] pub fn new(fs: impl FileRead, http: impl HttpSend) -> Self { Self { fs: Arc::new(fs), http: Arc::new(http), + env: Arc::new(OsEnv), } } + /// Set the environment for the context. Use this if you want to mock the environment. + #[inline] + pub fn with_env(mut self, env: impl Env) -> Self { + self.env = Arc::new(env); + self + } + /// Read the file content entirely in `Vec`. #[inline] pub async fn file_read(&self, path: &str) -> Result> { diff --git a/crates/reqsign/src/env.rs b/crates/reqsign/src/env.rs new file mode 100644 index 0000000..d8dc99a --- /dev/null +++ b/crates/reqsign/src/env.rs @@ -0,0 +1,135 @@ +use std::fmt::Debug; +use std::{ffi::OsString, path::PathBuf}; + +/// Permits parameterizing the home functions via the _from variants +pub trait Env: Debug + 'static { + /// Get an environment variable, as per std::env::var_os. + fn var_os(&self, key: &str) -> Option; + + /// Return the path to the users home dir, returns `None` if any error occurs. + fn home_dir(&self) -> Option; +} + +/// Implements Env for the OS context, both Unix style and Windows. +#[derive(Debug, Copy, Clone)] +pub struct OsEnv; + +impl Env for OsEnv { + fn var_os(&self, key: &str) -> Option { + std::env::var_os(key) + } + + #[cfg(any(unix, target_os = "redox"))] + fn home_dir(&self) -> Option { + #[allow(deprecated)] + std::env::home_dir() + } + + #[cfg(windows)] + fn home_dir(&self) -> Option { + windows::home_dir_inner + } +} + +/// Implements Env for the mock context. +#[cfg(test)] +#[derive(Debug, Clone)] +pub struct MockEnv { + pub home_dir: Option, + pub envs: std::collections::HashMap, +} + +#[cfg(test)] +impl Env for MockEnv { + fn var_os(&self, key: &str) -> Option { + self.envs.get(key).cloned() + } + + fn home_dir(&self) -> Option { + self.home_dir.clone() + } +} + +#[cfg(target_os = "windows")] +mod windows { + use std::env; + use std::ffi::OsString; + use std::os::windows::ffi::OsStringExt; + use std::path::PathBuf; + use std::ptr; + use std::slice; + + use windows_sys::Win32::Foundation::S_OK; + use windows_sys::Win32::System::Com::CoTaskMemFree; + use windows_sys::Win32::UI::Shell::{ + FOLDERID_Profile, SHGetKnownFolderPath, KF_FLAG_DONT_VERIFY, + }; + + pub fn home_dir_inner() -> Option { + env::var_os("USERPROFILE") + .filter(|s| !s.is_empty()) + .map(PathBuf::from) + .or_else(home_dir_crt) + } + + #[cfg(not(target_vendor = "uwp"))] + fn home_dir_crt() -> Option { + unsafe { + let mut path = ptr::null_mut(); + match SHGetKnownFolderPath( + &FOLDERID_Profile, + KF_FLAG_DONT_VERIFY as u32, + std::ptr::null_mut(), + &mut path, + ) { + S_OK => { + let path_slice = slice::from_raw_parts(path, wcslen(path)); + let s = OsString::from_wide(&path_slice); + CoTaskMemFree(path.cast()); + Some(PathBuf::from(s)) + } + _ => { + // Free any allocated memory even on failure. A null ptr is a no-op for `CoTaskMemFree`. + CoTaskMemFree(path.cast()); + None + } + } + } + } + + #[cfg(target_vendor = "uwp")] + fn home_dir_crt() -> Option { + None + } + + extern "C" { + fn wcslen(buf: *const u16) -> usize; + } + + #[cfg(not(target_vendor = "uwp"))] + #[cfg(test)] + mod tests { + use super::home_dir_inner; + use std::env; + use std::ops::Deref; + use std::path::{Path, PathBuf}; + + #[test] + fn test_with_without() { + let olduserprofile = env::var_os("USERPROFILE").unwrap(); + + env::remove_var("HOME"); + env::remove_var("USERPROFILE"); + + assert_eq!(home_dir_inner(), Some(PathBuf::from(olduserprofile))); + + let home = Path::new(r"C:\Users\foo tar baz"); + + env::set_var("HOME", home.as_os_str()); + assert_ne!(home_dir_inner().as_ref().map(Deref::deref), Some(home)); + + env::set_var("USERPROFILE", home.as_os_str()); + assert_eq!(home_dir_inner().as_ref().map(Deref::deref), Some(home)); + } + } +} diff --git a/crates/reqsign/src/lib.rs b/crates/reqsign/src/lib.rs index 6b83c21..fee1543 100644 --- a/crates/reqsign/src/lib.rs +++ b/crates/reqsign/src/lib.rs @@ -22,5 +22,7 @@ mod fs; pub use fs::FileRead; mod http; pub use http::HttpSend; +mod env; +pub use env::Env; mod context; pub use context::Context;