Upload main.rs

src/main.rs

@@ -1,6 +1,7 @@
 use axum::{
     routing::{get, post},
     http::StatusCode,
+    http::{HeaderMap, header::{RANGE, CONTENT_RANGE, ACCEPT_RANGES, CONTENT_LENGTH}},
     response::{IntoResponse, Html},
     extract::{Json, State, Query, Form},
     Router,
@@ -14,6 +15,7 @@ use std::sync::Arc;
 use bytes::Bytes;
 use futures_util::StreamExt;
 use axum::body::Body;
+use axum::http::HeaderValue;
 use blowfish::Blowfish;
 use cipher::{KeyIvInit, BlockDecryptMut};
 use cipher::block_padding::NoPadding;
@@ -88,6 +90,55 @@ fn extract_media_url(media_json: &serde_json::Value) -> String {
     String::new()
 }
 
+fn extract_media_format(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 {
+        let has_url = m
+            .get("sources")
+            .and_then(|v| v.as_array())
+            .and_then(|a| a.get(0))
+            .and_then(|s| s.get("url"))
+            .and_then(|u| u.as_str())
+            .map(|u| !u.trim().is_empty())
+            .unwrap_or(false);
+        if !has_url {
+            continue;
+        }
+        if let Some(s) = m.get("format").and_then(|v| v.as_str()) {
+            let s = s.trim();
+            if !s.is_empty() {
+                return s.to_string();
+            }
+        }
+        if let Some(s) = m.get("format_name").and_then(|v| v.as_str()) {
+            let s = s.trim();
+            if !s.is_empty() {
+                return s.to_string();
+            }
+        }
+        if let Some(n) = m.get("format").and_then(|v| v.as_i64()) {
+            return n.to_string();
+        }
+    }
+    String::new()
+}
+
+fn total_from_content_range(v: &str) -> Option<u64> {
+    let v = v.trim();
+    let (_, rest) = v.split_once('/')?;
+    if rest.trim() == "*" {
+        return None;
+    }
+    rest.trim().parse::<u64>().ok()
+}
+
 async fn media_url_for_track(
     client: &mut APIClient,
     id: u32,
@@ -96,7 +147,7 @@ async fn media_url_for_track(
     let resp: Result<DeezerTrackList, APIError> = client
         .api_call(
             "song.getListData",
-            &json!({"sng_ids":[id],"array_default":["SNG_ID","TRACK_TOKEN"]}),
+            &json!({"sng_ids":[id],"array_default":["SNG_ID","TRACK_TOKEN","DURATION"]}),
         )
         .await;
 
@@ -172,6 +223,23 @@ async fn upstream_media_url(id: u32, formats: &Vec<Format>) -> Result<String, St
     Ok(url)
 }
 
+async fn public_track_duration(id: u32) -> Option<u32> {
+    let client = reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(10))
+        .build()
+        .ok()?;
+    let r = client
+        .get(format!("https://api.deezer.com/track/{id}"))
+        .send()
+        .await
+        .ok()?;
+    if !r.status().is_success() {
+        return None;
+    }
+    let v: serde_json::Value = r.json().await.ok()?;
+    v.get("duration").and_then(|d| d.as_u64()).map(|d| d as u32)
+}
+
 #[derive(Deserialize)]
 struct DeezerTrackList {
     data: Vec<DeezerTrack>
@@ -180,7 +248,8 @@ struct DeezerTrackList {
 #[derive(Deserialize)]
 #[allow(non_snake_case)]
 struct DeezerTrack {
-    TRACK_TOKEN: String
+    TRACK_TOKEN: String,
+    DURATION: Option<u32>
 }
 
 async fn root() -> &'static str {
@@ -316,7 +385,131 @@ struct StreamParams {
     arl: Option<String>,
 }
 
-async fn stream(State(state): State<AppState>, Query(q): Query<StreamParams>) -> impl IntoResponse {
+async fn stream_info(State(state): State<AppState>, Query(q): Query<StreamParams>) -> impl IntoResponse {
+    let formats = match q.format.as_deref().unwrap_or("AUTO") {
+        "FLAC" => vec![Format::FLAC, Format::MP3_320, Format::MP3_128],
+        "MP3_320" => vec![Format::MP3_320, Format::MP3_128],
+        "MP3_MISC" => vec![Format::MP3_MISC, Format::MP3_128],
+        "MP3_128" => vec![Format::MP3_128],
+        "AUTO" | _ => vec![Format::FLAC, Format::MP3_320, Format::MP3_128],
+    };
+
+    let requested = q.format.clone().unwrap_or_else(|| "AUTO".to_string());
+    let mut url: String = String::new();
+    let mut used: String = String::new();
+    let mut last_err: Option<String> = None;
+
+    if let Ok(text) = upstream_get_url_text(&formats, &vec![q.id]).await {
+        if let Ok(v) = serde_json::from_str::<serde_json::Value>(&text) {
+            url = extract_media_url(&v);
+            used = extract_media_format(&v);
+        }
+    }
+
+    if url.is_empty() {
+        if let Some(arl) = q.arl.clone().filter(|s| !s.trim().is_empty()) {
+            let client = api_client_for_arl(&state, Some(arl)).await;
+            let res = {
+                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","DURATION"]}),
+                    )
+                    .await;
+                let track_list = resp.map_err(|e| e.to_string())?;
+                if track_list.data.is_empty() {
+                    return Err("No valid ID".to_string());
+                }
+                let track_token = track_list.data[0].TRACK_TOKEN.as_str();
+                let media_resp = client
+                    .get_media(&formats, vec![track_token])
+                    .await
+                    .map_err(|e| e.to_string())?;
+                let media_json: serde_json::Value = media_resp
+                    .json()
+                    .await
+                    .map_err(|_| "Bad media response".to_string())?;
+                let url = extract_media_url(&media_json);
+                let used = extract_media_format(&media_json);
+                if url.is_empty() {
+                    return Err("No url in media response".to_string());
+                }
+                Ok::<(String, String), String>((url, used))
+            }
+            .await;
+            match res {
+                Ok((u, f)) => {
+                    url = u;
+                    used = f;
+                }
+                Err(e) => last_err = Some(e),
+            }
+        }
+    }
+
+    if url.is_empty() {
+        return json_error(
+            StatusCode::SERVICE_UNAVAILABLE,
+            last_err.unwrap_or_else(|| "stream_info:no_url".to_string()),
+        );
+    }
+
+    let duration = public_track_duration(q.id).await.unwrap_or(0);
+    let http = reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(15))
+        .build()
+        .unwrap();
+    let r = http
+        .get(url.clone())
+        .header(RANGE, "bytes=0-0")
+        .send()
+        .await;
+
+    let mut total_bytes: u64 = 0;
+    let mut mime: String = String::new();
+    if let Ok(r) = r {
+        if let Some(ct) = r.headers().get(CONTENT_TYPE).and_then(|v| v.to_str().ok()) {
+            mime = ct.to_string();
+        }
+        if let Some(cr) = r
+            .headers()
+            .get(CONTENT_RANGE)
+            .and_then(|v| v.to_str().ok())
+            .and_then(total_from_content_range)
+        {
+            total_bytes = cr;
+        } else if let Some(cl) = r.headers().get(CONTENT_LENGTH).and_then(|v| v.to_str().ok()) {
+            total_bytes = cl.parse::<u64>().unwrap_or(0);
+        }
+    }
+
+    let bitrate_kbps = if duration > 0 && total_bytes > 0 {
+        ((total_bytes as f64) * 8.0 / (duration as f64) / 1000.0).round() as u64
+    } else {
+        0
+    };
+
+    (
+        StatusCode::OK,
+        Json(json!({
+            "id": q.id,
+            "requested": requested,
+            "used": used,
+            "duration": duration,
+            "bytes": total_bytes,
+            "bitrate_kbps": bitrate_kbps,
+            "mime": mime
+        })),
+    )
+        .into_response()
+}
+
+async fn stream(
+    State(state): State<AppState>,
+    Query(q): Query<StreamParams>,
+    headers: HeaderMap,
+) -> impl IntoResponse {
     let formats = match q.format.as_deref().unwrap_or("AUTO") {
         "FLAC" => vec![Format::FLAC, Format::MP3_320, Format::MP3_128],
         "MP3_320" => vec![Format::MP3_320, Format::MP3_128],
@@ -394,12 +587,65 @@ async fn stream(State(state): State<AppState>, Query(q): Query<StreamParams>) ->
         }
     };
 
+    fn parse_range_header(v: &str) -> Option<(u64, Option<u64>)> {
+        let v = v.trim();
+        let v = v.strip_prefix("bytes=")?;
+        let mut it = v.splitn(2, '-');
+        let a = it.next()?.trim();
+        let b = it.next().unwrap_or("").trim();
+        if a.is_empty() {
+            return None;
+        }
+        let start = a.parse::<u64>().ok()?;
+        let end = if b.is_empty() {
+            None
+        } else {
+            let e = b.parse::<u64>().ok()?;
+            Some(e)
+        };
+        Some((start, end))
+    }
+
+    fn parse_total_from_content_range(v: &str) -> Option<u64> {
+        let v = v.trim();
+        let (_, rest) = v.split_once('/')?;
+        if rest.trim() == "*" {
+            return None;
+        }
+        rest.trim().parse::<u64>().ok()
+    }
+
+    let req_range = headers
+        .get(RANGE)
+        .and_then(|v| v.to_str().ok())
+        .and_then(parse_range_header);
+
     let http = reqwest::Client::builder()
         .timeout(std::time::Duration::from_secs(60))
         .build()
         .unwrap();
 
-    let upstream = match http.get(url).send().await {
+    let block_size: u64 = 2048;
+    let mut block_index: u64 = 0;
+    let mut drop_bytes: usize = 0;
+    let mut remaining: Option<u64> = None;
+
+    let mut req = http.get(url);
+    if let Some((start, end)) = req_range {
+        let aligned_start = start - (start % block_size);
+        drop_bytes = (start - aligned_start) as usize;
+        block_index = aligned_start / block_size;
+        if let Some(e) = end {
+            if e >= start {
+                remaining = Some(e - start + 1);
+            }
+            req = req.header(RANGE, format!("bytes={}-{}", aligned_start, e));
+        } else {
+            req = req.header(RANGE, format!("bytes={}-", aligned_start));
+        }
+    }
+
+    let upstream = match req.send().await {
         Ok(r) => r,
         Err(_) => {
             return (StatusCode::SERVICE_UNAVAILABLE, "Upstream request failed".to_string())
@@ -417,11 +663,18 @@ async fn stream(State(state): State<AppState>, Query(q): Query<StreamParams>) ->
 
     let key = blowfish_key(q.id);
     let mut carry: Vec<u8> = Vec::with_capacity(4096);
-    let mut block_index: u64 = 0;
+    let upstream_headers = upstream.headers().clone();
+    let upstream_total = upstream_headers
+        .get(CONTENT_RANGE)
+        .and_then(|v| v.to_str().ok())
+        .and_then(parse_total_from_content_range);
 
     let mut upstream = upstream.bytes_stream();
 
     let out = async_stream::stream! {
+        let mut drop_bytes = drop_bytes;
+        let mut remaining = remaining;
+        let mut done = false;
         while let Some(chunk_res) = upstream.next().await {
             let chunk = match chunk_res {
                 Ok(c) => c,
@@ -439,17 +692,98 @@ async fn stream(State(state): State<AppState>, Query(q): Query<StreamParams>) ->
                     let _ = decrypt_stripe(&mut block, &key);
                 }
                 block_index += 1;
+                if drop_bytes > 0 {
+                    if drop_bytes >= block.len() {
+                        drop_bytes -= block.len();
+                        continue;
+                    }
+                    block = block.split_off(drop_bytes);
+                    drop_bytes = 0;
+                }
+                if let Some(rem) = remaining {
+                    if rem == 0 {
+                        done = true;
+                        break;
+                    }
+                    if (block.len() as u64) > rem {
+                        block.truncate(rem as usize);
+                        remaining = Some(0);
+                        yield Ok::<Bytes, BoxErr>(Bytes::from(block));
+                        done = true;
+                        break;
+                    } else {
+                        remaining = Some(rem - block.len() as u64);
+                    }
+                }
                 yield Ok::<Bytes, BoxErr>(Bytes::from(block));
             }
+            if done { break; }
         }
 
-        if !carry.is_empty() {
-            yield Ok::<Bytes, BoxErr>(Bytes::from(std::mem::take(&mut carry)));
+        if !done && !carry.is_empty() {
+            let mut tail = std::mem::take(&mut carry);
+            if drop_bytes > 0 {
+                if drop_bytes < tail.len() {
+                    tail = tail.split_off(drop_bytes);
+                    drop_bytes = 0;
+                } else {
+                    tail.clear();
+                }
+            }
+            if !tail.is_empty() {
+                if let Some(rem) = remaining {
+                    if rem > 0 {
+                        if (tail.len() as u64) > rem {
+                            tail.truncate(rem as usize);
+                        }
+                        yield Ok::<Bytes, BoxErr>(Bytes::from(tail));
+                    }
+                } else {
+                    yield Ok::<Bytes, BoxErr>(Bytes::from(tail));
+                }
+            }
         }
     };
 
     let body = Body::from_stream(out);
-    (StatusCode::OK, [(CONTENT_TYPE, "audio/mpeg")], body).into_response()
+    let mut resp = axum::response::Response::new(body);
+    resp.headers_mut()
+        .insert(CONTENT_TYPE, HeaderValue::from_static("audio/mpeg"));
+    resp.headers_mut()
+        .insert(ACCEPT_RANGES, HeaderValue::from_static("bytes"));
+
+    if let Some((start, end)) = req_range {
+        if let Some(total) = upstream_total {
+            let end_for_hdr = end.unwrap_or_else(|| total.saturating_sub(1));
+            if end_for_hdr >= start {
+                let _ = resp.headers_mut().insert(
+                    CONTENT_RANGE,
+                    HeaderValue::from_str(&format!("bytes {}-{}/{}", start, end_for_hdr, total))
+                        .unwrap_or_else(|_| HeaderValue::from_static("bytes 0-0/*")),
+                );
+                let len = end_for_hdr.saturating_sub(start).saturating_add(1);
+                let _ = resp
+                    .headers_mut()
+                    .insert(CONTENT_LENGTH, HeaderValue::from_str(&len.to_string()).unwrap());
+                *resp.status_mut() = StatusCode::PARTIAL_CONTENT;
+            }
+        } else if let Some(end_for_hdr) = end {
+            if end_for_hdr >= start {
+                let _ = resp.headers_mut().insert(
+                    CONTENT_RANGE,
+                    HeaderValue::from_str(&format!("bytes {}-{}/{}", start, end_for_hdr, "*"))
+                        .unwrap_or_else(|_| HeaderValue::from_static("bytes 0-0/*")),
+                );
+                let len = end_for_hdr.saturating_sub(start).saturating_add(1);
+                let _ = resp
+                    .headers_mut()
+                    .insert(CONTENT_LENGTH, HeaderValue::from_str(&len.to_string()).unwrap());
+                *resp.status_mut() = StatusCode::PARTIAL_CONTENT;
+            }
+        }
+    }
+
+    resp
 }
 
 #[derive(Debug, Deserialize)]
@@ -1354,6 +1688,7 @@ async fn main() {
         .route("/get_url", post(get_url))
         .route("/fetch", get(fetch))
         .route("/stream", get(stream))
+        .route("/stream_info", get(stream_info))
         .route("/user_data", post(user_data))
         .route("/playlists", post(playlists))
         .route("/playlist", post(playlist_tracks))