Upload 2 files

src/main.rs

@@ -10,7 +10,7 @@ use http::{Method, header::CONTENT_TYPE};
 use serde_json::json;
 use serde::Deserialize;
 use std::collections::HashMap;
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 use bytes::Bytes;
 use futures_util::StreamExt;
 use axum::body::Body;
@@ -22,6 +22,7 @@ use reqwest::{cookie::Jar, Url};
 use reqwest::cookie::CookieStore;
 use reqwest::header::ACCEPT;
 use std::time::{SystemTime, UNIX_EPOCH};
+use tokio::sync::{Mutex, RwLock};
 
 type BoxErr = Box<dyn std::error::Error + Send + Sync + 'static>;
 
@@ -30,7 +31,8 @@ use api::{APIClient, APIError, Format};
 
 #[derive(Clone)]
 struct AppState {
-    api: Arc<RwLock<APIClient>>,
+    api: Arc<Mutex<APIClient>>,
+    api_by_arl: Arc<RwLock<HashMap<String, Arc<Mutex<APIClient>>>>>,
     pair: Arc<RwLock<HashMap<String, PairSession>>>,
 }
 
@@ -38,6 +40,23 @@ fn json_error(status: StatusCode, message: impl Into<String>) -> axum::response:
     (status, Json(json!({ "error": message.into() }))).into_response()
 }
 
+async fn api_client_for_arl(state: &AppState, arl: Option<String>) -> Arc<Mutex<APIClient>> {
+    let Some(arl) = arl.map(|s| s.trim().to_string()).filter(|s| !s.is_empty()) else {
+        return state.api.clone();
+    };
+
+    let mut map = state.api_by_arl.write().await;
+    if let Some(c) = map.get(&arl) {
+        return c.clone();
+    }
+    if map.len() >= 32 {
+        map.clear();
+    }
+    let c = Arc::new(Mutex::new(APIClient::new_with_arl(arl.clone())));
+    map.insert(arl, c.clone());
+    c
+}
+
 #[derive(Deserialize)]
 struct DeezerTrackList {
     data: Vec<DeezerTrack>
@@ -68,38 +87,36 @@ async fn get_url(State(state): State<AppState>, Json(req): Json<RequestParams>)
         return (StatusCode::BAD_REQUEST, "ID list cannot be empty".to_string());
     }
 
-    let mut client = if let Some(arl) = req.arl.clone().filter(|s| !s.trim().is_empty()) {
-        APIClient::new_with_arl(arl)
-    } else {
-        state.api.read().unwrap().clone()
-    };
-    let old_license = client.license_token.clone();
-
-    let resp: Result<DeezerTrackList, APIError> = client.api_call("song.getListData", &json!({"sng_ids":req.ids,"array_default":["SNG_ID","TRACK_TOKEN"]})).await;
-    let track_list;
-    match resp {
-        Ok(t) => track_list = t,
-        Err(e) => return (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
-    };
+    let client = api_client_for_arl(&state, req.arl.clone()).await;
+    let media_resp = {
+        let mut client = client.lock().await;
+        let resp: Result<DeezerTrackList, APIError> = client
+            .api_call(
+                "song.getListData",
+                &json!({"sng_ids":req.ids,"array_default":["SNG_ID","TRACK_TOKEN"]}),
+            )
+            .await;
+        let track_list = match resp {
+            Ok(t) => t,
+            Err(e) => return (StatusCode::SERVICE_UNAVAILABLE, e.to_string()),
+        };
 
-    if track_list.data.is_empty() {
-        return (StatusCode::BAD_REQUEST, "No valid IDs found".to_string());
-    }
+        if track_list.data.is_empty() {
+            return (StatusCode::BAD_REQUEST, "No valid IDs found".to_string());
+        }
 
-    let track_tokens: Vec<&str> = track_list.data.iter().map(|t| t.TRACK_TOKEN.as_str()).collect();
+        let track_tokens: Vec<&str> = track_list
+            .data
+            .iter()
+            .map(|t| t.TRACK_TOKEN.as_str())
+            .collect();
 
-    let media_result = client.get_media(&req.formats, track_tokens).await;
-    let media_resp;
-    match media_result {
-        Ok(r) => media_resp = r,
-        Err(_) => return (StatusCode::INTERNAL_SERVER_ERROR, "Error while getting response from media.deezer.com".to_string()),
+        match client.get_media(&req.formats, track_tokens).await {
+            Ok(r) => r,
+            Err(e) => return (StatusCode::SERVICE_UNAVAILABLE, e.to_string()),
+        }
     };
 
-    if req.arl.is_none() && client.license_token != old_license {
-        let mut client_write = state.api.write().unwrap();
-        *client_write = client;
-    }
-
     (StatusCode::OK, media_resp.text().await.unwrap())
 }
 
@@ -183,47 +200,74 @@ async fn stream(State(state): State<AppState>, Query(q): Query<StreamParams>) ->
         "AUTO" | _ => vec![Format::FLAC, Format::MP3_320, Format::MP3_128],
     };
 
