Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimize /etc/hosts writes #259

Merged
merged 6 commits into from
Jun 2, 2023
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions client/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -280,18 +280,20 @@ fn update_hosts_file(
hosts_path: PathBuf,
peers: &[Peer],
) -> Result<(), WrappedIoError> {
log::info!("updating {} with the latest peers.", "/etc/hosts".yellow());
mcginty marked this conversation as resolved.
Show resolved Hide resolved

let mut hosts_builder = HostsBuilder::new(format!("innernet {interface}"));
for peer in peers {
hosts_builder.add_hostname(
peer.contents.ip,
&format!("{}.{}.wg", peer.contents.name, interface),
);
}
if let Err(e) = hosts_builder.write_to(&hosts_path).with_path(hosts_path) {
log::warn!("failed to update hosts ({})", e);
}
match hosts_builder.write_to(&hosts_path).with_path(hosts_path) {
Ok(has_written) if has_written => {
evaporei marked this conversation as resolved.
Show resolved Hide resolved
log::info!("updated {} with the latest peers.", "/etc/hosts".yellow())
strohel marked this conversation as resolved.
Show resolved Hide resolved
},
Ok(_) => {},
Err(e) => log::warn!("failed to update hosts ({})", e),
};

Ok(())
}
Expand Down
66 changes: 43 additions & 23 deletions hostsfile/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use std::{
collections::HashMap,
collections::BTreeMap,
fmt,
fs::OpenOptions,
io::{self, BufRead, BufReader, ErrorKind, Write},
Expand Down Expand Up @@ -81,7 +81,7 @@ impl std::error::Error for Error {
/// ```
pub struct HostsBuilder {
tag: String,
hostname_map: HashMap<IpAddr, Vec<String>>,
hostname_map: BTreeMap<IpAddr, Vec<String>>,
}

impl HostsBuilder {
Expand All @@ -90,7 +90,7 @@ impl HostsBuilder {
pub fn new<S: Into<String>>(tag: S) -> Self {
Self {
tag: tag.into(),
hostname_map: HashMap::new(),
hostname_map: BTreeMap::new(),
}
}

Expand All @@ -116,7 +116,8 @@ impl HostsBuilder {

/// 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.
pub fn write(&self) -> io::Result<()> {
/// Returns true if the hosts file has changed.
pub fn write(&self) -> io::Result<bool> {
self.write_to(Self::default_path()?)
}

Expand Down Expand Up @@ -178,7 +179,9 @@ impl HostsBuilder {
///
/// 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).
pub fn write_to<P: AsRef<Path>>(&self, hosts_path: P) -> io::Result<()> {
///
/// 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.
Expand Down Expand Up @@ -206,9 +209,33 @@ impl HostsBuilder {
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);
}

let insert = match (begin, end) {
(Some(begin), Some(end)) => {
lines.drain(begin..end + 1);
let old_section: Vec<String> = lines.drain(begin..end + 1).collect();

let sections_are_eq = old_section == lines_to_insert;

if sections_are_eq {
evaporei marked this conversation as resolved.
Show resolved Hide resolved
return Ok(false);
}

begin
},
(None, None) => {
Expand All @@ -233,21 +260,12 @@ impl HostsBuilder {
for line in &lines[..insert] {
writeln!(s, "{line}")?;
}
if !self.hostname_map.is_empty() {
writeln!(s, "{begin_marker}")?;
for (ip, hostnames) in &self.hostname_map {
if cfg!(windows) {
// windows only allows one hostname per line
for hostname in hostnames {
writeln!(s, "{ip} {hostname}")?;
}
} else {
// assume the same format as Unix
writeln!(s, "{} {}", ip, hostnames.join(" "))?;
}
}
writeln!(s, "{end_marker}")?;

// Append hostnames_map section
for line in lines_to_insert {
writeln!(s, "{line}")?;
}

for line in &lines[insert..] {
writeln!(s, "{line}")?;
}
Expand All @@ -260,8 +278,9 @@ impl HostsBuilder {
_ => {
log::debug!("wrote hosts file with the write-and-swap strategy");
},
}
Ok(())
};

Ok(true)
}

fn write_and_swap(temp_path: &Path, hosts_path: &Path, contents: &[u8]) -> io::Result<()> {
Expand Down Expand Up @@ -314,7 +333,8 @@ mod tests {
temp_file.write_all(b"preexisting\ncontent").unwrap();
let mut builder = HostsBuilder::new("foo");
builder.add_hostname([1, 1, 1, 1].into(), "whatever");
builder.write_to(&temp_path).unwrap();
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}");
Expand Down