refactor: optimize cache control headers file type detection (#175)

mfontanini · web-flow · commit b2cff1b1661b · 2023-02-22T07:46:12.000+01:00
diff --git a/src/control_headers.rs b/src/control_headers.rs
@@ -21,16 +21,12 @@ pub fn append_headers(uri: &str, resp: &mut Response<Body>) {
     // Default max-age value in seconds (one day)
     let mut max_age = MAX_AGE_ONE_DAY;
 
-    if CACHE_EXT_ONE_HOUR
-        .iter()
-        .any(|x| uri.ends_with(&[".", *x].concat()))
-    {
-        max_age = MAX_AGE_ONE_HOUR;
-    } else if CACHE_EXT_ONE_YEAR
-        .iter()
-        .any(|x| uri.ends_with(&[".", *x].concat()))
-    {
-        max_age = MAX_AGE_ONE_YEAR;
+    if let Some(extension) = uri_file_extension(uri) {
+        if CACHE_EXT_ONE_HOUR.binary_search(&extension).is_ok() {
+            max_age = MAX_AGE_ONE_HOUR;
+        } else if CACHE_EXT_ONE_YEAR.binary_search(&extension).is_ok() {
+            max_age = MAX_AGE_ONE_YEAR;
+        }
     }
 
     let cache_control = CacheControl::new()
@@ -44,13 +40,20 @@ fn duration_from_secs(secs: u64) -> std::time::Duration {
     std::time::Duration::from_secs(std::cmp::min(secs, u32::MAX as u64))
 }
 
+/// Gets the file extension for a URI.
+///
+/// This assumes the extension contains a single dot. e.g. for "/file.tar.gz" it returns "gz".
+fn uri_file_extension(uri: &str) -> Option<&str> {
+    uri.rsplit_once('.').map(|(_, rest)| rest)
+}
+
 #[cfg(test)]
 mod tests {
     use hyper::{Body, Response, StatusCode};
 
     use super::{
-        append_headers, CACHE_EXT_ONE_HOUR, CACHE_EXT_ONE_YEAR, MAX_AGE_ONE_DAY, MAX_AGE_ONE_HOUR,
-        MAX_AGE_ONE_YEAR,
+        append_headers, uri_file_extension, CACHE_EXT_ONE_HOUR, CACHE_EXT_ONE_YEAR,
+        MAX_AGE_ONE_DAY, MAX_AGE_ONE_HOUR, MAX_AGE_ONE_YEAR,
     };
 
     #[tokio::test]
@@ -101,4 +104,11 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn find_uri_extension() {
+        assert_eq!(uri_file_extension("/potato.zip"), Some("zip"));
+        assert_eq!(uri_file_extension("/potato."), Some(""));
+        assert_eq!(uri_file_extension("/"), None);
+    }
 }
