haki_dl/
source.rs

1//! Source and segment loading utilities.
2
3use std::collections::BTreeMap;
4use std::path::{Path, PathBuf};
5
6use crate::config::DownloadOptions;
7use crate::error::{Error, Result};
8use crate::http::{DefaultHttpClient, HttpRequest, HttpResponse};
9use crate::manifest::ExtractorType;
10use crate::processor::{
11    ContentProcessor, DefaultDashContentProcessor, DefaultHlsContentProcessor, ParserConfig,
12};
13
14pub(crate) const HTTP_LIVE_TS_MARKER: &str = "<HAKI_LIVE_TS>";
15pub(crate) const BINARY_DATA_MARKER: &str = "<HAKI_BINARY_DATA>";
16
17/// Loaded source classification.
18#[derive(Clone, Copy, Debug, Eq, PartialEq)]
19pub enum LoadedSourceKind {
20    /// HLS text manifest.
21    Hls,
22    /// MPEG-DASH text manifest.
23    Dash,
24    /// Smooth Streaming manifest.
25    Mss,
26    /// Direct MPEG-TS live stream.
27    HttpLiveTs,
28    /// Direct binary input that is not a supported streaming manifest.
29    BinaryData,
30}
31
32impl LoadedSourceKind {
33    /// Returns the extractor family for this source.
34    pub fn extractor_type(self) -> ExtractorType {
35        match self {
36            Self::Hls => ExtractorType::Hls,
37            Self::Dash => ExtractorType::MpegDash,
38            Self::Mss => ExtractorType::Mss,
39            Self::HttpLiveTs => ExtractorType::HttpLive,
40            Self::BinaryData => ExtractorType::HttpLive,
41        }
42    }
43
44    fn raw_file_name(self) -> &'static str {
45        match self {
46            Self::Hls => "raw.m3u8",
47            Self::Dash => "raw.mpd",
48            Self::Mss => "raw.ism",
49            Self::HttpLiveTs => "raw.txt",
50            Self::BinaryData => "raw.bin",
51        }
52    }
53}
54
55/// Loaded source text and retained raw file map.
56#[derive(Clone, Debug, Eq, PartialEq)]
57pub struct LoadedSource {
58    /// Source kind.
59    pub kind: LoadedSourceKind,
60    /// Source text or live-TS marker.
61    pub text: String,
62    /// Original input URL/path.
63    pub original_url: String,
64    /// Final URL/path after redirects.
65    pub final_url: String,
66    /// Raw files that should be written to the temp directory later.
67    pub raw_files: BTreeMap<String, String>,
68    /// Debug lines captured while loading the source.
69    pub debug_logs: Vec<String>,
70}
71
72/// Compatibility source loader.
73pub struct SourceLoader {
74    http: DefaultHttpClient,
75    content_processors: Vec<Box<dyn ContentProcessor>>,
76}
77
78impl Default for SourceLoader {
79    fn default() -> Self {
80        Self {
81            http: DefaultHttpClient::new(),
82            content_processors: vec![
83                Box::<DefaultHlsContentProcessor>::default(),
84                Box::<DefaultDashContentProcessor>::default(),
85            ],
86        }
87    }
88}
89
90impl SourceLoader {
91    /// Creates a loader with default processors.
92    pub fn new() -> Self {
93        Self::default()
94    }
95
96    /// Creates a loader from public download options.
97    pub fn from_options(options: &DownloadOptions) -> Self {
98        Self::new().with_http(DefaultHttpClient::from_options(options))
99    }
100
101    /// Creates a loader with a custom HTTP client.
102    pub fn with_http(mut self, http: DefaultHttpClient) -> Self {
103        self.http = http;
104        self
105    }
106
107    /// Loads a manifest, local source, or direct live-TS source.
108    pub async fn load_source(
109        &self,
110        input: &str,
111        config: &mut ParserConfig,
112    ) -> Result<LoadedSource> {
113        let (mut text, original_url, final_url, mut debug_logs) = if input.starts_with("file:") {
114            let path = file_uri_to_path(input);
115            (
116                read_source_text_file(&path).await?,
117                input.to_string(),
118                input.to_string(),
119                Vec::new(),
120            )
121        } else if input.starts_with("http://") || input.starts_with("https://") {
122            let response = self.get_web_source(input, &config.headers).await?;
123            (response.0, input.to_string(), response.1, response.2)
124        } else if tokio::fs::metadata(input).await.is_ok() {
125            let path = tokio::fs::canonicalize(input).await?;
126            let file_uri =
127                path_to_file_uri(&path).unwrap_or_else(|| path.to_string_lossy().to_string());
128            (
129                read_source_text_file(&path).await?,
130                file_uri.clone(),
131                file_uri,
132                Vec::new(),
133            )
134        } else {
135            return Err(Error::http("source input could not be loaded"));
136        };
137
138        if text.trim().is_empty() {
139            return Err(Error::http("source input was empty"));
140        }
141        text = text.trim().trim_start_matches('\u{feff}').to_string();
142        config.original_url = original_url.clone();
143        config.url = final_url.clone();
144        let kind = detect_source_kind(&text)?;
145        let mut raw_files = BTreeMap::new();
146        raw_files.insert(kind.raw_file_name().to_string(), text.clone());
147        for processor in &self.content_processors {
148            if processor.can_process(kind.extractor_type(), &text, config) {
149                text = processor.process(&text, config)?;
150            }
151        }
152        debug_logs.retain(|line| !line.trim().is_empty());
153        Ok(LoadedSource {
154            kind,
155            text,
156            original_url,
157            final_url,
158            raw_files,
159            debug_logs,
160        })
161    }
162
163    /// Loads one segment-like byte source.
164    pub async fn load_segment_bytes(&self, url: &str, config: &ParserConfig) -> Result<Vec<u8>> {
165        if url.starts_with("file:") {
166            return Ok(tokio::fs::read(file_uri_to_path(url)).await?);
167        }
168        if let Some(value) = url.strip_prefix("base64://") {
169            return base64_decode(value);
170        }
171        if let Some(value) = url.strip_prefix("hex://") {
172            return hex_to_bytes(value);
173        }
174        if tokio::fs::metadata(url).await.is_ok() {
175            return Ok(tokio::fs::read(url).await?);
176        }
177        let mut request = HttpRequest::new(url);
178        request.headers = config.headers.clone();
179        Ok(self.http.send(request).await?.body)
180    }
181
182    async fn get_web_source(
183        &self,
184        url: &str,
185        headers: &BTreeMap<String, String>,
186    ) -> Result<(String, String, Vec<String>)> {
187        let mut request = HttpRequest::new(url);
188        request.headers = headers.clone();
189        let response = self.http.send_source(request).await?;
190        let mut debug_logs = response.debug_logs.clone();
191        if is_mpeg2_ts_buffer(&response.body) {
192            debug_logs.push("Detected MPEG-TS stream".to_string());
193            return Ok((HTTP_LIVE_TS_MARKER.to_string(), url.to_string(), debug_logs));
194        }
195        if let Some(encoding) = response.headers.get("content-encoding")
196            && !encoding.trim().is_empty()
197        {
198            debug_logs.push(format!("Detected compression: {encoding}"));
199        }
200        let charset = response_charset(&response);
201        if charset.is_some() || has_text_bom(&response.body) {
202            let text = decode_text_bytes(&response.body, charset.as_deref())?;
203            return Ok((text, response.final_url, debug_logs));
204        }
205        if looks_like_binary(&sample(&response.body)) {
206            debug_logs.push("Heuristic detection: binary data".to_string());
207            return Ok((
208                BINARY_DATA_MARKER.to_string(),
209                response.final_url,
210                debug_logs,
211            ));
212        }
213        let text = decode_text_bytes(&response.body, charset.as_deref())?;
214        Ok((text, response.final_url, debug_logs))
215    }
216}
217
218async fn read_source_text_file(path: &std::path::Path) -> Result<String> {
219    decode_text_bytes(&tokio::fs::read(path).await?, None)
220}
221
222/// Writes retained raw files under a target directory.
223pub async fn write_raw_files(
224    raw_files: &BTreeMap<String, String>,
225    directory: &std::path::Path,
226) -> Result<Vec<PathBuf>> {
227    tokio::fs::create_dir_all(directory).await?;
228    let mut written = Vec::with_capacity(raw_files.len());
229    for (name, text) in raw_files {
230        let path = directory.join(name);
231        if tokio::fs::metadata(&path).await.is_err() {
232            tokio::fs::write(&path, n_utf8_text_file_bytes(text)).await?;
233        }
234        written.push(path);
235    }
236    Ok(written)
237}
238
239fn n_utf8_text_file_bytes(text: &str) -> Vec<u8> {
240    let mut bytes = Vec::with_capacity(3 + text.len());
241    bytes.extend_from_slice(&[0xef, 0xbb, 0xbf]);
242    bytes.extend_from_slice(text.as_bytes());
243    bytes
244}
245
246fn detect_source_kind(text: &str) -> Result<LoadedSourceKind> {
247    let trimmed = text.trim();
248    if trimmed.starts_with("#EXTM3U") {
249        Ok(LoadedSourceKind::Hls)
250    } else if trimmed.contains("</MPD>") && trimmed.contains("<MPD") {
251        Ok(LoadedSourceKind::Dash)
252    } else if trimmed.contains("</SmoothStreamingMedia>")
253        && trimmed.contains("<SmoothStreamingMedia")
254    {
255        Ok(LoadedSourceKind::Mss)
256    } else if trimmed == HTTP_LIVE_TS_MARKER {
257        Ok(LoadedSourceKind::HttpLiveTs)
258    } else if trimmed == BINARY_DATA_MARKER {
259        Ok(LoadedSourceKind::BinaryData)
260    } else {
261        Err(Error::compatibility("source input type is not supported"))
262    }
263}
264
265fn response_charset(response: &HttpResponse) -> Option<String> {
266    response
267        .headers
268        .get("content-type")
269        .and_then(|value| content_type_charset(value))
270}
271
272fn decode_text_bytes(bytes: &[u8], charset: Option<&str>) -> Result<String> {
273    if bytes.starts_with(&[0xef, 0xbb, 0xbf]) {
274        return Ok(String::from_utf8_lossy(&bytes[3..]).to_string());
275    }
276    if bytes.starts_with(&[0xff, 0xfe]) {
277        return decode_utf16(&bytes[2..], false);
278    }
279    if bytes.starts_with(&[0xfe, 0xff]) {
280        return decode_utf16(&bytes[2..], true);
281    }
282    match charset {
283        Some("utf-16") => decode_utf16_auto(bytes),
284        Some("utf-16le") => decode_utf16(bytes, false),
285        Some("utf-16be") => decode_utf16(bytes, true),
286        Some(label) => Ok(decode_with_label(label, bytes)),
287        None => Ok(String::from_utf8_lossy(bytes).to_string()),
288    }
289}
290
291fn has_text_bom(bytes: &[u8]) -> bool {
292    bytes.starts_with(&[0xef, 0xbb, 0xbf])
293        || bytes.starts_with(&[0xff, 0xfe])
294        || bytes.starts_with(&[0xfe, 0xff])
295}
296
297fn decode_with_label(label: &str, bytes: &[u8]) -> String {
298    match encoding_rs::Encoding::for_label(label.as_bytes()) {
299        Some(encoding) => {
300            let (text, _, _) = encoding.decode(bytes);
301            text.into_owned()
302        }
303        None => String::from_utf8_lossy(bytes).to_string(),
304    }
305}
306
307fn content_type_charset(content_type: &str) -> Option<String> {
308    content_type.split(';').find_map(|part| {
309        let (key, value) = part.trim().split_once('=')?;
310        if key.eq_ignore_ascii_case("charset") {
311            Some(value.trim_matches('"').to_ascii_lowercase())
312        } else {
313            None
314        }
315    })
316}
317
318fn decode_utf16_auto(bytes: &[u8]) -> Result<String> {
319    if bytes.starts_with(&[0xff, 0xfe]) {
320        return decode_utf16(&bytes[2..], false);
321    }
322    if bytes.starts_with(&[0xfe, 0xff]) {
323        return decode_utf16(&bytes[2..], true);
324    }
325    decode_utf16(bytes, false)
326}
327
328fn decode_utf16(bytes: &[u8], big_endian: bool) -> Result<String> {
329    let mut units = Vec::with_capacity(bytes.len() / 2);
330    for pair in bytes.chunks_exact(2) {
331        let unit = if big_endian {
332            u16::from_be_bytes([pair[0], pair[1]])
333        } else {
334            u16::from_le_bytes([pair[0], pair[1]])
335        };
336        units.push(unit);
337    }
338    let mut text = String::from_utf16_lossy(&units);
339    if bytes.len() & 1 != 0 {
340        text.push(char::REPLACEMENT_CHARACTER);
341    }
342    Ok(text)
343}
344
345fn sample(bytes: &[u8]) -> Vec<u8> {
346    bytes.iter().take(4096).copied().collect()
347}
348
349fn looks_like_binary(data: &[u8]) -> bool {
350    if data.is_empty() {
351        return false;
352    }
353    let mut non_text = 0_usize;
354    let mut total = 0_usize;
355    let mut index = 0_usize;
356    while index < data.len() {
357        let byte = data[index];
358        total += 1;
359        if byte == 0 {
360            return true;
361        }
362        if (0x20..=0x7e).contains(&byte) || matches!(byte, 0x09 | 0x0a | 0x0d) {
363            index += 1;
364            continue;
365        }
366        let seq_len = utf8_sequence_length(byte);
367        if seq_len > 1
368            && index + seq_len <= data.len()
369            && valid_utf8_sequence(&data[index..index + seq_len])
370        {
371            index += seq_len;
372            continue;
373        }
374        non_text += 1;
375        index += 1;
376    }
377    (non_text as f64 / total as f64) > 0.3
378}
379
380fn utf8_sequence_length(byte: u8) -> usize {
381    if byte & 0x80 == 0x00 {
382        1
383    } else if byte & 0xe0 == 0xc0 {
384        2
385    } else if byte & 0xf0 == 0xe0 {
386        3
387    } else if byte & 0xf8 == 0xf0 {
388        4
389    } else {
390        1
391    }
392}
393
394fn valid_utf8_sequence(seq: &[u8]) -> bool {
395    if seq.len() <= 1 {
396        return false;
397    }
398    seq.iter().skip(1).all(|byte| byte & 0xc0 == 0x80)
399}
400
401fn is_mpeg2_ts_buffer(buffer: &[u8]) -> bool {
402    const PACKET_SIZE: usize = 188;
403    if buffer.len() < PACKET_SIZE {
404        return false;
405    }
406    let packet_count = std::cmp::min(buffer.len() / PACKET_SIZE, 5);
407    let sync_count = (0..packet_count)
408        .filter(|index| buffer[index * PACKET_SIZE] == 0x47)
409        .count();
410    sync_count >= 3
411}
412
413fn file_uri_to_path(value: &str) -> PathBuf {
414    if let Ok(url) = reqwest::Url::parse(value)
415        && url.scheme() == "file"
416        && let Ok(path) = url.to_file_path()
417    {
418        return path;
419    }
420    let stripped = value
421        .trim_start_matches("file://")
422        .trim_start_matches("file:");
423    #[cfg(windows)]
424    {
425        let mut normalized = stripped.trim_start_matches('/');
426        if let Some(rest) = normalized.strip_prefix("?/") {
427            normalized = rest;
428        }
429        PathBuf::from(normalized)
430    }
431    #[cfg(not(windows))]
432    {
433        PathBuf::from(stripped)
434    }
435}
436
437fn path_to_file_uri(path: &Path) -> Option<String> {
438    let text = path.to_str()?;
439    #[cfg(windows)]
440    {
441        let text = text.replace('\\', "/");
442        let normalized = text.strip_prefix("//?/").unwrap_or(&text);
443        Some(format!("file:///{normalized}"))
444    }
445    #[cfg(not(windows))]
446    {
447        Some(format!("file://{text}"))
448    }
449}
450
451fn hex_to_bytes(value: &str) -> Result<Vec<u8>> {
452    let value = value.trim();
453    let value = value
454        .strip_prefix("0x")
455        .or_else(|| value.strip_prefix("0X"))
456        .unwrap_or(value);
457    if value.len() & 1 != 0 {
458        return Err(Error::config("hex length must be even"));
459    }
460    let mut bytes = Vec::with_capacity(value.len() / 2);
461    let mut chars = value.chars();
462    while let Some(high) = chars.next() {
463        let low = chars
464            .next()
465            .ok_or_else(|| Error::config("hex length must be even"))?;
466        let high = hex_value(high).ok_or_else(|| Error::config("hex is invalid"))?;
467        let low = hex_value(low).ok_or_else(|| Error::config("hex is invalid"))?;
468        bytes.push((high << 4) | low);
469    }
470    Ok(bytes)
471}
472
473fn hex_value(ch: char) -> Option<u8> {
474    match ch {
475        '0'..='9' => Some(ch as u8 - b'0'),
476        'a'..='f' => Some(ch as u8 - b'a' + 10),
477        'A'..='F' => Some(ch as u8 - b'A' + 10),
478        _ => None,
479    }
480}
481
482fn base64_decode(value: &str) -> Result<Vec<u8>> {
483    crate::base64::decode_base64(value).map_err(Error::config)
484}
haki_dl/source.rs

haki_dl/
source.rs
