innernet/hostsfile/src/lib.rs

345 lines
11 KiB
Rust
Raw Permalink Normal View History

use std::{
collections::BTreeMap,
fmt,
fs::OpenOptions,
io::{self, BufRead, BufReader, ErrorKind, Write},
net::IpAddr,
path::{Path, PathBuf},
result,
time::{SystemTime, UNIX_EPOCH},
};
2021-03-29 17:22:14 +00:00
pub type Result<T> = result::Result<T, Box<dyn std::error::Error>>;
/// A custom error struct for this crate.
#[derive(Debug, Clone)]
pub struct Error(String);
impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
None
}
}
/// `HostsBuilder` manages a section of /etc/hosts file that contains a list of IP to hostname
/// mappings. A hosts file can have multiple sections that are distinguished by tag names.
///
/// # Examples
///
/// ```no_run
/// use hostsfile::{HostsBuilder, Result};
/// # fn main() -> Result<()> {
/// let mut hosts = HostsBuilder::new("dns");
/// hosts.add_hostname("8.8.8.8".parse().unwrap(), "google-dns1");
/// hosts.add_hostname("8.8.4.4".parse().unwrap(), "google-dns2");
/// hosts.write_to("/tmp/hosts")?;
/// # Ok(())
/// # }
/// ```
///
/// `/tmp/hosts` will have a section:
///
/// ```text
/// # DO NOT EDIT dns BEGIN
/// 8.8.8.8 google-dns1
/// 8.8.4.4 google-dns2
/// # DO NOT EDIT dns END
/// ```
///
/// Another run of `HostsBuilder` with the same tag name overrides the section.
///
/// ```no_run
/// use hostsfile::{HostsBuilder, Result};
/// # fn main() -> Result<()> {
/// let mut hosts = HostsBuilder::new("dns");
/// hosts.add_hostnames("1.1.1.1".parse().unwrap(), &["cloudflare-dns", "apnic-dns"]);
/// hosts.write_to("/tmp/hosts")?;
/// # Ok(())
/// # }
/// ```
///
/// `/tmp/hosts` will have a section:
///
/// ```text
/// # DO NOT EDIT dns BEGIN
/// 1.1.1.1 cloudflare-dns apnic-dns
/// # DO NOT EDIT dns END
/// ```
///
/// On Windows the host file format is slightly different in this case:
/// ```text
/// # DO NOT EDIT dns BEGIN
/// 1.1.1.1 cloudflare-dns
/// 1.1.1.1 apnic-dns
/// # DO NOT EDIT dns END
/// ```
2021-03-29 17:22:14 +00:00
pub struct HostsBuilder {
tag: String,
hostname_map: BTreeMap<IpAddr, Vec<String>>,
2021-03-29 17:22:14 +00:00
}
impl HostsBuilder {
/// Creates a new `HostsBuilder` with the given tag name. It corresponds to a section in the
/// hosts file containing a list of IP to hostname mappings.
pub fn new<S: Into<String>>(tag: S) -> Self {
Self {
tag: tag.into(),
hostname_map: BTreeMap::new(),
2021-03-29 17:22:14 +00:00
}
}
/// Adds a mapping of `ip` to `hostname`. If there hostnames associated with the IP already,
/// the hostname will be appended to the list.
pub fn add_hostname<S: ToString>(&mut self, ip: IpAddr, hostname: S) {
let hostnames_dest = self.hostname_map.entry(ip).or_default();
2021-03-29 17:22:14 +00:00
hostnames_dest.push(hostname.to_string());
}
/// Adds a mapping of `ip` to a list of `hostname`s. If there hostnames associated with the IP
/// already, the new hostnames will be appended to the list.
pub fn add_hostnames<I: IntoIterator<Item = impl ToString>>(
&mut self,
ip: IpAddr,
hostnames: I,
) {
let hostnames_dest = self.hostname_map.entry(ip).or_default();
2021-03-29 17:22:14 +00:00
for hostname in hostnames.into_iter() {
hostnames_dest.push(hostname.to_string());
}
}
/// Inserts a new section to the system's default hosts file. If there is a section with the
/// same tag name already, it will be replaced with the new list instead.
/// Returns true if the hosts file has changed.
pub fn write(&self) -> io::Result<bool> {
self.write_to(Self::default_path()?)
}
/// Returns the default hosts path based on the current OS.
pub fn default_path() -> io::Result<PathBuf> {
2021-03-29 17:22:14 +00:00
let hosts_file = if cfg!(unix) {
PathBuf::from("/etc/hosts")
} else if cfg!(windows) {
PathBuf::from(
// according to https://support.microsoft.com/en-us/topic/how-to-reset-the-hosts-file-back-to-the-default-c2a43f9d-e176-c6f3-e4ef-3500277a6dae
// the location depends on the environment variable %WinDir%.
format!(
"{}\\System32\\Drivers\\Etc\\hosts",
std::env::var("WinDir").map_err(|_| io::Error::new(
ErrorKind::Other,
"WinDir environment variable missing".to_owned()
))?
),
)
2021-03-29 17:22:14 +00:00
} else {
return Err(io::Error::new(
ErrorKind::Other,
"unsupported operating system.".to_owned(),
));
2021-03-29 17:22:14 +00:00
};
if !hosts_file.exists() {
return Err(ErrorKind::NotFound.into());
2021-03-29 17:22:14 +00:00
}
Ok(hosts_file)
2021-03-29 17:22:14 +00:00
}
pub fn get_temp_path(hosts_path: &Path) -> io::Result<PathBuf> {
let hosts_dir = hosts_path.parent().ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"hosts path missing a parent folder",
)
})?;
let start = SystemTime::now();
let since_the_epoch = start
.duration_since(UNIX_EPOCH)
.expect("Time went backwards");
let mut temp_filename = hosts_path
.file_name()
.ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "hosts path missing a filename")
})?
.to_os_string();
temp_filename.push(format!(".tmp{}", since_the_epoch.as_millis()));
Ok(hosts_dir.with_file_name(temp_filename))
}
2021-03-29 17:22:14 +00:00
/// Inserts a new section to the specified hosts file. If there is a section with the same tag
/// name already, it will be replaced with the new list instead.
///
/// `hosts_path` is the *full* path to write to, including the filename.
///
/// On Windows, the format of one hostname per line will be used, all other systems will use
/// the same format as Unix and Unix-like systems (i.e. allow multiple hostnames per line).
///
/// Returns true if the hosts file has changed.
pub fn write_to<P: AsRef<Path>>(&self, hosts_path: P) -> io::Result<bool> {
let hosts_path = hosts_path.as_ref();
if hosts_path.is_dir() {
// TODO(jake): use io::ErrorKind::IsADirectory when it's stable.
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"hosts path was a directory",
));
}
let temp_path = Self::get_temp_path(hosts_path)?;
2021-03-29 17:22:14 +00:00
let begin_marker = format!("# DO NOT EDIT {} BEGIN", &self.tag);
let end_marker = format!("# DO NOT EDIT {} END", &self.tag);
let hosts_file = OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(hosts_path)?;
let mut lines = BufReader::new(hosts_file)
.lines()
.map(|line| line.unwrap())
.collect::<Vec<_>>();
2021-03-29 17:22:14 +00:00
let begin = lines.iter().position(|line| line.trim() == begin_marker);
let end = lines.iter().position(|line| line.trim() == end_marker);
let mut lines_to_insert = vec![];
if !self.hostname_map.is_empty() {
lines_to_insert.push(begin_marker);
for (ip, hostnames) in &self.hostname_map {
if cfg!(windows) {
// windows only allows one hostname per line
for hostname in hostnames {
lines_to_insert.push(format!("{ip} {hostname}"));
}
} else {
// assume the same format as Unix
lines_to_insert.push(format!("{} {}", ip, hostnames.join(" ")));
}
}
lines_to_insert.push(end_marker);
}
2021-03-29 17:22:14 +00:00
let insert = match (begin, end) {
(Some(begin), Some(end)) => {
let old_section: Vec<String> = lines.drain(begin..end + 1).collect();
if old_section == lines_to_insert {
return Ok(false);
}
2021-03-29 17:22:14 +00:00
begin
},
(None, None) => {
// Insert a blank line before a new section.
if let Some(last_line) = lines.iter().last() {
2021-05-06 03:40:00 +00:00
if !last_line.is_empty() {
2021-03-29 17:22:14 +00:00
lines.push("".to_string());
}
}
lines.len()
},
_ => {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("start or end marker missing in {:?}", &hosts_path),
));
2021-03-29 17:22:14 +00:00
},
};
let mut s = vec![];
2021-03-29 17:22:14 +00:00
for line in &lines[..insert] {
writeln!(s, "{line}")?;
2021-03-29 17:22:14 +00:00
}
// Append hostnames_map section
for line in lines_to_insert {
writeln!(s, "{line}")?;
2021-03-29 17:22:14 +00:00
}
2021-03-29 17:22:14 +00:00
for line in &lines[insert..] {
writeln!(s, "{line}")?;
2021-03-29 17:22:14 +00:00
}
match Self::write_and_swap(&temp_path, hosts_path, &s) {
Err(_) => {
Self::write_clobber(hosts_path, &s)?;
log::debug!("wrote hosts file with the clobber fallback strategy");
},
_ => {
log::debug!("wrote hosts file with the write-and-swap strategy");
},
};
Ok(true)
}
fn write_and_swap(temp_path: &Path, hosts_path: &Path, contents: &[u8]) -> io::Result<()> {
// Copy the file we plan on modifying so its permissions and metadata are preserved.
2022-09-24 03:43:33 +00:00
std::fs::copy(hosts_path, temp_path)?;
Self::write_clobber(temp_path, contents)?;
std::fs::rename(temp_path, hosts_path)?;
Ok(())
}
fn write_clobber(hosts_path: &Path, contents: &[u8]) -> io::Result<()> {
OpenOptions::new()
.create(true)
.read(true)
.write(true)
.truncate(true)
.open(hosts_path)?
.write_all(contents)?;
2021-03-29 17:22:14 +00:00
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::Path;
#[test]
fn test_temp_path_good() {
let hosts_path = Path::new("/etc/hosts");
let temp_path = HostsBuilder::get_temp_path(hosts_path).unwrap();
println!("{temp_path:?}");
assert!(temp_path
.file_name()
.unwrap()
.to_str()
.unwrap()
.starts_with("hosts.tmp"));
}
#[test]
fn test_temp_path_invalid() {
let hosts_path = Path::new("/");
assert!(HostsBuilder::get_temp_path(hosts_path).is_err());
}
#[test]
fn test_write() {
let (mut temp_file, temp_path) = tempfile::NamedTempFile::new().unwrap().into_parts();
temp_file.write_all(b"preexisting\ncontent").unwrap();
let mut builder = HostsBuilder::new("foo");
builder.add_hostname([1, 1, 1, 1].into(), "whatever");
assert!(builder.write_to(&temp_path).unwrap());
assert!(!builder.write_to(&temp_path).unwrap());
let contents = std::fs::read_to_string(&temp_path).unwrap();
println!("contents: {contents}");
assert!(contents.starts_with("preexisting\ncontent"));
assert!(contents.contains("# DO NOT EDIT foo BEGIN"));
assert!(contents.contains("1.1.1.1 whatever"));
}
}