-    let mut client = if let Some(arl) = q.arl.clone().filter(|s| !s.trim().is_empty()) {
-        APIClient::new_with_arl(arl)
-    } else {
-        state.api.read().unwrap().clone()
-    };
+    fn extract_media_url(media_json: &serde_json::Value) -> String {
+        let data = media_json.get("data").and_then(|v| v.as_array());
+        let Some(data0) = data.and_then(|a| a.get(0)) else {
+            return String::new();
+        };
+        let media_arr = data0.get("media").and_then(|v| v.as_array());
+        let Some(media_arr) = media_arr else {
+            return String::new();
+        };
+        for m in media_arr {
+            if let Some(sources) = m.get("sources").and_then(|v| v.as_array()) {
+                for s in sources {
+                    if let Some(url) = s.get("url").and_then(|v| v.as_str()) {
+                        let url = url.trim();
+                        if !url.is_empty() {
+                            return url.to_string();
+                        }
+                    }
+                }
+            }
+        }
+        String::new()
+    }
 
-    let resp: Result<DeezerTrackList, APIError> = client
-        .api_call(
-            "song.getListData",
-            &json!({"sng_ids":[q.id],"array_default":["SNG_ID","TRACK_TOKEN"]}),
-        )
-        .await;
+    let client = api_client_for_arl(&state, q.arl.clone()).await;
+    let url = {
+        let mut client = client.lock().await;
+        let resp: Result<DeezerTrackList, APIError> = client
+            .api_call(
+                "song.getListData",
+                &json!({"sng_ids":[q.id],"array_default":["SNG_ID","TRACK_TOKEN"]}),
+            )
+            .await;
 
-    let track_list = match resp {
-        Ok(t) => t,
-        Err(e) => return (StatusCode::BAD_GATEWAY, e.to_string()).into_response(),
-    };
+        let track_list = match resp {
+            Ok(t) => t,
+            Err(e) => return (StatusCode::SERVICE_UNAVAILABLE, e.to_string()).into_response(),
+        };
 
-    if track_list.data.is_empty() {
-        return (StatusCode::BAD_REQUEST, "No valid ID".to_string()).into_response();
-    }
+        if track_list.data.is_empty() {
+            return (StatusCode::BAD_REQUEST, "No valid ID".to_string()).into_response();
+        }
 
-    let track_token = track_list.data[0].TRACK_TOKEN.as_str();
-    let media_resp = match client.get_media(&formats, vec![track_token]).await {
-        Ok(r) => r,
-        Err(e) => return (StatusCode::BAD_GATEWAY, e.to_string()).into_response(),
-    };
+        let track_token = track_list.data[0].TRACK_TOKEN.as_str();
+        let media_resp = match client.get_media(&formats, vec![track_token]).await {
+            Ok(r) => r,
+            Err(e) => return (StatusCode::SERVICE_UNAVAILABLE, e.to_string()).into_response(),
+        };
 
-    let media_json: serde_json::Value = match media_resp.json().await {
-        Ok(v) => v,
-        Err(_) => return (StatusCode::BAD_GATEWAY, "Bad media response".to_string()).into_response(),
-    };
+        let media_json: serde_json::Value = match media_resp.json().await {
+            Ok(v) => v,
+            Err(_) => {
+                return (StatusCode::SERVICE_UNAVAILABLE, "Bad media response".to_string())
+                    .into_response()
+            }
+        };
 
-    let url = media_json["data"][0]["media"][0]["sources"][0]["url"]
-        .as_str()
-        .unwrap_or("")
-        .to_string();
+        let url = extract_media_url(&media_json);
 
-    if url.is_empty() {
-        return (StatusCode::BAD_GATEWAY, "No url in media response".to_string()).into_response();
-    }
+        if url.is_empty() {
+            return (
+                StatusCode::SERVICE_UNAVAILABLE,
+                "No url in media response".to_string(),
+            )
+                .into_response();
+        }
+        url
+    };
 
     let http = reqwest::Client::builder()
         .timeout(std::time::Duration::from_secs(60))
