pitchfork_cli/proxy/
hosts.rs1use crate::settings::settings;
12use std::sync::OnceLock;
13
14const MARKER_START: &str = "# pitchfork-start";
16const MARKER_END: &str = "# pitchfork-end";
17static BLANK_LINES_RE: OnceLock<regex::Regex> = OnceLock::new();
18
19fn sync_hosts_file_with_slugs(bind_ip: &str, tld: &str, slug_names: &[String]) {
27 let entries: Vec<String> = slug_names
28 .iter()
29 .map(|slug| format!("{bind_ip} {slug}.{tld}"))
30 .collect();
31 write_hosts_block(&entries);
32}
33
34pub fn sync_hosts_from_settings() {
39 let slugs = crate::pitchfork_toml::PitchforkToml::read_global_slugs();
40 let slug_names: Vec<String> = slugs.keys().cloned().collect();
41 sync_hosts_from_settings_with_slugs(&slug_names);
42}
43
44pub fn sync_hosts_from_settings_with_slugs(slug_names: &[String]) {
51 let s = settings();
52 if s.proxy.enable && s.proxy.sync_hosts {
53 let lan_enabled = s.proxy.lan || !s.proxy.lan_ip.is_empty();
54 let tld = if lan_enabled { "local" } else { &s.proxy.tld };
55 let ip = if lan_enabled {
58 if !s.proxy.lan_ip.is_empty() {
60 s.proxy.lan_ip.clone()
61 } else {
62 s.proxy.host.clone()
67 }
68 } else {
69 s.proxy.host.clone()
70 };
71 sync_hosts_file_with_slugs(&ip, tld, slug_names);
72 }
73}
74
75pub fn clean_hosts_file() {
79 write_hosts_block(&[]);
80}
81
82fn write_hosts_block(entries: &[String]) {
84 let path = crate::env::PITCHFORK_HOSTS_FILE.clone();
85
86 let content = match std::fs::read_to_string(&path) {
87 Ok(c) => c,
88 Err(e) => {
89 if !entries.is_empty() {
90 log::warn!(
91 "Failed to read {} for hosts sync: {e}. \
92 Set proxy.sync_hosts = false to suppress this warning.",
93 path.display()
94 );
95 }
96 return;
97 }
98 };
99
100 let cleaned = remove_block(&content);
101
102 let new_content = if entries.is_empty() {
103 cleaned
104 } else {
105 let block = build_block(entries);
106 format!("{}\n{block}\n", cleaned.trim_end())
107 };
108
109 let parent = path.parent().unwrap_or(std::path::Path::new("/etc"));
111 let tmp_path = parent.join(format!(".pitchfork-hosts-tmp-{}", std::process::id()));
112
113 if let Err(e) = std::fs::write(&tmp_path, &new_content) {
114 log::warn!(
115 "Failed to write {} for hosts sync: {e}. \
116 Writing to /etc/hosts may require sudo. \
117 Set proxy.sync_hosts = false to suppress this warning.",
118 tmp_path.display()
119 );
120 let _ = std::fs::remove_file(&tmp_path);
121 return;
122 }
123
124 if let Err(e) = std::fs::rename(&tmp_path, &path) {
125 log::warn!(
126 "Failed to rename {} to {}: {e}. \
127 Writing to /etc/hosts may require sudo. \
128 Set proxy.sync_hosts = false to suppress this warning.",
129 tmp_path.display(),
130 path.display()
131 );
132 let _ = std::fs::remove_file(&tmp_path);
133 }
134}
135
136fn build_block(entries: &[String]) -> String {
138 if entries.is_empty() {
139 return String::new();
140 }
141 let lines = entries.join("\n");
142 format!("{MARKER_START}\n{lines}\n{MARKER_END}")
143}
144
145fn remove_block(content: &str) -> String {
148 let start_idx = match content.find(MARKER_START) {
149 Some(i) => i,
150 None => return content.to_string(),
151 };
152 let end_idx = match content[start_idx..].find(MARKER_END) {
153 Some(i) => start_idx + i + MARKER_END.len(),
154 None => return content.to_string(),
155 };
156 let before = &content[..start_idx];
157 let after = &content[end_idx..];
158 let result = format!("{before}{after}");
159 let re = BLANK_LINES_RE.get_or_init(|| regex::Regex::new(r"\n{3,}").unwrap());
161 re.replace_all(&result, "\n\n").trim_end().to_string() + "\n"
162}
163
164#[cfg(test)]
165mod tests {
166 use super::*;
167
168 #[test]
169 fn test_build_block() {
170 let entries = vec![
171 "127.0.0.1 myapp.localhost".to_string(),
172 "127.0.0.1 api.myapp.localhost".to_string(),
173 ];
174 let block = build_block(&entries);
175 assert!(block.starts_with("# pitchfork-start\n"));
176 assert!(block.ends_with("\n# pitchfork-end"));
177 assert!(block.contains("127.0.0.1 myapp.localhost"));
178 assert!(block.contains("127.0.0.1 api.myapp.localhost"));
179 }
180
181 #[test]
182 fn test_build_block_empty() {
183 assert!(build_block(&[]).is_empty());
184 }
185
186 #[test]
187 fn test_remove_block() {
188 let content =
189 "127.0.0.1 localhost\n# pitchfork-start\n127.0.0.1 myapp.localhost\n# pitchfork-end\n";
190 let cleaned = remove_block(content);
191 assert!(!cleaned.contains("pitchfork-start"));
192 assert!(!cleaned.contains("myapp.localhost"));
193 assert!(cleaned.contains("127.0.0.1 localhost"));
194 }
195
196 #[test]
197 fn test_remove_block_no_markers() {
198 let content = "127.0.0.1 localhost\n";
199 let cleaned = remove_block(content);
200 assert_eq!(cleaned, content);
201 }
202
203 #[test]
204 fn test_remove_block_normalizes_blank_lines() {
205 let content = "127.0.0.1 localhost\n\n\n# pitchfork-start\n127.0.0.1 myapp.localhost\n# pitchfork-end\n\n\n";
206 let cleaned = remove_block(content);
207 assert!(!cleaned.contains("\n\n\n"));
208 }
209
210 #[test]
211 fn test_remove_block_ignores_end_marker_before_start_marker() {
212 let content = "127.0.0.1 localhost\n# pitchfork-end\n# pitchfork-start\n127.0.0.1 myapp.localhost\n# pitchfork-end\n";
213 let cleaned = remove_block(content);
214 assert_eq!(cleaned, "127.0.0.1 localhost\n# pitchfork-end\n");
215 }
216}