Add quirks support for youtube-nocookie.com and youtube embed URLs (#…

…1563)
lycheeverse · Nov 8, 2024 · f4a35ff · f4a35ff
1 parent a28c92b
commit f4a35ff
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 8 deletions.
diff --git a/lychee-bin/tests/cli.rs b/lychee-bin/tests/cli.rs
@@ -892,8 +892,9 @@ mod cli {
 
         // Ensure clean state
         if cache_file.exists() {
+            println!("Removing cache file before test: {:?}", cache_file);
             fs::remove_file(&cache_file)?;
-            tokio::time::sleep(Duration::from_millis(100)).await;
+            tokio::time::sleep(Duration::from_millis(500)).await;
         }
 
         // Setup mock servers
@@ -925,7 +926,7 @@ mod cli {
         let _output = cmd.output()?;
 
         // Wait for cache file to be written
-        for _ in 0..10 {
+        for _ in 0..20 {
             if cache_file.exists() {
                 break;
             }
@@ -934,6 +935,8 @@ mod cli {
 
         // Check cache contents
         let data = fs::read_to_string(&cache_file)?;
+        println!("Cache file contents: {}", data);
+
         assert!(
             data.contains(&format!("{}/,200", mock_server_ok.uri())),
             "Missing OK entry in cache"
@@ -955,7 +958,13 @@ mod cli {
             )));
 
         // Clean up
-        fs::remove_file(&cache_file)?;
+        fs::remove_file(&cache_file).map_err(|e| {
+            anyhow::anyhow!(
+                "Failed to remove cache file: {:?}, error: {}",
+                cache_file,
+                e
+            )
+        })?;
 
         Ok(())
     }

diff --git a/lychee-lib/src/quirks/mod.rs b/lychee-lib/src/quirks/mod.rs
@@ -13,7 +13,7 @@ use std::collections::HashMap;
 static CRATES_PATTERN: Lazy<Regex> =
     Lazy::new(|| Regex::new(r"^(https?://)?(www\.)?crates.io").unwrap());
 static YOUTUBE_PATTERN: Lazy<Regex> =
-    Lazy::new(|| Regex::new(r"^(https?://)?(www\.)?(youtube\.com)").unwrap());
+    Lazy::new(|| Regex::new(r"^(https?://)?(www\.)?youtube(-nocookie)?\.com").unwrap());
 static YOUTUBE_SHORT_PATTERN: Lazy<Regex> =
     Lazy::new(|| Regex::new(r"^(https?://)?(www\.)?(youtu\.?be)").unwrap());
 
@@ -48,13 +48,21 @@ impl Default for Quirks {
             Quirk {
                 pattern: &YOUTUBE_PATTERN,
                 rewrite: |mut request| {
-                    if request.url().path() != "/watch" {
-                        return request;
-                    }
-                    if let Some(id) = query(&request).get("v") {
+                    // Extract video id if it's a video page
+                    let video_id = match request.url().path() {
+                        "/watch" => query(&request).get("v").map(ToOwned::to_owned),
+                        path if path.starts_with("/embed/") => {
+                            path.strip_prefix("/embed/").map(ToOwned::to_owned)
+                        }
+                        _ => return request,
+                    };
+
+                    // Only rewrite to thumbnail if we got a video id
+                    if let Some(id) = video_id {
                         *request.url_mut() =
                             Url::parse(&format!("https://img.youtube.com/vi/{id}/0.jpg")).unwrap();
                     }
+
                     request
                 },
             },
@@ -146,6 +154,19 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_youtube_video_nocookie_request() {
+        let url = Url::parse("https://www.youtube-nocookie.com/embed/BIguvia6AvM").unwrap();
+        let request = Request::new(Method::GET, url);
+        let modified = Quirks::default().apply(request);
+        let expected_url = Url::parse("https://img.youtube.com/vi/BIguvia6AvM/0.jpg").unwrap();
+
+        assert_eq!(
+            MockRequest(modified),
+            MockRequest::new(Method::GET, expected_url)
+        );
+    }
+
     #[test]
     fn test_youtube_video_shortlink_request() {
         let url = Url::parse("https://youtu.be/Rvu7N4wyFpk?t=42").unwrap();