implemented cache

.gitignore

@@ -1,2 +1,4 @@
 /target
 /Cargo.lock
+
+rusty-tube.json

Cargo.toml

@@ -26,3 +26,4 @@ chrono = {version = "0.4.19", features = ["serde"]}
 env_logger = "0.9.0"
 test-log = "0.2.11"
 rstest = "0.15.0"
+temp_testdir = "0.2.3"

src/cache.rs

@@ -1,17 +1,25 @@
-use std::{future::Future, sync::Arc};
+use std::{
+    fs::File,
+    future::Future,
+    io::BufReader,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 
 use anyhow::Result;
 use chrono::{DateTime, Duration, Utc};
+use log::{error, info};
 use serde::{Deserialize, Serialize};
 use tokio::sync::Mutex;
 
-#[derive(Default, Debug)]
+#[derive(Default, Debug, Clone)]
 pub struct Cache {
+    file: Option<PathBuf>,
     data: Arc<Mutex<CacheData>>,
 }
 
 #[derive(Default, Debug, Clone, Serialize, Deserialize)]
-pub struct CacheData {
+struct CacheData {
     desktop_client: Option<CacheEntry<DesktopClientData>>,
     deobf: Option<CacheEntry<DeobfData>>,
 }
@@ -38,10 +46,10 @@ pub struct DesktopClientData {
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct DeobfData {
-    js_url: String,
+    pub js_url: String,
-    sig_fn: String,
+    pub sig_fn: String,
-    nsig_fn: String,
+    pub nsig_fn: String,
-    sts: String,
+    pub sts: String,
 }
 
 impl Cache {
@@ -52,10 +60,11 @@ impl Cache {
         let mut cache = self.data.lock().await;
 
         if cache.desktop_client.is_none()
-            || cache.desktop_client.as_ref().unwrap().last_update < Utc::now() - Duration::days(1)
+            || cache.desktop_client.as_ref().unwrap().last_update < Utc::now() - Duration::hours(24)
         {
             let cdata = updater.await?;
             cache.desktop_client = Some(CacheEntry::from(cdata.clone()));
+            self.save(&cache);
             Ok(cdata)
         } else {
             Ok(cache.desktop_client.as_ref().unwrap().data.clone())
@@ -68,33 +77,102 @@ impl Cache {
     {
         let mut cache = self.data.lock().await;
         if cache.deobf.is_none()
-            || cache.deobf.as_ref().unwrap().last_update < Utc::now() - Duration::days(1)
+            || cache.deobf.as_ref().unwrap().last_update < Utc::now() - Duration::hours(1)
         {
             let deobf_data = updater.await?;
             cache.deobf = Some(CacheEntry::from(deobf_data.clone()));
+            self.save(&cache);
             Ok(deobf_data)
         } else {
             Ok(cache.deobf.as_ref().unwrap().data.clone())
         }
     }
 
-    pub async fn to_json(&self) -> Result<String, serde_json::Error> {
+    pub async fn to_json(&self) -> Result<String> {
         let cache = self.data.lock().await;
-        serde_json::to_string(&cache.clone())
+        Ok(serde_json::to_string(&cache.clone())?)
     }
 
-    pub async fn from_json(&self, json: &str) -> Result<(), serde_json::Error> {
+    pub async fn to_json_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
-        let cd = serde_json::from_str::<CacheData>(json)?;
+        let cache = self.data.lock().await;
+        Ok(serde_json::to_writer(&File::create(path)?, &cache.clone())?)
+    }
 
-        let mut cache = self.data.lock().await;
+    pub fn from_json(json: &str) -> Self {
-        *cache = cd;
+        let data: CacheData = match serde_json::from_str(json) {
+            Ok(cd) => cd,
+            Err(e) => {
+                error!(
+                    "Could not load cache from json, falling back to default. Error: {}",
+                    e
+                );
+                CacheData::default()
+            }
+        };
+        Cache {
+            data: Arc::new(Mutex::new(data)),
+            file: None,
+        }
+    }
 
-        Ok(())
+    pub fn from_json_file<P: AsRef<Path>>(path: P) -> Self {
+        let file = match File::open(path.as_ref()) {
+            Ok(file) => file,
+            Err(e) => {
+                if e.kind() == std::io::ErrorKind::NotFound {
+                    info!(
+                        "Cache json file at {} not found, will be created",
+                        path.as_ref().to_string_lossy()
+                    )
+                } else {
+                    error!(
+                        "Could not open cache json file, falling back to default. Error: {}",
+                        e
+                    );
+                }
+                return Cache {
+                    file: Some(path.as_ref().to_path_buf()),
+                    ..Default::default()
+                };
+            }
+        };
+        let data: CacheData = match serde_json::from_reader(BufReader::new(file)) {
+            Ok(data) => data,
+            Err(e) => {
+                error!(
+                    "Could not load cache from json, falling back to default. Error: {}",
+                    e
+                );
+                return Cache {
+                    file: Some(path.as_ref().to_path_buf()),
+                    ..Default::default()
+                };
+            }
+        };
+        Cache {
+            data: Arc::new(Mutex::new(data)),
+            file: Some(path.as_ref().to_path_buf()),
+        }
+    }
+
+    fn save(&self, cache: &CacheData) {
+        match self.file.as_ref() {
+            Some(file) => match File::create(file) {
+                Ok(file) => match serde_json::to_writer(file, cache) {
+                    Ok(_) => {}
+                    Err(e) => error!("Could not write cache to json. Error: {}", e),
+                },
+                Err(e) => error!("Could not open cache json file. Error: {}", e),
+            },
+            None => {}
+        }
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use temp_testdir::TempDir;
+
     use super::*;
 
     #[tokio::test]
@@ -143,9 +221,69 @@ mod tests {
         );
 
         let json = cache.to_json().await.unwrap();
+        let new_cache = Cache::from_json(&json);
 
-        let new_cache = Cache::default();
+        assert_eq!(
-        new_cache.from_json(&json).await.unwrap();
+            new_cache
+                .get_desktop_client_data(async {
+                    Ok(DesktopClientData {
+                        client_version: "".to_owned(),
+                    })
+                })
+                .await
+                .unwrap(),
+            cdata
+        );
+
+        assert_eq!(
+            new_cache
+                .get_deobf_data(async {
+                    Ok(DeobfData {
+                        js_url: "".to_owned(),
+                        nsig_fn: "".to_owned(),
+                        sig_fn: "".to_owned(),
+                        sts: "".to_owned(),
+                    })
+                })
+                .await
+                .unwrap(),
+            deobf_data
+        );
+    }
+
+    #[tokio::test]
+    async fn test_file() {
+        let temp = TempDir::default();
+        let mut file_path = PathBuf::from(temp.as_ref());
+        file_path.push("cache.json");
+
+        let cache = Cache::from_json_file(file_path.clone());
+
+        let cdata = cache
+            .get_desktop_client_data(async {
+                Ok(DesktopClientData {
+                    client_version: "1.2.3".to_owned(),
+                })
+            })
+            .await
+            .unwrap();
+
+        let deobf_data = cache
+            .get_deobf_data(async {
+                Ok(DeobfData {
+                    js_url:
+                        "https://www.youtube.com/s/player/011af516/player_ias.vflset/en_US/base.js"
+                            .to_owned(),
+                    sig_fn: "t_sig_fn".to_owned(),
+                    nsig_fn: "t_nsig_fn".to_owned(),
+                    sts: "t_sts".to_owned(),
+                })
+            })
+            .await
+            .unwrap();
+
+        assert!(file_path.exists());
+        let new_cache = Cache::from_json_file(file_path.clone());
 
         assert_eq!(
             new_cache

src/client/mod.rs

@@ -3,7 +3,7 @@ mod response;
 
 use std::{sync::Arc, time::Instant};
 
-use anyhow::{anyhow, bail, Context, Result};
+use anyhow::{anyhow, Context, Result};
 use async_trait::async_trait;
 use fancy_regex::Regex;
 use log::{debug, warn};
@@ -13,7 +13,11 @@ use reqwest::{header, Client, ClientBuilder, Method, Request, RequestBuilder, Re
 use serde::Serialize;
 use tokio::sync::Mutex;
 
-use crate::{deobfuscate::Deobfuscator, util};
+use crate::{
+    cache::{Cache, DesktopClientData},
+    deobfuscate::Deobfuscator,
+    util,
+};
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub enum ClientType {
@@ -114,6 +118,7 @@ const IOS_DEVICE_MODEL: &str = "iPhone14,5";
 
 pub struct RustyTube {
     pub locale: Arc<Locale>,
+    cache: Cache,
     desktop_client: Arc<DesktopClient>,
     android_client: Arc<AndroidClient>,
     ios_client: Arc<IosClient>,
@@ -128,19 +133,25 @@ pub struct Locale {
 impl RustyTube {
     #[must_use]
     pub fn new() -> Self {
-        Self::new_with_ua("en", "US")
+        Self::new_with_ua("en", "US", Some("rusty-tube.json".to_owned()))
     }
 
     #[must_use]
-    pub fn new_with_ua(lang: &str, country: &str) -> Self {
+    pub fn new_with_ua(lang: &str, country: &str, cache_file: Option<String>) -> Self {
         let locale = Arc::new(Locale {
             lang: lang.to_owned(),
             country: country.to_owned(),
         });
 
+        let cache = match cache_file.as_ref() {
+            Some(cache_file) => Cache::from_json_file(cache_file),
+            None => Cache::default(),
+        };
+
         Self {
             locale: locale.clone(),
-            desktop_client: Arc::new(DesktopClient::new(locale.clone())),
+            cache: cache.clone(),
+            desktop_client: Arc::new(DesktopClient::new(locale.clone(), cache)),
             android_client: Arc::new(AndroidClient::new(locale.clone())),
             ios_client: Arc::new(IosClient::new(locale)),
         }
@@ -159,48 +170,29 @@ impl RustyTube {
 
 #[async_trait]
 pub trait YTClient {
-    // fn new(locale: Arc<Locale>) -> Self;
-
     async fn get_context(&self, localized: bool) -> ContextYT;
     async fn request_builder(&self, method: Method, url: &str) -> RequestBuilder;
     async fn exec_request(&self, request: Request) -> Result<Response>;
     async fn exec_request_text(&self, request: Request) -> Result<String>;
 }
 
+async fn exec_request(http: Client, request: Request) -> Result<Response> {
+    Ok(http.execute(request).await?.error_for_status()?)
+}
+
+async fn exec_request_text(http: Client, request: Request) -> Result<String> {
+    Ok(exec_request(http, request).await?.text().await?)
+}
+
 pub struct DesktopClient {
     locale: Arc<Locale>,
     http: Client,
-    data: Mutex<DesktopClientData>,
+    cache: Cache,
     consent_cookie_yes: String,
     consent_cookie_no: String,
     deobf: Deobfuscator,
 }
 
-#[derive(Debug)]
-struct DesktopClientData {
-    last_update: Option<Instant>,
-    client_version: String,
-}
-
-impl Default for DesktopClientData {
-    fn default() -> Self {
-        Self {
-            last_update: None,
-            client_version: DESKTOP_CLIENT_VERSION.to_owned(),
-        }
-    }
-}
-
-impl DesktopClientData {
-    fn is_old(&self) -> bool {
-        self.last_update.is_none()
-            || Instant::now()
-                .duration_since(self.last_update.unwrap())
-                .as_secs()
-                > 86400
-    }
-}
-
 #[async_trait]
 impl YTClient for DesktopClient {
     async fn get_context(&self, localized: bool) -> ContextYT {
@@ -253,7 +245,7 @@ impl YTClient for DesktopClient {
 }
 
 impl DesktopClient {
-    fn new(locale: Arc<Locale>) -> Self {
+    fn new(locale: Arc<Locale>, cache: Cache) -> Self {
         let mut rng = rand::thread_rng();
 
         let http = ClientBuilder::new()
@@ -263,12 +255,12 @@ impl DesktopClient {
             .build()
             .expect("unable to build the HTTP client");
 
-        let deobf = Deobfuscator::new(http.clone());
+        let deobf = Deobfuscator::new(http.clone(), cache.clone());
 
         Self {
             locale,
             http,
-            data: Mutex::new(DesktopClientData::default()),
+            cache,
             consent_cookie_yes: format!(
                 "{}={}{}",
                 CONSENT_COOKIE,
@@ -285,19 +277,21 @@ impl DesktopClient {
         }
     }
 
-    async fn extract_client_version_from_swjs(&self) -> Result<Option<String>> {
+    async fn extract_client_version_from_swjs(
-        let swjs = self
+        http: Client,
-            .exec_request_text(
+        consent_cookie: &str,
-                self.http
+    ) -> Result<String> {
-                    .get("https://www.youtube.com/sw.js")
+        let swjs = exec_request_text(
-                    .header(header::ORIGIN, "https://www.youtube.com")
+            http.clone(),
-                    .header(header::REFERER, "https://www.youtube.com")
+            http.get("https://www.youtube.com/sw.js")
-                    .header(header::COOKIE, self.consent_cookie_yes.to_owned())
+                .header(header::ORIGIN, "https://www.youtube.com")
-                    .build()
+                .header(header::REFERER, "https://www.youtube.com")
-                    .unwrap(),
+                .header(header::COOKIE, consent_cookie)
-            )
+                .build()
-            .await
+                .unwrap(),
-            .context("Failed to download sw.js")?;
+        )
+        .await
+        .context("Failed to download sw.js")?;
 
         static CLIENT_VERSION_PATTERNS: Lazy<[Regex; 3]> = Lazy::new(|| {
             [
@@ -307,41 +301,30 @@ impl DesktopClient {
             ]
         });
 
-        Ok(util::get_cg_from_regexes(
+        util::get_cg_from_regexes(CLIENT_VERSION_PATTERNS.iter(), &swjs, 1)
-            CLIENT_VERSION_PATTERNS.iter(),
+            .ok_or(anyhow!("Could not find desktop client version in sw.js"))
-            &swjs,
-            1,
-        ))
     }
 
     async fn get_client_version(&self) -> String {
-        let mut client_data = self.data.lock().await;
+        let http = self.http.clone();
+        let consent_cookie = self.consent_cookie_yes.clone();
 
-        if client_data.is_old() {
+        let client_data = self
-            let client_version = self.extract_client_version_from_swjs().await;
+            .cache
-            let new_version = match client_version {
+            .get_desktop_client_data(async move {
-                Ok(client_version) => match client_version {
+                let client_version =
-                    Some(client_version) => {
+                    Self::extract_client_version_from_swjs(http, &consent_cookie).await?;
-                        debug!("Updated desktop client version to {}", client_version);
+                Ok(DesktopClientData { client_version })
-                        client_version
+            })
-                    }
+            .await;
-                    None => {
-                        warn!("Could not find desktop client version in sw.js");
-                        DESKTOP_CLIENT_VERSION.to_owned()
-                    }
-                },
-                Err(e) => {
-                    warn!("Could not extract desktop client version, Error: {}", e);
-                    DESKTOP_CLIENT_VERSION.to_owned()
-                }
-            };
 
-            *client_data = DesktopClientData {
+        match client_data {
-                client_version: new_version,
+            Ok(client_data) => client_data.client_version,
-                last_update: Some(Instant::now()),
+            Err(e) => {
+                warn!("{}", e);
+                DESKTOP_CLIENT_VERSION.to_owned()
             }
         }
-        client_data.client_version.to_owned()
     }
 }
 

src/client/player.rs

@@ -2,7 +2,7 @@ use anyhow::{anyhow, bail, Context, Result};
 use reqwest::Method;
 use serde::Serialize;
 
-use super::{response, ContextYT, ClientType, RustyTube, YTClient};
+use super::{response, ClientType, ContextYT, RustyTube, YTClient};
 use crate::util;
 
 // REQUEST
@@ -74,8 +74,7 @@ impl RustyTube {
             }
         };
 
-        let resp = self
+        let resp = client
-            .desktop_client
             .request_builder(Method::POST, "player")
             .await
             .json(&request_body)
@@ -97,7 +96,10 @@ mod tests {
     #[test(tokio::test)]
     async fn t_fetch_stream() {
         let rt = RustyTube::new();
-        let stream = rt.fetch_player("ZeerrnuLi5E", ClientType::Desktop).await.unwrap();
+        let stream = rt
+            .fetch_player("ZeerrnuLi5E", ClientType::Android)
+            .await
+            .unwrap();
 
         dbg!(stream);
     }

src/deobfuscate.rs

@@ -4,88 +4,63 @@ use log::debug;
 use once_cell::sync::Lazy;
 use reqwest::Client;
 use std::result::Result::Ok;
-use std::time::Instant;
-use tokio::sync::RwLock;
 
+use crate::cache::{Cache, DeobfData};
 use crate::util;
 
 pub struct Deobfuscator {
     http: Client,
-    cache: RwLock<JSCache>,
+    cache: Cache,
-}
-
-#[derive(Default)]
-struct JSCache {
-    last_update: Option<Instant>,
-    js_url: String,
-    sig_fn: String,
-    nsig_fn: String,
-    sts: String,
 }
 
 impl Deobfuscator {
     #[must_use]
-    pub fn new(http: Client) -> Self {
+    pub fn new(http: Client, cache: Cache) -> Self {
-        Self {
+        Self { http, cache }
-            http,
-            cache: RwLock::new(JSCache::default()),
-        }
     }
 
-    async fn update(&self) -> Result<()> {
+    async fn get_deobf_data(&self) -> Result<DeobfData> {
-        let mut cache = self.cache.write().await;
+        let http = self.http.clone();
 
-        if cache.is_stale() {
+        self.cache
-            let url = get_player_js_url(&self.http)
+            .get_deobf_data(async move {
-                .await
+                let js_url = get_player_js_url(&http)
-                .context("Failed to retrieve player.js URL")?;
+                    .await
+                    .context("Failed to retrieve player.js URL")?;
 
-            let player_js = get_response(&self.http, &url)
+                let player_js = get_response(&http, &js_url)
-                .await
+                    .await
-                .context("Failed to download player.js")?;
+                    .context("Failed to download player.js")?;
 
-            debug!("Downloaded player.js from {}", url);
+                debug!("Downloaded player.js from {}", js_url);
 
-            let sig_fn = get_sig_fn(&player_js)?;
+                let sig_fn = get_sig_fn(&player_js)?;
-            let nsig_fn = get_nsig_fn(&player_js)?;
+                let nsig_fn = get_nsig_fn(&player_js)?;
-            let sts = get_sts(&player_js)?;
+                let sts = get_sts(&player_js)?;
 
-            *cache = JSCache {
+                Ok(DeobfData {
-                last_update: Some(Instant::now()),
+                    js_url,
-                js_url: url.to_owned(),
+                    nsig_fn,
-                sig_fn,
+                    sig_fn,
-                nsig_fn,
+                    sts,
-                sts,
+                })
-            };
+            })
-        }
+            .await
-        Ok(())
     }
 
     pub async fn deobfuscate_sig(&self, sig: &str) -> Result<String> {
-        self.update().await?;
+        let deobf_data = self.get_deobf_data().await?;
-        let cache = self.cache.read().await;
+        deobfuscate_sig(sig, &deobf_data.sig_fn)
-        deobfuscate_sig(sig, &cache.sig_fn)
     }
 
     pub async fn deobfuscate_nsig(&self, nsig: &str) -> Result<String> {
-        self.update().await?;
+        let deobf_data = self.get_deobf_data().await?;
-        let cache = self.cache.read().await;
+        deobfuscate_nsig(nsig, &deobf_data.nsig_fn)
-        deobfuscate_nsig(nsig, &cache.nsig_fn)
     }
 
     pub async fn get_sts(&self) -> Result<String> {
-        self.update().await?;
+        let deobf_data = self.get_deobf_data().await?;
-        let cache = self.cache.read().await;
+        Ok(deobf_data.sts)
-        Ok(cache.sts.to_owned())
-    }
-}
-
-impl JSCache {
-    fn is_stale(&self) -> bool {
-        match self.last_update {
-            Some(last_update) => Instant::now().duration_since(last_update).as_secs() > 3600,
-            None => true,
-        }
     }
 }
 
@@ -407,7 +382,8 @@ c[36](c[8],c[32]),c[20](c[25],c[10]),c[2](c[22],c[8]),c[32](c[20],c[16]),c[32](c
     #[test(tokio::test)]
     async fn t_update() {
         let client = Client::new();
-        let deobf = Deobfuscator::new(client);
+        let cache = Cache::default();
+        let deobf = Deobfuscator::new(client, cache);
 
         let deobf_sig = deobf.deobfuscate_sig("GOqGOqGOq0QJ8wRAIgaryQHfplJ9xJSKFywyaSMHuuwZYsoMTAvRvfm51qIGECIA5061zWeyfMPX9hEl_U6f9J0tr7GTJMKyPf5XNrJb5fb5i").await.unwrap();
         println!("{}", deobf_sig);
@@ -416,7 +392,8 @@ c[36](c[8],c[32]),c[20](c[25],c[10]),c[2](c[22],c[8]),c[32](c[20],c[16]),c[32](c
     #[test(tokio::test)]
     async fn t_parallel() {
         let client = Client::new();
-        let deobf = Deobfuscator::new(client);
+        let cache = Cache::default();
+        let deobf = Deobfuscator::new(client, cache);
         let deobf_arc = Arc::new(deobf);
 
         let (deobf_sig, deobf_nsig) = tokio::join!(