use crate::settings::settings;
use std::sync::OnceLock;
const MARKER_START: &str = "# pitchfork-start";
const MARKER_END: &str = "# pitchfork-end";
static BLANK_LINES_RE: OnceLock<regex::Regex> = OnceLock::new();
fn hosts_path() -> std::path::PathBuf {
if cfg!(windows) {
let system_root = std::env::var("SystemRoot").unwrap_or_else(|_| r"C:\Windows".to_string());
std::path::PathBuf::from(system_root)
.join("System32")
.join("drivers")
.join("etc")
.join("hosts")
} else {
std::path::PathBuf::from("/etc/hosts")
}
}
pub fn sync_hosts_file(bind_ip: &str, tld: &str) {
let slugs = crate::pitchfork_toml::PitchforkToml::read_global_slugs();
let mut entries: Vec<String> = Vec::new();
for slug in slugs.keys() {
entries.push(format!("{bind_ip} {slug}.{tld}"));
}
write_hosts_block(&entries);
}
pub fn sync_hosts_from_settings() {
let s = settings();
if s.proxy.enable && s.proxy.sync_hosts {
sync_hosts_file(&s.proxy.host, &s.proxy.tld);
}
}
pub fn clean_hosts_file() {
write_hosts_block(&[]);
}
fn write_hosts_block(entries: &[String]) {
let path = hosts_path();
let content = match std::fs::read_to_string(&path) {
Ok(c) => c,
Err(e) => {
if !entries.is_empty() {
log::warn!(
"Failed to read {} for hosts sync: {e}. \
Set proxy.sync_hosts = false to suppress this warning.",
path.display()
);
}
return;
}
};
let cleaned = remove_block(&content);
let new_content = if entries.is_empty() {
cleaned
} else {
let block = build_block(entries);
format!("{}\n{block}\n", cleaned.trim_end())
};
let parent = path.parent().unwrap_or(std::path::Path::new("/etc"));
let tmp_path = parent.join(format!(".pitchfork-hosts-tmp-{}", std::process::id()));
if let Err(e) = std::fs::write(&tmp_path, &new_content) {
log::warn!(
"Failed to write {} for hosts sync: {e}. \
Writing to /etc/hosts may require sudo. \
Set proxy.sync_hosts = false to suppress this warning.",
tmp_path.display()
);
let _ = std::fs::remove_file(&tmp_path);
return;
}
if let Err(e) = std::fs::rename(&tmp_path, &path) {
log::warn!(
"Failed to rename {} to {}: {e}. \
Writing to /etc/hosts may require sudo. \
Set proxy.sync_hosts = false to suppress this warning.",
tmp_path.display(),
path.display()
);
let _ = std::fs::remove_file(&tmp_path);
}
}
fn build_block(entries: &[String]) -> String {
if entries.is_empty() {
return String::new();
}
let lines = entries.join("\n");
format!("{MARKER_START}\n{lines}\n{MARKER_END}")
}
fn remove_block(content: &str) -> String {
let start_idx = match content.find(MARKER_START) {
Some(i) => i,
None => return content.to_string(),
};
let end_idx = match content[start_idx..].find(MARKER_END) {
Some(i) => start_idx + i + MARKER_END.len(),
None => return content.to_string(),
};
let before = &content[..start_idx];
let after = &content[end_idx..];
let result = format!("{before}{after}");
let re = BLANK_LINES_RE.get_or_init(|| regex::Regex::new(r"\n{3,}").unwrap());
re.replace_all(&result, "\n\n").trim_end().to_string() + "\n"
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_build_block() {
let entries = vec![
"127.0.0.1 myapp.localhost".to_string(),
"127.0.0.1 api.myapp.localhost".to_string(),
];
let block = build_block(&entries);
assert!(block.starts_with("# pitchfork-start\n"));
assert!(block.ends_with("\n# pitchfork-end"));
assert!(block.contains("127.0.0.1 myapp.localhost"));
assert!(block.contains("127.0.0.1 api.myapp.localhost"));
}
#[test]
fn test_build_block_empty() {
assert!(build_block(&[]).is_empty());
}
#[test]
fn test_remove_block() {
let content =
"127.0.0.1 localhost\n# pitchfork-start\n127.0.0.1 myapp.localhost\n# pitchfork-end\n";
let cleaned = remove_block(content);
assert!(!cleaned.contains("pitchfork-start"));
assert!(!cleaned.contains("myapp.localhost"));
assert!(cleaned.contains("127.0.0.1 localhost"));
}
#[test]
fn test_remove_block_no_markers() {
let content = "127.0.0.1 localhost\n";
let cleaned = remove_block(content);
assert_eq!(cleaned, content);
}
#[test]
fn test_remove_block_normalizes_blank_lines() {
let content = "127.0.0.1 localhost\n\n\n# pitchfork-start\n127.0.0.1 myapp.localhost\n# pitchfork-end\n\n\n";
let cleaned = remove_block(content);
assert!(!cleaned.contains("\n\n\n"));
}
#[test]
fn test_remove_block_ignores_end_marker_before_start_marker() {
let content = "127.0.0.1 localhost\n# pitchfork-end\n# pitchfork-start\n127.0.0.1 myapp.localhost\n# pitchfork-end\n";
let cleaned = remove_block(content);
assert_eq!(cleaned, "127.0.0.1 localhost\n# pitchfork-end\n");
}
}