Skip to content

Commit

Permalink
feat: Add Env in context
Browse files Browse the repository at this point in the history
Signed-off-by: Xuanwo <[email protected]>
  • Loading branch information
Xuanwo committed Sep 23, 2024
1 parent 5156461 commit c44a681
Show file tree
Hide file tree
Showing 4 changed files with 151 additions and 0 deletions.
3 changes: 3 additions & 0 deletions crates/reqsign/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
11 changes: 11 additions & 0 deletions crates/reqsign/src/context.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use crate::env::{Env, OsEnv};
use crate::{FileRead, HttpSend};
use anyhow::Result;
use bytes::Bytes;
Expand All @@ -8,17 +9,27 @@ use std::sync::Arc;
pub struct Context {
fs: Arc<dyn FileRead>,
http: Arc<dyn HttpSend>,
env: Arc<dyn Env>,
}

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<u8>`.
#[inline]
pub async fn file_read(&self, path: &str) -> Result<Vec<u8>> {
Expand Down
135 changes: 135 additions & 0 deletions crates/reqsign/src/env.rs
Original file line number Diff line number Diff line change
@@ -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<OsString>;

/// Return the path to the users home dir, returns `None` if any error occurs.
fn home_dir(&self) -> Option<PathBuf>;
}

/// 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<OsString> {
std::env::var_os(key)
}

#[cfg(any(unix, target_os = "redox"))]
fn home_dir(&self) -> Option<PathBuf> {
#[allow(deprecated)]
std::env::home_dir()
}

#[cfg(windows)]
fn home_dir(&self) -> Option<PathBuf> {
windows::home_dir_inner
}
}

/// Implements Env for the mock context.
#[cfg(test)]
#[derive(Debug, Clone)]
pub struct MockEnv {
pub home_dir: Option<PathBuf>,
pub envs: std::collections::HashMap<String, OsString>,
}

#[cfg(test)]
impl Env for MockEnv {
fn var_os(&self, key: &str) -> Option<OsString> {
self.envs.get(key).cloned()
}

fn home_dir(&self) -> Option<PathBuf> {
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<PathBuf> {
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<PathBuf> {
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<PathBuf> {
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));
}
}
}
2 changes: 2 additions & 0 deletions crates/reqsign/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

0 comments on commit c44a681

Please sign in to comment.