@@ -232,11 +276,18 @@ async fn stream(State(state): State<AppState>, Query(q): Query<StreamParams>) ->
 
     let upstream = match http.get(url).send().await {
         Ok(r) => r,
-        Err(_) => return (StatusCode::BAD_GATEWAY, "Upstream request failed".to_string()).into_response(),
+        Err(_) => {
+            return (StatusCode::SERVICE_UNAVAILABLE, "Upstream request failed".to_string())
+                .into_response()
+        }
     };
 
     if !upstream.status().is_success() {
-        return (StatusCode::BAD_GATEWAY, "Upstream status not OK".to_string()).into_response();
+        return (
+            StatusCode::SERVICE_UNAVAILABLE,
+            "Upstream status not OK".to_string(),
+        )
+            .into_response();
     }
 
     let key = blowfish_key(q.id);
@@ -281,16 +332,19 @@ struct UserDataReq {
     arl: String,
 }
 
-async fn user_data(Json(req): Json<UserDataReq>) -> impl IntoResponse {
+async fn user_data(State(state): State<AppState>, Json(req): Json<UserDataReq>) -> impl IntoResponse {
     let arl = req.arl.trim().to_string();
     if arl.is_empty() {
         return json_error(StatusCode::BAD_REQUEST, "ARL required");
     }
 
-    let mut client = APIClient::new_with_arl(arl);
-    let data = match client.user_data().await {
-        Ok(v) => v,
-        Err(e) => return json_error(StatusCode::BAD_GATEWAY, e.to_string()),
+    let client = api_client_for_arl(&state, Some(arl)).await;
+    let data = {
+        let mut client = client.lock().await;
+        match client.user_data().await {
+            Ok(v) => v,
+            Err(e) => return json_error(StatusCode::SERVICE_UNAVAILABLE, e.to_string()),
+        }
     };
 
     let token = data["checkForm"].as_str().unwrap_or("").to_string();
@@ -306,33 +360,37 @@ struct PlaylistsReq {
     nb: Option<u32>,
 }
 
-async fn playlists(Json(req): Json<PlaylistsReq>) -> impl IntoResponse {
+async fn playlists(State(state): State<AppState>, Json(req): Json<PlaylistsReq>) -> impl IntoResponse {
     let arl = req.arl.trim().to_string();
     if arl.is_empty() {
         return json_error(StatusCode::BAD_REQUEST, "ARL required");
     }
 
-    let mut client = APIClient::new_with_arl(arl);
-    let ud = match client.user_data().await {
-        Ok(v) => v,
-        Err(e) => return json_error(StatusCode::BAD_GATEWAY, e.to_string()),
-    };
+    let client = api_client_for_arl(&state, Some(arl)).await;
+    let res = {
+        let mut client = client.lock().await;
+        let ud = match client.user_data().await {
+            Ok(v) => v,
+            Err(e) => return json_error(StatusCode::SERVICE_UNAVAILABLE, e.to_string()),
+        };
 
-    let uid = ud["USER"]["USER_ID"].as_i64().unwrap_or(0);
-    if uid <= 0 {
-        return json_error(StatusCode::BAD_GATEWAY, "No USER_ID");
-    }
+        let uid = ud["USER"]["USER_ID"].as_i64().unwrap_or(0);
+        if uid <= 0 {
+            return json_error(StatusCode::SERVICE_UNAVAILABLE, "No USER_ID");
+        }
 
-    let start = req.start.unwrap_or(0);
-    let nb = req.nb.unwrap_or(50);
+        let start = req.start.unwrap_or(0);
+        let nb = req.nb.unwrap_or(50);
 
-    let res: Result<serde_json::Value, APIError> = client
-        .api_call("playlist.getList", &json!({"user_id":uid,"nb":nb,"start":start}))
-        .await;
+        let res: Result<serde_json::Value, APIError> = client
+            .api_call("playlist.getList", &json!({"user_id":uid,"nb":nb,"start":start}))
+            .await;
+        res
+    };
 
     match res {
         Ok(v) => (StatusCode::OK, Json(v)).into_response(),
-        Err(e) => json_error(StatusCode::BAD_GATEWAY, e.to_string()),
+        Err(e) => json_error(StatusCode::SERVICE_UNAVAILABLE, e.to_string()),
     }
 }
 
@@ -344,32 +402,38 @@ struct PlaylistTracksReq {
     nb: Option<u32>,
 }
 
-async fn playlist_tracks(Json(req): Json<PlaylistTracksReq>) -> impl IntoResponse {
+async fn playlist_tracks(
+    State(state): State<AppState>,
+    Json(req): Json<PlaylistTracksReq>,
+) -> impl IntoResponse {
     let arl = req.arl.trim().to_string();
     if arl.is_empty() {
         return json_error(StatusCode::BAD_REQUEST, "ARL required");
     }
 
-    let mut client = APIClient::new_with_arl(arl);
-    let _ = client.user_data().await;
+    let client = api_client_for_arl(&state, Some(arl)).await;
 
     let start = req.start.unwrap_or(0);
     let nb = req.nb.unwrap_or(100);
 
-    let res: Result<serde_json::Value, APIError> = client
-        .api_call(
-            "playlist.getSongs",
-            &json!({
-                "playlist_id": req.playlist_id,
-                "start": start,
-                "nb": nb
-            }),
-        )
-        .await;
+    let res: Result<serde_json::Value, APIError> = {
+        let mut client = client.lock().await;
+        let _ = client.user_data().await;
+        client
+            .api_call(
+                "playlist.getSongs",
+                &json!({
+                    "playlist_id": req.playlist_id,
+                    "start": start,
+                    "nb": nb
+                }),
+            )
+            .await
+    };
 
     match res {
         Ok(v) => (StatusCode::OK, Json(v)).into_response(),
-        Err(e) => json_error(StatusCode::BAD_GATEWAY, e.to_string()),
+        Err(e) => json_error(StatusCode::SERVICE_UNAVAILABLE, e.to_string()),
     }
 }
 
@@ -660,7 +724,7 @@ async fn pair_start(State(state): State<AppState>) -> impl IntoResponse {
         .map(|d| d.as_millis())
         .unwrap_or(0);
     {
-        let mut map = state.pair.write().unwrap();
+        let mut map = state.pair.write().await;
         map.insert(
             code.clone(),
             PairSession {
@@ -682,7 +746,7 @@ async fn pair_status(State(state): State<AppState>, Query(q): Query<PairStatusQu
         return json_error(StatusCode::BAD_REQUEST, "code required");
     }
 
-    let map = state.pair.read().unwrap();
+    let map = state.pair.read().await;
     let Some(sess) = map.get(&code) else {
         return json_error(StatusCode::NOT_FOUND, "no such code");
     };
@@ -794,7 +858,7 @@ async fn pair_complete(State(state): State<AppState>, Form(f): Form<PairComplete
     }
 
     {
-        let map = state.pair.read().unwrap();
+        let map = state.pair.read().await;
         if !map.contains_key(&code) {
             return Html("<h3>Код не найден или истёк</h3>".to_string()).into_response();
         }
@@ -813,7 +877,7 @@ async fn pair_complete(State(state): State<AppState>, Form(f): Form<PairComplete
     let access_token = match oauth_access_token(&client, &email, &pass_md5).await {
         Ok(t) => t,
         Err(e) => {
-            let mut map = state.pair.write().unwrap();
+            let mut map = state.pair.write().await;
             if let Some(s) = map.get_mut(&code) {
                 s.error = Some(e.clone());
             }
@@ -830,7 +894,7 @@ async fn pair_complete(State(state): State<AppState>, Form(f): Form<PairComplete
     let me = match api_deezer_get(&client, "/user/me", &access_token, &[]).await {
         Ok(v) => v,
         Err(e) => {
-            let mut map = state.pair.write().unwrap();
+            let mut map = state.pair.write().await;
             if let Some(s) = map.get_mut(&code) {
                 s.error = Some(e.clone());
             }
@@ -876,7 +940,7 @@ async fn pair_complete(State(state): State<AppState>, Form(f): Form<PairComplete
     }
 
     {
-        let mut map = state.pair.write().unwrap();
+        let mut map = state.pair.write().await;
         if let Some(s) = map.get_mut(&code) {
             s.access_token = Some(access_token);
             s.user_id = uid;
@@ -905,7 +969,7 @@ async fn oauth_me(Json(req): Json<OauthMeReq>) -> impl IntoResponse {
         .unwrap();
     match api_deezer_get(&client, "/user/me", &token, &[]).await {
         Ok(v) => (StatusCode::OK, Json(v)).into_response(),
-        Err(e) => json_error(StatusCode::BAD_GATEWAY, e),
+        Err(e) => json_error(StatusCode::SERVICE_UNAVAILABLE, e),
     }
 }
 
@@ -929,17 +993,17 @@ async fn oauth_playlists(Json(req): Json<OauthPlaylistsReq>) -> impl IntoRespons
         .unwrap();
     let me = match api_deezer_get(&client, "/user/me", &token, &[]).await {
         Ok(v) => v,
-        Err(e) => return json_error(StatusCode::BAD_GATEWAY, e),
+        Err(e) => return json_error(StatusCode::SERVICE_UNAVAILABLE, e),
     };
     let uid = me.get("id").and_then(|x| x.as_i64()).unwrap_or(0);
     if uid <= 0 {
-        return json_error(StatusCode::BAD_GATEWAY, "api:no_user_id");
+        return json_error(StatusCode::SERVICE_UNAVAILABLE, "api:no_user_id");
     }
     let path = format!("/user/{uid}/playlists");
     let extra = [("index", index), ("limit", limit)];
     match api_deezer_get(&client, &path, &token, &extra).await {
         Ok(v) => (StatusCode::OK, Json(v)).into_response(),
-        Err(e) => json_error(StatusCode::BAD_GATEWAY, e),
+        Err(e) => json_error(StatusCode::SERVICE_UNAVAILABLE, e),
     }
 }
 
@@ -969,7 +1033,7 @@ async fn oauth_playlist_tracks(Json(req): Json<OauthPlaylistTracksReq>) -> impl
     let extra = [("index", index), ("limit", limit)];
     match api_deezer_get(&client, &path, &token, &extra).await {
         Ok(v) => (StatusCode::OK, Json(v)).into_response(),
-        Err(e) => json_error(StatusCode::BAD_GATEWAY, e),
+        Err(e) => json_error(StatusCode::SERVICE_UNAVAILABLE, e),
     }
 }
 
@@ -1001,21 +1065,21 @@ async fn oauth_arl(Json(req): Json<OauthArlReq>) -> impl IntoResponse {
 
     let ud = match gw_light_call(&client, "deezer.getUserData", "null", json!({})).await {
         Ok(v) => v,
-        Err(e) => return json_error(StatusCode::BAD_GATEWAY, e),
+        Err(e) => return json_error(StatusCode::SERVICE_UNAVAILABLE, e),
     };
 
     let check_form = ud["checkForm"].as_str().unwrap_or("").to_string();
     if check_form.is_empty() {
-        return json_error(StatusCode::BAD_GATEWAY, "no checkForm");
+        return json_error(StatusCode::SERVICE_UNAVAILABLE, "no checkForm");
     }
 
     let arl_v = match gw_light_call(&client, "user.getArl", &check_form, json!({})).await {
         Ok(v) => v,
-        Err(e) => return json_error(StatusCode::BAD_GATEWAY, e),
+        Err(e) => return json_error(StatusCode::SERVICE_UNAVAILABLE, e),
     };
     let arl = arl_v.as_str().unwrap_or("").to_string();
     if arl.len() < 20 {
-        return json_error(StatusCode::BAD_GATEWAY, "no arl");
+        return json_error(StatusCode::SERVICE_UNAVAILABLE, "no arl");
     }
 
     (StatusCode::OK, Json(json!({ "arl": arl }))).into_response()
@@ -1028,7 +1092,8 @@ async fn main() {
     let port: u16 = port.parse().unwrap_or(8000);
 
     let state = AppState {
-        api: Arc::new(RwLock::new(APIClient::new())),
+        api: Arc::new(Mutex::new(APIClient::new())),
+        api_by_arl: Arc::new(RwLock::new(HashMap::new())),
         pair: Arc::new(RwLock::new(HashMap::new())),
     };