use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
#[derive(Clone)]
pub struct SpanFields {
name: &'static str,
fields: Arc<Mutex<HashMap<String, String>>>,
}
impl SpanFields {
pub fn capturing(name: &'static str) -> Self {
Self {
name,
fields: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn get(&self, field: &str) -> Option<String> {
self.fields.lock().unwrap().get(field).cloned()
}
}
impl tracing::field::Visit for SpanFields {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
self.fields
.lock()
.unwrap()
.insert(field.name().to_string(), format!("{value:?}"));
}
}
impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for SpanFields {
fn on_new_span(
&self,
attrs: &tracing::span::Attributes<'_>,
_id: &tracing::Id,
_ctx: tracing_subscriber::layer::Context<'_, S>,
) {
if attrs.metadata().name() == self.name {
attrs.record(&mut self.clone());
}
}
fn on_record(
&self,
_id: &tracing::Id,
values: &tracing::span::Record<'_>,
_ctx: tracing_subscriber::layer::Context<'_, S>,
) {
values.record(&mut self.clone());
}
}
pub async fn capture_request_span<F, T>(body: F) -> SpanFields
where
F: std::future::Future<Output = T>,
{
use tracing_subscriber::layer::SubscriberExt;
let captured = SpanFields::capturing("request");
let subscriber = tracing_subscriber::registry().with(captured.clone());
let _guard = tracing::subscriber::set_default(subscriber);
body.await;
captured
}
pub struct TempDir(PathBuf);
impl TempDir {
pub fn new(label: &str) -> Self {
let path =
std::env::temp_dir().join(format!("acme-proxy-{label}-{}", uuid::Uuid::now_v7()));
std::fs::create_dir_all(&path).expect("temp directory must be creatable");
Self(path)
}
pub fn path(&self) -> &Path {
&self.0
}
pub fn join(&self, name: &str) -> PathBuf {
self.0.join(name)
}
pub fn write(&self, name: &str, contents: &str) -> PathBuf {
let path = self.join(name);
std::fs::write(&path, contents).expect("temp file must be writable");
path
}
}
impl AsRef<Path> for TempDir {
fn as_ref(&self) -> &Path {
&self.0
}
}
impl Drop for TempDir {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.0);
}
}
#[cfg(unix)]
pub fn write_script(dir: &TempDir, name: &str, body: &str) -> PathBuf {
use std::os::unix::fs::PermissionsExt;
let path = dir.write(name, body);
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o755))
.expect("script must be made executable");
path
}
pub struct EnvGuard {
keys: Vec<String>,
_lock: std::sync::MutexGuard<'static, ()>,
}
impl EnvGuard {
pub fn new(vars: &[(&str, &str)]) -> Self {
let _lock = crate::config::ENV_LOCK
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut keys = vec!["ACME_PROXY_CONFIG".to_string()];
unsafe {
std::env::set_var("ACME_PROXY_CONFIG", "/nonexistent/acme-proxy-config");
for (key, value) in vars {
std::env::set_var(key, value);
keys.push((*key).to_string());
}
}
Self { keys, _lock }
}
}
impl Drop for EnvGuard {
fn drop(&mut self) {
unsafe {
for key in &self.keys {
std::env::remove_var(key);
}
}
}
}