From 3623cc77b4ea1225915ce4d778ec53fe515cf895 Mon Sep 17 00:00:00 2001
From: Emi Simpson <emi@alchemi.dev>
Date: Tue, 8 Feb 2022 16:41:12 -0500
Subject: [PATCH] Handle non-unicode domains before the server starts

(also annotate all of the panics)
---
 Cargo.lock    |  1 +
 Cargo.toml    |  1 +
 src/main.rs   | 12 ++++++++++++
 src/server.rs | 13 +++++++------
 src/util.rs   | 22 ++++++++++++++++++++++
 5 files changed, 43 insertions(+), 6 deletions(-)
 create mode 100644 src/util.rs

diff --git a/Cargo.lock b/Cargo.lock
index d7267f7..50d35fb 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -54,6 +54,7 @@ name = "faery-ring"
 version = "0.1.0"
 dependencies = [
  "argh",
+ "ascii",
  "tiny_http",
 ]
 
diff --git a/Cargo.toml b/Cargo.toml
index 5a3f038..77e2a12 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,6 +16,7 @@ maintenance = { status = "passively-maintained" }
 [dependencies]
 tiny_http = "0.10.0"
 argh = "0.1.7"
+ascii = "1.0.0"
 
 [profile.release]
 lto = "fat"
diff --git a/src/main.rs b/src/main.rs
index b6cb273..b0dae90 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,5 @@
 pub mod server;
+pub mod util;
 
 use std::io::Read;
 use std::process::exit;
@@ -66,6 +67,17 @@ fn main() {
 	let domains: Vec<_> = read_into.split('\n')
 		.filter(|s| !s.is_empty())
 		.collect();
+	let domains = match util::to_ascii_list(&domains) {
+		Ok(d) => d,
+		Err((line, e)) => {
+			let line = domains[line]; // unwrap safe guaranteed by to_ascii_list()
+			let character = line.chars().nth(e.valid_up_to()).unwrap(); // unwrap safe
+			eprintln!(
+				"The domains file can only contain ASCII characters, but it looks like the domain {line} has a non-ASCII character ({character})"
+			);
+			exit(1005);
+		}
+	};
 
 	//
 	// Step 2: Run the server
diff --git a/src/server.rs b/src/server.rs
index 565e692..f255068 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -2,6 +2,7 @@ use std::borrow::Cow;
 use std::io::Cursor;
 use std::process::exit;
 
+use ascii::AsciiStr;
 use tiny_http::Header;
 use tiny_http::Request;
 use tiny_http::Response;
@@ -12,7 +13,7 @@ use tiny_http::StatusCode;
 ///
 /// Runs the server in a single thread, bound to 0.0.0.0 on the provided port.  Requests are
 /// handled first-come, first-serve
-pub fn go(domains: &[&str], port: u16) {
+pub fn go(domains: &[&AsciiStr], port: u16) {
 	let server = match Server::http(("0.0.0.0", port)) {
 		Ok(s) => s,
 		Err(e) => {
@@ -62,14 +63,14 @@ impl<'a> Route<'a> {
 	}
 
 	/// Generate a Response for this route
-	pub fn render<'b>(&self, domains: &[&'b str]) -> Response<Cursor<Cow<'b, [u8]>>> {
+	pub fn render<'b>(&self, domains: &[&'b AsciiStr]) -> Response<Cursor<Cow<'b, [u8]>>> {
 		match self {
 			Self::PrevOrNext(this_domain, offset) => {
 
 				let destination = domains.iter()
-					.position(|d| d.starts_with(this_domain))
+					.position(|d| d.as_str().starts_with(this_domain))
 					.map(|i| (i as isize + offset).rem_euclid(domains.len() as isize))
-					.map(|i| domains[i as usize]);
+					.map(|i| domains[i as usize]); // safe to unwrap:  we just did a rem_euclid
 
 				if let Some(url) = destination {
 					let destination = format!(
@@ -82,11 +83,11 @@ impl<'a> Route<'a> {
 							Header::from_bytes(
 								b"Location".to_owned(),
 								destination.as_bytes(),
-							).unwrap(),
+							).unwrap(), // safe to unwrap:  both of these are guaranteed ASCII
 							Header::from_bytes(
 								b"Cache-Control".to_owned(),
 								b"max-age=86400,stale-while-revalidate=172800,stale-if-error=3155760000".to_owned(),
-							).unwrap(),
+							).unwrap(),// safe to unwrap:  both of these are guaranteed ASCII
 						],
 						Cursor::new(Cow::Borrowed(b"")),
 						Some(0),
diff --git a/src/util.rs b/src/util.rs
new file mode 100644
index 0000000..0a67661
--- /dev/null
+++ b/src/util.rs
@@ -0,0 +1,22 @@
+//! Several small utility methods
+
+use ascii::AsAsciiStrError;
+use ascii::AsciiStr;
+
+/// Convert a series of [`&str`][]s into [`AsciiStr`][]s
+///
+/// Usually used for the domain list.  The result is either a [`Vec`][] containing all
+/// mapped strings, or else a touple containing the index of the string with the first
+/// error and the error that occurred.
+pub fn to_ascii_list<'a>(strs: &'a [&'a str]) -> Result<Vec<&'a AsciiStr>, (usize, AsAsciiStrError)> {
+	strs.iter()
+		.map(AsciiStr::from_ascii)
+		.enumerate()
+		.try_fold(Vec::with_capacity(strs.len()), |mut v, (i, s)| {
+			match s {
+				Ok(s) => v.push(s),
+				Err(e) => return Err((i, e)),
+			};
+			Ok(v)
+		})
+}