add: playlist get

Author: qingyiwebt

swiftstream/src/caching/stream_tracking.rs

@@ -11,7 +11,7 @@ use reqwest::Client;
 use tokio::{sync::RwLock, time::sleep};
 use url::Url;
 
-use crate::caching::CachePool;
+use crate::{caching::CachePool, transfer::parse_m3u8_async};
 
 pub struct StreamTrackingPool {
     tracking: RwLock<HashMap<String, Arc<TrackingItem>>>,
@@ -137,18 +137,16 @@ impl TrackingItem {
         &self,
         tracking_pool: &Arc<StreamTrackingPool>,
     ) -> Result<(), anyhow::Error> {
-        let response = tracking_pool.http_client.get(&self.origin).send().await?;
-        let data = response.bytes().await?;
+        let data = tracking_pool
+            .http_client
+            .get(&self.origin)
+            .send()
+            .await?
+            .bytes()
+            .await?;
 
         // parse
-        let playlist = tokio::task::spawn_blocking(move || {
-            let mut parser = mediastream_rs::Parser::new(Cursor::new(data));
-            if let Err(e) = parser.parse() {
-                return Err(e);
-            }
-            Ok(parser.get_result())
-        })
-        .await??;
+        let playlist = parse_m3u8_async(Cursor::new(data)).await?;
 
         self.prepare_all(tracking_pool, &self.origin, playlist)
             .await?;

swiftstream/src/config.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, fs::File, path::Path};
+use std::{fs::File, path::Path};
 
 use anyhow::Result;
 use serde::Deserialize;
@@ -12,13 +12,6 @@ pub struct Config {
     pub cache_expire: Option<u16>,
     pub track_expire: Option<u16>,
     pub track_interval: Option<u16>,
-    pub upstreams: HashMap<String, UpstreamConfig>,
-}
-
-#[derive(Debug, Deserialize)]
-#[serde(rename_all = "camelCase")]
-pub struct UpstreamConfig {
-    pub url: String,
 }
 
 pub fn load_config(path: impl AsRef<Path>) -> Result<Config> {

swiftstream/src/routes/media.rs

@@ -9,10 +9,10 @@ use mediastream_rs::format::M3uPlaylist;
 use serde::Deserialize;
 use url::Url;
 
-use crate::{AppStateRef, internal_error_with_log};
+use crate::{AppStateRef, internal_error_with_log, transfer::parse_m3u8_async};
 
 #[derive(Deserialize)]
-pub struct MediaRequest {
+pub struct MediaQuery {
     origin: String,
 }
 
@@ -46,28 +46,24 @@ async fn prepare_all(
 
 pub async fn get_media(
     State(state): State<AppStateRef>,
-    Query(query): Query<MediaRequest>,
+    Query(query): Query<MediaQuery>,
 ) -> Result<Response, StatusCode> {
-    let response = state
+    let data = state
         .http_client
         .get(&query.origin)
         .send()
         .await
+        .map_err(internal_error_with_log!())?
+        .bytes()
+        .await
         .map_err(internal_error_with_log!())?;
-    let data = response.bytes().await.map_err(internal_error_with_log!())?;
+
     state.tracking_pool.track(&query.origin).await;
 
     // parse
-    let mut playlist = tokio::task::spawn_blocking(move || {
-        let mut parser = mediastream_rs::Parser::new(Cursor::new(data));
-        if let Err(e) = parser.parse() {
-            return Err(e);
-        }
-        Ok(parser.get_result())
-    })
-    .await
-    .map_err(internal_error_with_log!())?
-    .map_err(internal_error_with_log!())?;
+    let mut playlist = parse_m3u8_async(Cursor::new(data))
+        .await
+        .map_err(internal_error_with_log!())?;
 
     prepare_all(&state, &mut playlist, query.origin)
         .await

swiftstream/src/routes/mod.rs

@@ -3,10 +3,12 @@ use axum::{Router, routing::get};
 use crate::AppStateRef;
 
 mod media;
+mod playlist;
 mod stream;
 
 pub fn get_routes(app_state: &AppStateRef) -> Router {
     Router::new()
+        .route("/playlist", get(playlist::get_playlist))
         .route("/stream", get(stream::get_stream))
         .route("/media", get(media::get_media))
         .with_state(app_state.clone())

swiftstream/src/routes/playlist.rs

@@ -0,0 +1,54 @@
+use std::io::Cursor;
+
+use axum::{
+    extract::{Query, State},
+    response::{IntoResponse, Response},
+};
+use reqwest::StatusCode;
+use serde::Deserialize;
+use url::Url;
+
+use crate::{AppStateRef, internal_error_with_log, transfer::parse_m3u8_async};
+
+#[derive(Deserialize)]
+pub struct PlaylistQuery {
+    pub origin: String,
+}
+pub async fn get_playlist(
+    State(state): State<AppStateRef>,
+    Query(query): Query<PlaylistQuery>,
+) -> Result<Response, StatusCode> {
+    let data = state
+        .http_client
+        .get(&query.origin)
+        .send()
+        .await
+        .map_err(internal_error_with_log!())?
+        .bytes()
+        .await
+        .map_err(internal_error_with_log!())?;
+
+    let mut playlist = parse_m3u8_async(Cursor::new(data))
+        .await
+        .map_err(internal_error_with_log!())?;
+
+    let base = Url::parse(&query.origin).map_err(internal_error_with_log!())?;
+
+    for media in playlist.medias.iter_mut() {
+        let media_location = media.location.clone();
+        let mut location = Url::parse(&media_location);
+        if location == Err(url::ParseError::RelativeUrlWithoutBase) {
+            location = base.join(&media_location);
+        }
+        let location = location.map_err(internal_error_with_log!())?.to_string();
+
+        media.location = format!(
+            "{}/media?origin={}",
+            state.config.base_url,
+            urlencoding::encode(&location)
+        )
+        .into();
+    }
+
+    Ok(playlist.to_string().into_response())
+}

swiftstream/src/routes/stream.rs

@@ -18,13 +18,13 @@ use crate::{
 };
 
 #[derive(Deserialize)]
-pub struct StreamRequest {
+pub struct StreamQuery {
     pub origin: String,
 }
 
 pub async fn get_stream(
     State(state): State<AppStateRef>,
-    Query(query): Query<StreamRequest>,
+    Query(query): Query<StreamQuery>,
     headers: HeaderMap,
 ) -> Result<Response, StatusCode> {
     let data = state.cache_pool.get(&query.origin).await;
@@ -56,10 +56,10 @@ pub async fn get_stream(
                 .into_iter()
                 .map(|x| (x, Cursor::new(data.bytes.clone())))
                 .map(|mut x| match x.0 {
-                    HttpRange::Prefix(len) => x.1.take(len),
-                    HttpRange::Suffix(len) => {
-                        _ = x.1.seek(SeekFrom::End(-(len).try_into().unwrap_or(0)));
-                        x.1.take(len)
+                    HttpRange::Suffix(len) => x.1.take(len),
+                    HttpRange::Prefix(len) => {
+                        _ = x.1.seek(SeekFrom::Start(len));
+                        x.1.take(u64::MAX)
                     }
                     HttpRange::Range(from, to) => {
                         _ = x.1.seek(SeekFrom::Start(from));

swiftstream/src/transfer/mediastream_parse.rs

@@ -0,0 +1,50 @@
+use std::{
+    fmt::Display,
+    io::{BufRead, Seek},
+};
+
+use mediastream_rs::{ParseError, format::M3uPlaylist};
+use std::error::Error;
+use tokio::task::JoinError;
+
+#[derive(Debug)]
+pub enum ParseM3U8Error {
+    ParseError(ParseError),
+    JoinError(JoinError),
+}
+
+impl Display for ParseM3U8Error {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::JoinError(e) => e.fmt(f),
+            Self::ParseError(e) => e.fmt(f),
+        }
+    }
+}
+
+impl Error for ParseM3U8Error {}
+
+impl From<JoinError> for ParseM3U8Error {
+    fn from(value: JoinError) -> Self {
+        Self::JoinError(value)
+    }
+}
+
+impl From<ParseError> for ParseM3U8Error {
+    fn from(value: ParseError) -> Self {
+        Self::ParseError(value)
+    }
+}
+
+pub async fn parse_m3u8_async(
+    stream: impl BufRead + Seek + Send + 'static,
+) -> Result<M3uPlaylist, ParseM3U8Error> {
+    Ok(tokio::task::spawn_blocking(move || {
+        let mut parser = mediastream_rs::Parser::new(stream);
+        if let Err(e) = parser.parse() {
+            return Err(e);
+        }
+        Ok(parser.get_result())
+    })
+    .await??)
+}

swiftstream/src/transfer/mod.rs

@@ -1,3 +1,5 @@
+mod mediastream_parse;
 mod range;
 
+pub use mediastream_parse::*;
 pub use range::*;