use std::sync::{atomic::AtomicU32, Arc, RwLock}; use once_cell::sync::Lazy; use rand::Rng; use regex::Regex; use reqwest::{header, Client}; use crate::{ client::YOUTUBE_MUSIC_HOME_URL, error::{Error, ExtractionError}, util, }; /// To increase privacy and possibly circumvent rate limits, RustyPipe uses multiple /// visitor data IDs. These are held in this cache object. /// /// On instantiation, the cache is empty, so for the first requests new visitor data IDs /// have to be requested. For subsequent requests a random ID from the cache is picked. /// After req_limit requests, a new token is requested asynchronously and added to the cache /// to prevent the IDs from being overused. /// /// The cache holds a maximum of 100 visitor data IDs. If more are added, the oldest ones /// are evicted. #[derive(Clone)] pub struct VisitorDataCache { inner: Arc, } struct VisitorDataCacheRef { req_counter: AtomicU32, visitor_data: RwLock>, http: Client, } static VISITOR_DATA_REGEX: Lazy = Lazy::new(|| Regex::new(r#""visitorData":"([\w\d_\-%]+?)""#).unwrap()); /// Number of requests after which a new token is requested const REQ_LIMIT: u32 = 10; /// Maximum size of the cache (-1) const MAX_SIZE: usize = 99; impl VisitorDataCache { pub fn new(http: Client) -> Self { Self { inner: VisitorDataCacheRef { req_counter: Default::default(), visitor_data: Default::default(), http, } .into(), } } async fn get_visitor_data(&self) -> Result { tracing::debug!("getting YT visitor data"); let resp = self .inner .http .get(YOUTUBE_MUSIC_HOME_URL) .header(header::ORIGIN, YOUTUBE_MUSIC_HOME_URL) .header(header::REFERER, YOUTUBE_MUSIC_HOME_URL) .send() .await?; let vdata = resp .headers() .get_all(header::SET_COOKIE) .iter() .find_map(|c| { if let Ok(cookie) = c.to_str() { if let Some(after) = cookie.strip_prefix("__Secure-YEC=") { return after .split_once(';') .map(|s| s.0.to_owned()) .filter(|s| !s.is_empty()); } } None }); match vdata { Some(vdata) => Ok(vdata), None => { if resp.status().is_success() { // Extract visitor data from html let html = resp.text().await?; util::get_cg_from_regex(&VISITOR_DATA_REGEX, &html, 1).ok_or(Error::Extraction( ExtractionError::InvalidData( "Could not find visitor data on html page".into(), ), )) } else { Err(Error::Extraction(ExtractionError::InvalidData( format!("Could not get visitor data, status: {}", resp.status()).into(), ))) } } } } pub async fn new_visitor_data(&self) -> Result { self.inner .req_counter .store(0, std::sync::atomic::Ordering::SeqCst); let vd = self.get_visitor_data().await.unwrap(); let mut vds = self.inner.visitor_data.write().unwrap(); for _ in 0..(vds.len().saturating_sub(MAX_SIZE)) { let rem = vds.remove(0); tracing::debug!("visitor data {rem} removed from cache"); } vds.push(vd.to_owned()); tracing::debug!("visitor data {} added to cache ({} ids)", vd, vds.len()); Ok(vd) } pub async fn get(&self) -> Result { // Request new visitor data in the background every 10 requests if self .inner .req_counter .fetch_add(1, std::sync::atomic::Ordering::SeqCst) >= REQ_LIMIT { let nc = self.clone(); tokio::spawn(async move { nc.new_visitor_data().await }); } { let vds = self.inner.visitor_data.read().unwrap(); if !vds.is_empty() { let mut rng = rand::thread_rng(); let vd = vds[rng.gen_range(0..vds.len())].to_owned(); tracing::debug!("visitor data {vd} picked from cache"); return Ok(vd); } } // Fetch new visitor data if the cache is empty self.new_visitor_data().await } } #[cfg(test)] mod tests { use std::time::Duration; use crate::client::DEFAULT_UA; use super::*; use tracing_test::traced_test; #[tokio::test] #[traced_test] async fn get_visitor_data() { let cache = VisitorDataCache::new(Client::builder().user_agent(DEFAULT_UA).build().unwrap()); // Get initial visitor data let v1 = cache.get().await.unwrap(); // Run as many request as necessary to fetch second visitor data for _ in 0..=REQ_LIMIT { let got = cache.get().await.unwrap(); assert_eq!(got, v1); } // Second visitor data does not arrive instantly, request immediately after returns the first data let vds_len = cache.inner.visitor_data.read().unwrap().len(); assert_eq!(vds_len, 1); // Wait for the second visitor data to arrive tokio::time::sleep(Duration::from_millis(1000)).await; let vds_len = cache.inner.visitor_data.read().unwrap().len(); assert_eq!(vds_len, 2); } }