Improved the program, and also fixed the sig function not working

src/jobs.rs

@@ -1,12 +1,11 @@
-use regex::Regex;
 use rquickjs::{async_with, AsyncContext, AsyncRuntime};
 use std::{num::NonZeroUsize, sync::Arc, thread::available_parallelism};
 use tokio::{io::AsyncWriteExt, runtime::Handle, sync::Mutex, task::block_in_place};
 use tub::Pool;
 
-use crate::consts::{
+use crate::{
-    NSIG_FUNCTION_ARRAY, NSIG_FUNCTION_NAME, REGEX_HELPER_OBJ_NAME, REGEX_PLAYER_ID,
+    consts::NSIG_FUNCTION_NAME,
-    REGEX_SIGNATURE_FUNCTION, REGEX_SIGNATURE_TIMESTAMP, TEST_YOUTUBE_VIDEO,
+    player::{fetch_update, FetchUpdateStatus},
 };
 
 pub enum JobOpcode {
@@ -41,11 +40,11 @@ impl From<u8> for JobOpcode {
 }
 
 pub struct PlayerInfo {
-    nsig_function_code: String,
+    pub nsig_function_code: String,
-    sig_function_code: String,
+    pub sig_function_code: String,
-    sig_function_name: String,
+    pub sig_function_name: String,
-    signature_timestamp: u64,
+    pub signature_timestamp: u64,
-    player_id: u32,
+    pub player_id: u32,
 }
 
 pub struct JavascriptInterpreter {
@@ -81,7 +80,7 @@ impl JavascriptInterpreter {
 }
 
 pub struct GlobalState {
-    player_info: Mutex<PlayerInfo>,
+    pub player_info: Mutex<PlayerInfo>,
     js_runtime_pool: Pool<Arc<JavascriptInterpreter>>,
 }
 
@@ -128,180 +127,26 @@ pub async fn process_fetch_update<W>(
     let mut writer;
 
     let global_state = state.clone();
-    let response = match reqwest::get(TEST_YOUTUBE_VIDEO).await {
+
-        Ok(req) => req.text().await.unwrap(),
+    match fetch_update(global_state).await {
-        Err(x) => {
+        Ok(_x) => {
-            println!("Could not fetch the test video: {}", x);
             writer = cloned_writer.lock().await;
-            write_failure!(writer, request_id);
+            writer.write_u32(request_id).await;
-            return;
+            // sync code to tell the client the player had updated
+            writer.write_u16(0xF44F).await;
+            println!("Successfully updated the player");
         }
-    };
+        Err(FetchUpdateStatus::PlayerAlreadyUpdated) => {
-
-    let player_id_str = match REGEX_PLAYER_ID.captures(&response).unwrap().get(1) {
-        Some(result) => result.as_str(),
-        None => {
             writer = cloned_writer.lock().await;
-            write_failure!(writer, request_id);
+            writer.write_u32(request_id).await;
-            return;
+            writer.write_u16(0xFFFF).await;
+        }
+        Err(_x) => {
+            writer = cloned_writer.lock().await;
+            writer.write_u32(request_id).await;
+            writer.write_u16(0).await;
         }
-    };
-
-    let player_id: u32 = u32::from_str_radix(player_id_str, 16).unwrap();
-
-    let mut current_player_info = global_state.player_info.lock().await;
-    let current_player_id = current_player_info.player_id;
-    // release the mutex for other tasks
-    drop(current_player_info);
-
-    if player_id == current_player_id {
-        // Player is already up to date
-        writer = cloned_writer.lock().await;
-        writer.write_u32(request_id).await;
-        writer.write_u16(0xFFFF).await;
-        return;
     }
-
-    // Download the player script
-    let player_js_url: String = format!(
-        "https://www.youtube.com/s/player/{:08x}/player_ias.vflset/en_US/base.js",
-        player_id
-    );
-    let player_javascript = match reqwest::get(player_js_url).await {
-        Ok(req) => req.text().await.unwrap(),
-        Err(x) => {
-            println!("Could not fetch the player JS: {}", x);
-            writer = cloned_writer.lock().await;
-            write_failure!(writer, request_id);
-            return;
-        }
-    };
-
-    let nsig_function_array = NSIG_FUNCTION_ARRAY.captures(&player_javascript).unwrap();
-    let nsig_array_name = nsig_function_array.get(1).unwrap().as_str();
-    let nsig_array_value = nsig_function_array
-        .get(2)
-        .unwrap()
-        .as_str()
-        .parse::<usize>()
-        .unwrap();
-
-    let mut nsig_array_context_regex: String = String::new();
-    nsig_array_context_regex += "var ";
-    nsig_array_context_regex += nsig_array_name;
-    nsig_array_context_regex += "\\s*=\\s*\\[(.+?)][;,]";
-
-    let nsig_array_context = match Regex::new(&nsig_array_context_regex) {
-        Ok(x) => x,
-        Err(x) => {
-            println!("Error: nsig regex compilation failed: {}", x);
-            writer = cloned_writer.lock().await;
-            write_failure!(writer, request_id);
-            return;
-        }
-    };
-
-    let array_content = nsig_array_context
-        .captures(&player_javascript)
-        .unwrap()
-        .get(1)
-        .unwrap()
-        .as_str()
-        .split(',');
-
-    let array_values: Vec<&str> = array_content.collect();
-
-    let nsig_function_name = array_values.get(nsig_array_value).unwrap();
-
-    // Extract nsig function code
-    let mut nsig_function_code_regex_str: String = String::new();
-    nsig_function_code_regex_str += nsig_function_name;
-    nsig_function_code_regex_str +=
-        "=\\s*function([\\S\\s]*?\\}\\s*return [\\w$]+?\\.join\\(\"\"\\)\\s*\\};)";
-
-    let nsig_function_code_regex = Regex::new(&nsig_function_code_regex_str).unwrap();
-
-    let mut nsig_function_code = String::new();
-    nsig_function_code += "function ";
-    nsig_function_code += NSIG_FUNCTION_NAME;
-    nsig_function_code += nsig_function_code_regex
-        .captures(&player_javascript)
-        .unwrap()
-        .get(1)
-        .unwrap()
-        .as_str();
-
-    // Extract signature function name
-    let sig_function_name = REGEX_SIGNATURE_FUNCTION
-        .captures(&player_javascript)
-        .unwrap()
-        .get(1)
-        .unwrap()
-        .as_str();
-
-    let mut sig_function_body_regex_str: String = String::new();
-    sig_function_body_regex_str += sig_function_name;
-    sig_function_body_regex_str += "=function\\([a-zA-Z0-9_]+\\)\\{.+?\\}";
-
-    let sig_function_body_regex = Regex::new(&sig_function_body_regex_str).unwrap();
-
-    let sig_function_body = sig_function_body_regex
-        .captures(&player_javascript)
-        .unwrap()
-        .get(0)
-        .unwrap()
-        .as_str();
-
-    // Get the helper object
-    let helper_object_name = REGEX_HELPER_OBJ_NAME
-        .captures(sig_function_body)
-        .unwrap()
-        .get(1)
-        .unwrap()
-        .as_str();
-
-    let mut helper_object_body_regex_str = String::new();
-    helper_object_body_regex_str += "var ";
-    helper_object_body_regex_str += helper_object_name;
-    helper_object_body_regex_str += "=\\{(?>.|\\n)+?\\}\\};";
-
-    let helper_object_body_regex = Regex::new(&helper_object_body_regex_str).unwrap();
-    let helper_object_body = helper_object_body_regex
-        .captures(&player_javascript)
-        .unwrap()
-        .get(0)
-        .unwrap()
-        .as_str();
-
-    let mut sig_code = String::new();
-    sig_code += helper_object_body;
-    sig_code += sig_function_body;
-
-    // Get signature timestamp
-    let signature_timestamp: u64 = REGEX_SIGNATURE_TIMESTAMP
-        .captures(&player_javascript)
-        .unwrap()
-        .get(1)
-        .unwrap()
-        .as_str()
-        .parse()
-        .unwrap();
-
-    current_player_info = global_state.player_info.lock().await;
-    current_player_info.player_id = player_id;
-    current_player_info.nsig_function_code = nsig_function_code;
-    current_player_info.sig_function_code = sig_code;
-    current_player_info.sig_function_name = sig_function_name.to_string();
-    current_player_info.signature_timestamp = signature_timestamp;
-
-    writer = cloned_writer.lock().await;
-    writer.write_u32(request_id).await;
-    // sync code to tell the client the player had updated
-    writer.write_u16(0xF44F).await;
-
-    writer.flush().await;
-
-    println!("Successfully updated the player");
 }
 
 pub async fn process_decrypt_n_signature<W>(

src/main.rs

@@ -1,10 +1,13 @@
 mod consts;
 mod jobs;
+mod player;
 
 use consts::DEFAULT_SOCK_PATH;
 use jobs::{process_decrypt_n_signature, process_fetch_update, GlobalState, JobOpcode};
+use player::fetch_update;
 use std::{env::args, io::Error, sync::Arc};
 use tokio::{
+    fs::remove_file,
     io::{self, AsyncReadExt, BufReader, BufWriter},
     net::{UnixListener, UnixStream},
     sync::Mutex,
@@ -51,8 +54,26 @@ async fn main() {
     // have to please rust
     let state: Arc<GlobalState> = Arc::new(GlobalState::new());
 
-    let socket = UnixListener::bind(socket_url).unwrap();
+    let socket = match UnixListener::bind(socket_url) {
+        Ok(x) => x,
+        Err(x) => {
+            if x.kind() == std::io::ErrorKind::AddrInUse {
+                remove_file(socket_url).await.unwrap();
+                UnixListener::bind(socket_url).unwrap()
+            } else {
+                println!("Error occurred while trying to bind: {}", x);
+                return;
+            }
+        }
+    };
 
+    println!("Fetching player");
+    match fetch_update(state.clone()).await {
+        Ok(()) => println!("Successfully fetched player"),
+        Err(x) => {
+            println!("Error occured while trying to fetch the player: {:?}", x);
+        }
+    }
     loop {
         let (socket, _addr) = socket.accept().await.unwrap();
 

src/player.rs

@@ -0,0 +1,178 @@
+use std::sync::Arc;
+
+use regex::Regex;
+
+use crate::{
+    consts::{
+        NSIG_FUNCTION_ARRAY, NSIG_FUNCTION_NAME, REGEX_HELPER_OBJ_NAME, REGEX_PLAYER_ID,
+        REGEX_SIGNATURE_FUNCTION, REGEX_SIGNATURE_TIMESTAMP, TEST_YOUTUBE_VIDEO,
+    },
+    jobs::GlobalState,
+};
+
+// TODO: too lazy to make proper debugging print
+#[derive(Debug)]
+pub enum FetchUpdateStatus {
+    CannotFetchTestVideo,
+    CannotMatchPlayerID,
+    CannotFetchPlayerJS,
+    NsigRegexCompileFailed,
+    PlayerAlreadyUpdated,
+}
+
+pub async fn fetch_update(state: Arc<GlobalState>) -> Result<(), FetchUpdateStatus> {
+    let global_state = state.clone();
+    let response = match reqwest::get(TEST_YOUTUBE_VIDEO).await {
+        Ok(req) => req.text().await.unwrap(),
+        Err(x) => {
+            println!("Could not fetch the test video: {}", x);
+            return Err(FetchUpdateStatus::CannotFetchTestVideo);
+        }
+    };
+
+    let player_id_str = match REGEX_PLAYER_ID.captures(&response).unwrap().get(1) {
+        Some(result) => result.as_str(),
+        None => return Err(FetchUpdateStatus::CannotMatchPlayerID),
+    };
+
+    let player_id: u32 = u32::from_str_radix(player_id_str, 16).unwrap();
+
+    let mut current_player_info = global_state.player_info.lock().await;
+    let current_player_id = current_player_info.player_id;
+    // release the mutex for other tasks
+    drop(current_player_info);
+
+    if player_id == current_player_id {
+        return Err(FetchUpdateStatus::PlayerAlreadyUpdated);
+    }
+
+    // Download the player script
+    let player_js_url: String = format!(
+        "https://www.youtube.com/s/player/{:08x}/player_ias.vflset/en_US/base.js",
+        player_id
+    );
+    let player_javascript = match reqwest::get(player_js_url).await {
+        Ok(req) => req.text().await.unwrap(),
+        Err(x) => {
+            println!("Could not fetch the player JS: {}", x);
+            return Err(FetchUpdateStatus::CannotFetchPlayerJS);
+        }
+    };
+
+    let nsig_function_array = NSIG_FUNCTION_ARRAY.captures(&player_javascript).unwrap();
+    let nsig_array_name = nsig_function_array.get(1).unwrap().as_str();
+    let nsig_array_value = nsig_function_array
+        .get(2)
+        .unwrap()
+        .as_str()
+        .parse::<usize>()
+        .unwrap();
+
+    let mut nsig_array_context_regex: String = String::new();
+    nsig_array_context_regex += "var ";
+    nsig_array_context_regex += nsig_array_name;
+    nsig_array_context_regex += "\\s*=\\s*\\[(.+?)][;,]";
+
+    let nsig_array_context = match Regex::new(&nsig_array_context_regex) {
+        Ok(x) => x,
+        Err(x) => {
+            println!("Error: nsig regex compilation failed: {}", x);
+            return Err(FetchUpdateStatus::NsigRegexCompileFailed);
+        }
+    };
+
+    let array_content = nsig_array_context
+        .captures(&player_javascript)
+        .unwrap()
+        .get(1)
+        .unwrap()
+        .as_str()
+        .split(',');
+
+    let array_values: Vec<&str> = array_content.collect();
+
+    let nsig_function_name = array_values.get(nsig_array_value).unwrap();
+
+    // Extract nsig function code
+    let mut nsig_function_code_regex_str: String = String::new();
+    nsig_function_code_regex_str += nsig_function_name;
+    nsig_function_code_regex_str +=
+        "=\\s*function([\\S\\s]*?\\}\\s*return [\\w$]+?\\.join\\(\"\"\\)\\s*\\};)";
+
+    let nsig_function_code_regex = Regex::new(&nsig_function_code_regex_str).unwrap();
+
+    let mut nsig_function_code = String::new();
+    nsig_function_code += "function ";
+    nsig_function_code += NSIG_FUNCTION_NAME;
+    nsig_function_code += nsig_function_code_regex
+        .captures(&player_javascript)
+        .unwrap()
+        .get(1)
+        .unwrap()
+        .as_str();
+
+    // Extract signature function name
+    let sig_function_name = REGEX_SIGNATURE_FUNCTION
+        .captures(&player_javascript)
+        .unwrap()
+        .get(1)
+        .unwrap()
+        .as_str();
+
+    let mut sig_function_body_regex_str: String = String::new();
+    sig_function_body_regex_str += sig_function_name;
+    sig_function_body_regex_str += "=function\\([a-zA-Z0-9_]+\\)\\{.+?\\}";
+
+    let sig_function_body_regex = Regex::new(&sig_function_body_regex_str).unwrap();
+
+    let sig_function_body = sig_function_body_regex
+        .captures(&player_javascript)
+        .unwrap()
+        .get(0)
+        .unwrap()
+        .as_str();
+
+    // Get the helper object
+    let helper_object_name = REGEX_HELPER_OBJ_NAME
+        .captures(sig_function_body)
+        .unwrap()
+        .get(1)
+        .unwrap()
+        .as_str();
+
+    let mut helper_object_body_regex_str = String::new();
+    helper_object_body_regex_str += "(var ";
+    helper_object_body_regex_str += helper_object_name;
+    helper_object_body_regex_str += "=\\{(?:.|\\n)+?\\}\\};)";
+
+    let helper_object_body_regex = Regex::new(&helper_object_body_regex_str).unwrap();
+    let helper_object_body = helper_object_body_regex
+        .captures(&player_javascript)
+        .unwrap()
+        .get(0)
+        .unwrap()
+        .as_str();
+
+    let mut sig_code = String::new();
+    sig_code += helper_object_body;
+    sig_code += sig_function_body;
+
+    // Get signature timestamp
+    let signature_timestamp: u64 = REGEX_SIGNATURE_TIMESTAMP
+        .captures(&player_javascript)
+        .unwrap()
+        .get(1)
+        .unwrap()
+        .as_str()
+        .parse()
+        .unwrap();
+
+    current_player_info = global_state.player_info.lock().await;
+    current_player_info.player_id = player_id;
+    current_player_info.nsig_function_code = nsig_function_code;
+    current_player_info.sig_function_code = sig_code;
+    current_player_info.sig_function_name = sig_function_name.to_string();
+    current_player_info.signature_timestamp = signature_timestamp;
+
+    Ok(())
+}