From d486f04905cc86c4e576af9d90c4e46c8e655824 Mon Sep 17 00:00:00 2001 From: euzu Date: Tue, 7 Jan 2025 14:00:13 +0100 Subject: [PATCH] reconnect stream wip --- src/api/api_utils.rs | 127 ++++++------------------------- src/api/model/buffered_stream.rs | 87 +++++++++++++++++++++ src/api/model/mod.rs | 3 +- src/api/model/shared_stream.rs | 61 ++++++++++++++- test/rest-api.http | 2 +- 5 files changed, 174 insertions(+), 106 deletions(-) create mode 100644 src/api/model/buffered_stream.rs diff --git a/src/api/api_utils.rs b/src/api/api_utils.rs index c7d944410..dc9db0b47 100644 --- a/src/api/api_utils.rs +++ b/src/api/api_utils.rs @@ -1,27 +1,24 @@ use crate::api::model::app_state::AppState; use crate::api::model::request::UserApiRequest; use crate::api::model::shared_stream::SharedStream; -use crate::model::api_proxy::{ProxyUserCredentials}; +use crate::debug_if_enabled; +use crate::model::api_proxy::ProxyUserCredentials; use crate::model::config::{ConfigInput, ConfigTarget}; +use crate::model::playlist::PlaylistItemType; use crate::utils::request_utils; use crate::utils::request_utils::mask_sensitive_info; +use actix_web::http::header::DATE; use actix_web::http::header::{HeaderValue, CACHE_CONTROL}; -use actix_web::http::header::{DATE}; use actix_web::{HttpRequest, HttpResponse}; -use async_std::prelude::Stream; -use async_std::stream::StreamExt; use bytes::Bytes; use chrono::Utc; -use log::{debug, error}; +use log::{ error}; use std::collections::HashMap; use std::path::Path; -use std::sync::Arc; -use std::time::Duration; -use tokio::sync::broadcast; -use tokio_stream::wrappers::BroadcastStream; +use tokio_stream::wrappers::{BroadcastStream}; use url::Url; -use crate::debug_if_enabled; -use crate::model::playlist::PlaylistItemType; +use crate::api::model::buffered_stream; +use crate::api::model::buffered_stream::get_stream_response_with_headers; pub async fn serve_file(file_path: &Path, req: &HttpRequest, mime_type: mime::Mime) -> HttpResponse { if file_path.exists() { @@ -73,61 +70,8 @@ async fn create_notify_stream( } } -/// Creates a shared stream and stores it in the shared state. -async fn create_shared_stream( - app_state: &AppState, - stream_url: &str, - header: HashMap>, - bytes_stream: S, -) where - S: Stream> + Unpin + 'static, -{ - // Create a broadcast channel for the shared stream - let (tx, _) = broadcast::channel(100); - let sender = Arc::new(tx); - - // Insert the shared stream into the shared state - app_state - .shared_streams - .lock() - .await - .insert( - stream_url.to_string(), - SharedStream { - data_stream: sender.clone(), - header, - }, - ); - - let shared_streams_map = Arc::clone(&app_state.shared_streams); - let mut source_stream = Box::pin(bytes_stream); - let streaming_url = stream_url.to_string(); - // Spawn a task to forward items from the source stream to the broadcast channel - actix_rt::spawn(async move { - while let Some(item) = source_stream.next().await { - if let Ok(data) = item { - if sender.receiver_count() > 0 { - // if let Err(err) = sender.send(data) { - // debug!("{err}") - // } - if sender.send(data).is_err() { - // ignore - } - actix_web::rt::time::sleep(Duration::from_millis(20)).await; - } else { - debug_if_enabled!("No active subscribers. Closing stream {}", mask_sensitive_info(&streaming_url)); - // Cleanup for removing unused shared streams - let mut shared_streams = shared_streams_map.lock().await; - shared_streams.remove(&streaming_url); - return; - } - } - } - }); -} pub async fn stream_response(app_state: &AppState, stream_url: &str, req: &HttpRequest, input: Option<&ConfigInput>, share_stream: bool) -> HttpResponse { - let req_headers: HashMap<&str, &[u8]> = req.headers().iter().map(|(k, v)| (k.as_str(), v.as_bytes())).collect(); debug_if_enabled!("Try to open stream {}", mask_sensitive_info(stream_url)); if share_stream { if let Some(value) = shared_stream_response(app_state, stream_url).await { @@ -136,55 +80,34 @@ pub async fn stream_response(app_state: &AppState, stream_url: &str, req: &HttpR } if let Ok(url) = Url::parse(stream_url) { - let client = request_utils::get_client_request(input, &url, Some(&req_headers)); - match client.send().await { - Ok(response) => { - let status = response.status(); - if status.is_success() { - let mut response_builder = HttpResponse::Ok(); - let mut header = HashMap::new(); - response.headers().iter().for_each(|(k, v)| { - let key = k.as_str().to_string(); - // ignore date it is dynamic - if !"date".eq(key.to_lowercase().as_str()) { - header.insert(k.as_str().to_string(), v.as_bytes().to_vec()); - } - response_builder.insert_header((k.as_str(), v.as_ref())); - }); - if share_stream { - create_shared_stream(app_state, stream_url, header, response.bytes_stream()).await; - if let Some(stream) = create_notify_stream(app_state, stream_url).await { - debug!("Creating shared channel {stream_url}"); - return response_builder.body(actix_web::body::BodyStream::new(stream)); - } - } else { - return response_builder.body(actix_web::body::BodyStream::new(response.bytes_stream())); - } - } - debug_if_enabled!("Failed to open stream got status {} for {}", status, mask_sensitive_info(stream_url)); - } - Err(err) => { - error!("Received failure from server {}: {}", mask_sensitive_info(stream_url), err); + let mut buffered_stream_handler = buffered_stream::BufferedStreamHandler::new(&url, req, input); + let stream = buffered_stream_handler.get_stream(); + return if share_stream { + SharedStream::register(app_state, stream_url, stream).await; + if let Some(broadcast_stream) = create_notify_stream(app_state, stream_url).await { + let body_stream = actix_web::body::BodyStream::new(broadcast_stream); + let mut response_builder = get_stream_response_with_headers(); + response_builder.body(body_stream) + } else { + HttpResponse::BadRequest().finish() } + } else { + let mut response_builder = get_stream_response_with_headers(); + let body_stream = actix_web::body::BodyStream::new(stream); + response_builder.body(body_stream) } - } else { - error!("Url is malformed {}", mask_sensitive_info(stream_url)); } + error!("Url is malformed {}", mask_sensitive_info(stream_url)); HttpResponse::BadRequest().finish() } async fn shared_stream_response(app_state: &AppState, stream_url: &str) -> Option { if let Some(stream) = create_notify_stream(app_state, stream_url).await { debug_if_enabled!("Using shared channel {}", mask_sensitive_info(stream_url)); - // return HttpResponse::Ok().body(actix_web::body::BodyStream::new(stream)); - if let Some(shared_stream) = app_state.shared_streams.lock().await.get(stream_url) { - let mut response_builder = HttpResponse::Ok(); - for (key, value) in &shared_stream.header { - response_builder.insert_header((key.as_str(), &value[..])); - } + if app_state.shared_streams.lock().await.get(stream_url).is_some() { + let mut response_builder = get_stream_response_with_headers(); let current_date = Utc::now().format("%a, %d %b %Y %H:%M:%S GMT").to_string(); response_builder.insert_header((DATE, current_date.as_bytes())); - response_builder.insert_header((CACHE_CONTROL, "no-cache".as_bytes())); // response_builder.insert_header((ACCEPT_RANGES, "bytes".as_bytes())); return Some(response_builder.body(actix_web::body::BodyStream::new(stream))); } diff --git a/src/api/model/buffered_stream.rs b/src/api/model/buffered_stream.rs new file mode 100644 index 000000000..1f77d1c08 --- /dev/null +++ b/src/api/model/buffered_stream.rs @@ -0,0 +1,87 @@ +use crate::model::config::ConfigInput; +use crate::utils::request_utils; +use actix_web::{HttpRequest, HttpResponse, HttpResponseBuilder}; +use async_std::stream::StreamExt; +use bytes::Bytes; +use core::time::Duration; +use reqwest::{Error, RequestBuilder}; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::mpsc; +use tokio_stream::Stream; +use tokio_stream::wrappers::ReceiverStream; +use url::Url; + +const BUFFER_SIZE: usize = 8092; + +pub struct BufferedStreamHandler { + client: Arc, +} + +impl BufferedStreamHandler { + pub fn new(url: &Url, req: &HttpRequest, input: Option<&ConfigInput>) -> Self { + let req_headers: HashMap<&str, &[u8]> = req.headers().iter().map(|(k, v)| (k.as_str(), v.as_bytes())).collect(); + let client = Arc::new(request_utils::get_client_request(input, url, Some(&req_headers))); + + BufferedStreamHandler { + client, + } + } + + pub fn get_stream(&mut self) -> impl Stream> + Unpin + 'static { + let (tx, rx) = mpsc::channel::>(BUFFER_SIZE); + let client = Arc::clone(&self.client); + actix_web::rt::spawn({ + async move { + loop { + let Some(client) = client.try_clone() else { break }; + println!("Connection to stream"); + match client.send().await { + Ok(response) => { + if !response.status().is_success() { + continue; + } + let mut byte_stream = response.bytes_stream(); + while let Some(chunk) = byte_stream.next().await { + match chunk { + Ok(chunk) => { + if tx.send(Ok(chunk)).await.is_err() { + println!("Stream finished, client disconnect ?"); + return; + } + } + Err(err) => { + println!("Stream disconnected, cant read from server {err}"); + break; + } + } + } + } + Err(err) => { + if err.is_timeout() { + actix_web::rt::time::sleep(Duration::from_secs(1)).await; + continue; + } + println!("Stream finished {err}"); + break; + } + } + actix_web::rt::time::sleep(Duration::from_secs(1)).await; + } + } + }); + + ReceiverStream::new(rx) + } +} + + +pub fn get_stream_response_with_headers() -> HttpResponseBuilder { + let mut response_builder = HttpResponse::Ok(); + response_builder.insert_header((actix_web::http::header::CONTENT_TYPE, "application/octet-stream")); + response_builder.insert_header((actix_web::http::header::CONTENT_LENGTH, 0)); + response_builder.insert_header((actix_web::http::header::CONNECTION, "keep-alive")); + response_builder.insert_header((actix_web::http::header::CACHE_CONTROL, "no-cache")); + + response_builder +} \ No newline at end of file diff --git a/src/api/model/mod.rs b/src/api/model/mod.rs index 615449ccf..c96d911b2 100644 --- a/src/api/model/mod.rs +++ b/src/api/model/mod.rs @@ -3,4 +3,5 @@ pub mod config; pub mod download; pub mod xtream; pub mod app_state; -pub mod shared_stream; \ No newline at end of file +pub mod shared_stream; +pub mod buffered_stream; diff --git a/src/api/model/shared_stream.rs b/src/api/model/shared_stream.rs index 9ba311785..80f7156fa 100644 --- a/src/api/model/shared_stream.rs +++ b/src/api/model/shared_stream.rs @@ -1,9 +1,66 @@ use bytes::Bytes; -use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; +use async_std::prelude::{Stream, StreamExt}; +use tokio::sync::broadcast; +use crate::api::model::app_state::AppState; +use crate::debug_if_enabled; +use crate::utils::request_utils::mask_sensitive_info; pub struct SharedStream { pub data_stream: Arc>, - pub header: HashMap>, } +impl SharedStream { + + /// Creates a shared stream and stores it in the shared state. + pub(crate) async fn register( + app_state: &AppState, + stream_url: &str, + bytes_stream: S, + ) where + S: Stream> + Unpin + 'static, + { + // Create a broadcast channel for the shared stream + let (tx, _) = broadcast::channel(100); + let sender = Arc::new(tx); + + // Insert the shared stream into the shared state + app_state + .shared_streams + .lock() + .await + .insert( + stream_url.to_string(), + SharedStream { + data_stream: sender.clone(), + }, + ); + + let shared_streams_map = Arc::clone(&app_state.shared_streams); + let mut source_stream = Box::pin(bytes_stream); + let streaming_url = stream_url.to_string(); + // Spawn a task to forward items from the source stream to the broadcast channel + actix_rt::spawn(async move { + while let Some(item) = source_stream.next().await { + if let Ok(data) = item { + if sender.receiver_count() > 0 { + // if let Err(err) = sender.send(data) { + // debug!("{err}") + // } + if sender.send(data).is_err() { + // ignore + } + actix_web::rt::time::sleep(Duration::from_millis(20)).await; + } else { + debug_if_enabled!("No active subscribers. Closing stream {}", mask_sensitive_info(&streaming_url)); + // Cleanup for removing unused shared streams + let mut shared_streams = shared_streams_map.lock().await; + shared_streams.remove(&streaming_url); + return; + } + } + } + }); + } +} diff --git a/test/rest-api.http b/test/rest-api.http index eb7690a1a..111fce0e3 100644 --- a/test/rest-api.http +++ b/test/rest-api.http @@ -27,7 +27,7 @@ Content-Type: application/x-www-form-urlencoded username=xt&password=xt ### m3u playlist reverse -GET {{local}}/m3u?username=m3u&password=m3u +GET {{local}}/m3u?username=xtr&password=xtr ### m3u playlist