ani_lib/
hls_relay.rs

1use std::{
2    collections::HashMap,
3    convert::Infallible,
4    net::SocketAddr,
5    sync::{
6        Arc, Mutex,
7        atomic::{AtomicUsize, Ordering},
8    },
9};
10
11use bytes::Bytes;
12use http_body_util::Full;
13use hyper::{
14    Method, Request, Response, StatusCode, body::Incoming, header, server::conn::http1,
15    service::service_fn,
16};
17use hyper_util::rt::TokioIo;
18use reqwest::Client;
19use sha2::{Digest, Sha256};
20use tokio::{net::TcpListener, sync::oneshot, task::JoinHandle};
21use tracing::{debug, info};
22use url::Url;
23
24use crate::{AniError, RequestHeaders, Result, StreamLink, SubtitleTrack};
25
26const MAX_PLAYLIST_BYTES: usize = 2 * 1024 * 1024;
27const MAX_SUBTITLE_BYTES: usize = 2 * 1024 * 1024;
28const MAX_RESOURCES: usize = 16_384;
29const SCAN_LIMIT: usize = 64 * 1024;
30const FALLBACK_SUBTITLE_DURATION_SECONDS: f64 = 30.0 * 60.0;
31
32#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
33enum ResourceKind {
34    EntryPlaylist,
35    Playlist,
36    Segment,
37    Subtitle,
38    SubtitlePlaylist,
39    Resource,
40}
41
42#[derive(Clone)]
43struct Registered {
44    url: Url,
45    headers: RequestHeaders,
46    kind: ResourceKind,
47    subtitles: Vec<SubtitleTrack>,
48    duration_seconds: Option<f64>,
49}
50
51struct State {
52    client: Client,
53    base: SocketAddr,
54    resources: Mutex<HashMap<String, Registered>>,
55    tokens_by_url: Mutex<HashMap<String, String>>,
56    counter: AtomicUsize,
57    secret: String,
58}
59
60/// A loopback-only HLS relay. Dropping it stops accepting new requests.
61pub struct HlsRelay {
62    shutdown: Option<oneshot::Sender<()>>,
63    task: JoinHandle<()>,
64}
65
66impl Drop for HlsRelay {
67    fn drop(&mut self) {
68        if let Some(shutdown) = self.shutdown.take() {
69            let _ = shutdown.send(());
70        }
71        // Dropping the join handle detaches the task; the shutdown signal makes
72        // its accept loop exit without cancelling an in-flight response.
73        let _ = &self.task;
74    }
75}
76
77/// Starts a relay and returns a stream whose URL points at its loopback endpoint.
78///
79/// Subtitles are exposed both as plain sidecar URLs on the returned
80/// `StreamLink::subtitles` (used by desktop players via `--sub-file`) and as
81/// synthetic HLS subtitle renditions in the relayed master playlist. Android
82/// players only receive a single intent URL, so they rely on the HLS
83/// rendition to discover subtitles at all.
84pub async fn relay_stream(stream: &StreamLink) -> Result<(HlsRelay, StreamLink)> {
85    relay_stream_inner(stream, true).await
86}
87
88/// Like [`relay_stream`], but never exposes subtitles as synthetic HLS
89/// subtitle renditions in the master playlist.
90///
91/// Wrapping a single, potentially very long, subtitle file as one oversized
92/// HLS media segment does not match how HLS clients expect WebVTT segments
93/// to be chunked (RFC 8216 section 3.5), and it produces unreliable cue
94/// timing in some HLS demuxers (see issue #18). Desktop players do not need
95/// this: they are launched with an explicit `--sub-file` argument that
96/// points directly at the relayed sidecar subtitle, which this function
97/// still populates on the returned `StreamLink::subtitles`.
98pub async fn relay_stream_without_hls_subtitles(
99    stream: &StreamLink,
100) -> Result<(HlsRelay, StreamLink)> {
101    relay_stream_inner(stream, false).await
102}
103
104async fn relay_stream_inner(
105    stream: &StreamLink,
106    expose_subtitles_in_hls: bool,
107) -> Result<(HlsRelay, StreamLink)> {
108    if !stream.hls {
109        return Err(AniError::InvalidInput {
110            input: stream.url.clone(),
111            reason: "only HLS streams can be relayed".to_string(),
112        });
113    }
114    let upstream = validate_upstream(&stream.url)?;
115    let listener = TcpListener::bind("127.0.0.1:0").await?;
116    let address = listener.local_addr()?;
117    info!(
118        bind_address = %address,
119        upstream_host = %upstream.host_str().unwrap_or("?"),
120        expose_subtitles_in_hls,
121        subtitle_tracks = stream.subtitles.len(),
122        "starting loopback HLS relay",
123    );
124    let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
125    let state = Arc::new(State {
126        client: Client::builder()
127            .redirect(reqwest::redirect::Policy::limited(8))
128            .build()?,
129        base: address,
130        resources: Mutex::new(HashMap::new()),
131        tokens_by_url: Mutex::new(HashMap::new()),
132        counter: AtomicUsize::new(0),
133        secret: format!("{}-{}", std::process::id(), unix_nanos()),
134    });
135    let token = register(
136        &state,
137        upstream,
138        stream.headers.clone(),
139        ResourceKind::EntryPlaylist,
140    )?;
141    let server_state = Arc::clone(&state);
142    let task = tokio::spawn(async move {
143        loop {
144            tokio::select! {
145                _ = &mut shutdown_rx => break,
146                accepted = listener.accept() => {
147                    let Ok((socket, _)) = accepted else { break };
148                    let state = Arc::clone(&server_state);
149                    tokio::spawn(async move {
150                        let service = service_fn(move |request| handle(Arc::clone(&state), request));
151                        let _ = http1::Builder::new().serve_connection(TokioIo::new(socket), service).await;
152                    });
153                }
154            }
155        }
156    });
157    let mut local = stream.clone();
158    local.url = local_url(address, &token);
159    local.headers = RequestHeaders::default();
160    let mut playlist_subtitles = Vec::with_capacity(local.subtitles.len());
161    for track in &mut local.subtitles {
162        let url = validate_upstream(&track.url)?;
163        let token = register(
164            &state,
165            url.clone(),
166            stream.headers.clone(),
167            ResourceKind::Subtitle,
168        )?;
169        track.url = local_url(address, &token);
170
171        if !expose_subtitles_in_hls {
172            continue;
173        }
174        let duration_seconds = probe_subtitle_duration(&state.client, &url, &stream.headers).await;
175        let playlist_token = register(
176            &state,
177            url,
178            stream.headers.clone(),
179            ResourceKind::SubtitlePlaylist,
180        )?;
181        let mut playlist_track = track.clone();
182        playlist_track.url = local_url(address, &playlist_token);
183        playlist_subtitles.push(playlist_track);
184        if let Some(entry) = state
185            .resources
186            .lock()
187            .expect("relay registry poisoned")
188            .get_mut(&playlist_token)
189        {
190            entry.duration_seconds = duration_seconds;
191        }
192    }
193    if !playlist_subtitles.is_empty()
194        && let Some(entry) = state
195            .resources
196            .lock()
197            .expect("relay registry poisoned")
198            .get_mut(&token)
199    {
200        entry.subtitles = playlist_subtitles;
201    }
202    debug!(
203        local_url = %local.url,
204        bind_address = %address,
205        "HLS relay is ready to serve the rewritten stream URL",
206    );
207    Ok((
208        HlsRelay {
209            shutdown: Some(shutdown_tx),
210            task,
211        },
212        local,
213    ))
214}
215
216async fn handle(
217    state: Arc<State>,
218    request: Request<Incoming>,
219) -> std::result::Result<Response<Full<Bytes>>, Infallible> {
220    let response = match handle_inner(&state, request).await {
221        Ok(response) => response,
222        Err(error) => response(
223            StatusCode::BAD_GATEWAY,
224            "text/plain",
225            error.to_string().into_bytes(),
226        ),
227    };
228    Ok(response)
229}
230
231async fn handle_inner(
232    state: &Arc<State>,
233    request: Request<Incoming>,
234) -> Result<Response<Full<Bytes>>> {
235    if request.method() == Method::OPTIONS {
236        return Ok(Response::builder()
237            .status(StatusCode::NO_CONTENT)
238            .header(header::ACCESS_CONTROL_ALLOW_ORIGIN, "*")
239            .header(header::ACCESS_CONTROL_ALLOW_METHODS, "GET, HEAD, OPTIONS")
240            .header(header::ACCESS_CONTROL_ALLOW_HEADERS, "Range, Content-Type")
241            .header(
242                header::ACCESS_CONTROL_EXPOSE_HEADERS,
243                "Content-Length, Content-Range, Accept-Ranges",
244            )
245            .body(Full::new(Bytes::new()))
246            .expect("valid CORS response"));
247    }
248
249    if request.method() != Method::GET && request.method() != Method::HEAD {
250        return Ok(response(
251            StatusCode::METHOD_NOT_ALLOWED,
252            "text/plain",
253            b"method not allowed".to_vec(),
254        ));
255    }
256    let token = request
257        .uri()
258        .path()
259        .strip_prefix("/r/")
260        .filter(|v| !v.is_empty());
261    let Some(token) = token else {
262        return Ok(response(
263            StatusCode::NOT_FOUND,
264            "text/plain",
265            b"not found".to_vec(),
266        ));
267    };
268    let registered = state
269        .resources
270        .lock()
271        .expect("relay registry poisoned")
272        .get(token)
273        .cloned();
274    let Some(registered) = registered else {
275        return Ok(response(
276            StatusCode::FORBIDDEN,
277            "text/plain",
278            b"invalid relay token".to_vec(),
279        ));
280    };
281    if registered.kind == ResourceKind::SubtitlePlaylist {
282        let segment_token = register(
283            state,
284            registered.url.clone(),
285            registered.headers.clone(),
286            ResourceKind::Subtitle,
287        )?;
288        let body = subtitle_playlist(
289            &local_url(state.base, &segment_token),
290            registered
291                .duration_seconds
292                .unwrap_or(FALLBACK_SUBTITLE_DURATION_SECONDS),
293        );
294        return Ok(if request.method() == Method::HEAD {
295            response(StatusCode::OK, "application/vnd.apple.mpegurl", Vec::new())
296        } else {
297            response(
298                StatusCode::OK,
299                "application/vnd.apple.mpegurl",
300                body.into_bytes(),
301            )
302        });
303    }
304    let mut upstream = state
305        .client
306        .request(request.method().clone(), registered.url.clone());
307    if let Some(value) = &registered.headers.referer {
308        upstream = upstream.header(header::REFERER, value);
309    }
310    if let Some(value) = &registered.headers.origin {
311        upstream = upstream.header(header::ORIGIN, value);
312    }
313    for (name, value) in &registered.headers.extra {
314        upstream = upstream.header(name, value);
315    }
316    let conditional_headers = [header::IF_NONE_MATCH, header::IF_MODIFIED_SINCE];
317    if registered.kind != ResourceKind::Segment
318        && let Some(value) = request.headers().get(header::RANGE)
319    {
320        upstream = upstream.header(header::RANGE, value);
321    }
322    for name in conditional_headers {
323        if let Some(value) = request.headers().get(&name) {
324            upstream = upstream.header(name, value);
325        }
326    }
327    let upstream = upstream.send().await?;
328    let status = upstream.status();
329    let upstream_headers = upstream.headers().clone();
330    let content_type = upstream
331        .headers()
332        .get(header::CONTENT_TYPE)
333        .and_then(|v| v.to_str().ok())
334        .unwrap_or("")
335        .to_owned();
336    if request.method() == Method::HEAD {
337        let corrected_type = corrected_content_type(&registered, &content_type, false);
338        let mut result = response(status, &corrected_type, Vec::new());
339        copy_upstream_headers(
340            &upstream_headers,
341            result.headers_mut(),
342            registered.kind != ResourceKind::Segment,
343        );
344        return Ok(result);
345    }
346    let bytes = upstream.bytes().await?.to_vec();
347    let playlist = matches!(
348        registered.kind,
349        ResourceKind::EntryPlaylist | ResourceKind::Playlist
350    ) || content_type.contains("mpegurl")
351        || registered
352            .url
353            .path()
354            .to_ascii_lowercase()
355            .ends_with(".m3u8")
356        || bytes.starts_with(b"#EXTM3U");
357    if playlist {
358        if bytes.len() > MAX_PLAYLIST_BYTES {
359            return Err(AniError::ProviderResponseSizeExceeded {
360                provider: "HLS relay upstream".to_string(),
361                size: bytes.len(),
362                limit: MAX_PLAYLIST_BYTES,
363            });
364        }
365        let text = String::from_utf8(bytes).map_err(|error| AniError::ProviderInvalidResponse {
366            provider: "HLS relay upstream".to_string(),
367            message: "provider playlist is not UTF-8".to_string(),
368            source: Some(Box::new(error)),
369        })?;
370        let rewritten = rewrite_playlist(state, &registered, &text)?;
371        let rewritten = if registered.kind == ResourceKind::EntryPlaylist {
372            expose_subtitles(state, &registered, &text, rewritten)?
373        } else {
374            rewritten
375        };
376        return Ok(response(
377            status,
378            "application/vnd.apple.mpegurl",
379            rewritten.into_bytes(),
380        ));
381    }
382    let (bytes, stripped) = if registered.kind == ResourceKind::Segment {
383        strip_png_wrapper(bytes)
384    } else {
385        (bytes, false)
386    };
387    let corrected_type = corrected_content_type(&registered, &content_type, stripped);
388    let mut result = response(status, &corrected_type, bytes);
389    copy_upstream_headers(&upstream_headers, result.headers_mut(), false);
390    Ok(result)
391}
392
393fn rewrite_playlist(state: &Arc<State>, parent: &Registered, body: &str) -> Result<String> {
394    let mut output = String::with_capacity(body.len() + 256);
395    let mut next_uri_is_playlist = false;
396    for line in body.lines() {
397        let trimmed = line.trim();
398        let rewritten = if trimmed.is_empty() {
399            line.to_owned()
400        } else if trimmed.starts_with('#') {
401            if trimmed
402                .to_ascii_uppercase()
403                .starts_with("#EXT-X-STREAM-INF:")
404            {
405                next_uri_is_playlist = true;
406            }
407            rewrite_uri_attributes(state, parent, line)?
408        } else {
409            let kind = if next_uri_is_playlist {
410                ResourceKind::Playlist
411            } else {
412                ResourceKind::Segment
413            };
414            next_uri_is_playlist = false;
415            relay_reference(state, parent, trimmed, kind)?
416        };
417        output.push_str(&rewritten);
418        output.push('\n');
419    }
420    Ok(output)
421}
422
423fn expose_subtitles(
424    state: &Arc<State>,
425    parent: &Registered,
426    original: &str,
427    rewritten: String,
428) -> Result<String> {
429    if parent.subtitles.is_empty() {
430        return Ok(rewritten);
431    }
432    if original.lines().any(|line| {
433        line.trim_start()
434            .to_ascii_uppercase()
435            .starts_with("#EXT-X-STREAM-INF:")
436    }) {
437        return Ok(add_subtitles_to_master(&rewritten, &parent.subtitles));
438    }
439
440    let media_token = register(
441        state,
442        parent.url.clone(),
443        parent.headers.clone(),
444        ResourceKind::Playlist,
445    )?;
446    let mut master = String::from("#EXTM3U\n#EXT-X-VERSION:3\n");
447    append_subtitle_renditions(&mut master, &parent.subtitles);
448    master.push_str("#EXT-X-STREAM-INF:BANDWIDTH=1,SUBTITLES=\"aniplay-subs\"\n");
449    master.push_str(&local_url(state.base, &media_token));
450    master.push('\n');
451    Ok(master)
452}
453
454fn add_subtitles_to_master(body: &str, subtitles: &[SubtitleTrack]) -> String {
455    let mut output = String::with_capacity(body.len() + subtitles.len() * 160);
456    let mut inserted = false;
457    for line in body.lines() {
458        output.push_str(line);
459        output.push('\n');
460        if !inserted && line.trim().eq_ignore_ascii_case("#EXTM3U") {
461            append_subtitle_renditions(&mut output, subtitles);
462            inserted = true;
463        }
464    }
465    if !inserted {
466        append_subtitle_renditions(&mut output, subtitles);
467    }
468    output
469        .lines()
470        .map(|line| {
471            if line
472                .trim_start()
473                .to_ascii_uppercase()
474                .starts_with("#EXT-X-STREAM-INF:")
475                && !line.to_ascii_uppercase().contains("SUBTITLES=")
476            {
477                format!("{line},SUBTITLES=\"aniplay-subs\"")
478            } else {
479                line.to_owned()
480            }
481        })
482        .collect::<Vec<_>>()
483        .join("\n")
484        + "\n"
485}
486
487fn append_subtitle_renditions(output: &mut String, subtitles: &[SubtitleTrack]) {
488    let default_index = subtitles
489        .iter()
490        .position(|subtitle| subtitle.default)
491        .unwrap_or(0);
492    for (index, subtitle) in subtitles.iter().enumerate() {
493        let name = hls_attribute(&subtitle.label);
494        let default = if index == default_index { "YES" } else { "NO" };
495        output.push_str(&format!(
496            "#EXT-X-MEDIA:TYPE=SUBTITLES,GROUP-ID=\"aniplay-subs\",NAME=\"{name}\",DEFAULT={default},AUTOSELECT=YES,FORCED=NO,URI=\"{}\"\n",
497            subtitle.url
498        ));
499    }
500}
501
502fn hls_attribute(value: &str) -> String {
503    value
504        .replace(['\r', '\n'], " ")
505        .replace('"', "'")
506        .trim()
507        .to_owned()
508}
509
510fn subtitle_playlist(segment_url: &str, duration_seconds: f64) -> String {
511    let duration_seconds = duration_seconds.max(1.0);
512    let target_duration = duration_seconds.ceil() as u64;
513    format!(
514        "#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-PLAYLIST-TYPE:VOD\n#EXT-X-TARGETDURATION:{target_duration}\n#EXT-X-MEDIA-SEQUENCE:0\n#EXTINF:{duration_seconds:.3},\n{segment_url}\n#EXT-X-ENDLIST\n"
515    )
516}
517
518async fn probe_subtitle_duration(
519    client: &Client,
520    url: &Url,
521    headers: &RequestHeaders,
522) -> Option<f64> {
523    let mut request = client.get(url.clone());
524    if let Some(value) = &headers.referer {
525        request = request.header(header::REFERER, value);
526    }
527    if let Some(value) = &headers.origin {
528        request = request.header(header::ORIGIN, value);
529    }
530    for (name, value) in &headers.extra {
531        request = request.header(name, value);
532    }
533    let response = request.send().await.ok()?.error_for_status().ok()?;
534    let bytes = response.bytes().await.ok()?;
535    if bytes.len() > MAX_SUBTITLE_BYTES {
536        return None;
537    }
538    subtitle_duration(std::str::from_utf8(&bytes).ok()?)
539}
540
541fn subtitle_duration(body: &str) -> Option<f64> {
542    body.lines()
543        .filter_map(|line| line.split_once("-->").map(|(_, end)| end))
544        .filter_map(parse_subtitle_timestamp)
545        .reduce(f64::max)
546        .map(|seconds| seconds + 1.0)
547}
548
549fn parse_subtitle_timestamp(value: &str) -> Option<f64> {
550    let timestamp = value
551        .trim()
552        .split_ascii_whitespace()
553        .next()?
554        .replace(',', ".");
555    let parts = timestamp
556        .split(':')
557        .map(str::parse::<f64>)
558        .collect::<std::result::Result<Vec<_>, _>>()
559        .ok()?;
560    match parts.as_slice() {
561        [minutes, seconds] => Some(minutes * 60.0 + seconds),
562        [hours, minutes, seconds] => Some(hours * 3600.0 + minutes * 60.0 + seconds),
563        _ => None,
564    }
565}
566
567fn rewrite_uri_attributes(state: &Arc<State>, parent: &Registered, line: &str) -> Result<String> {
568    let mut output = line.to_owned();
569    let mut offset = 0;
570    let upper = line.trim().to_ascii_uppercase();
571    let kind =
572        if upper.starts_with("#EXT-X-MEDIA:") || upper.starts_with("#EXT-X-I-FRAME-STREAM-INF:") {
573            ResourceKind::Playlist
574        } else {
575            ResourceKind::Resource
576        };
577    while let Some(found) = output[offset..].to_ascii_uppercase().find("URI=") {
578        let quote_at = offset + found + 4;
579        let Some(quote) = output.as_bytes().get(quote_at).copied() else {
580            break;
581        };
582        if quote != b'\'' && quote != b'"' {
583            offset = quote_at + 1;
584            continue;
585        }
586        let start = quote_at + 1;
587        let Some(end_rel) = output.as_bytes()[start..]
588            .iter()
589            .position(|value| *value == quote)
590        else {
591            break;
592        };
593        let end = start + end_rel;
594        let rewritten = relay_reference(state, parent, &output[start..end], kind)?;
595        output.replace_range(start..end, &rewritten);
596        offset = start + rewritten.len() + 1;
597    }
598    Ok(output)
599}
600
601fn relay_reference(
602    state: &Arc<State>,
603    parent: &Registered,
604    reference: &str,
605    kind: ResourceKind,
606) -> Result<String> {
607    let resolved =
608        parent
609            .url
610            .join(reference)
611            .map_err(|error| AniError::ProviderInvalidResponse {
612                provider: "HLS relay upstream".to_string(),
613                message: format!("invalid playlist URL: {error}"),
614                source: Some(Box::new(error)),
615            })?;
616    validate_url(&resolved)?;
617    let token = register(state, resolved, parent.headers.clone(), kind)?;
618    Ok(local_url(state.base, &token))
619}
620
621fn register(
622    state: &Arc<State>,
623    url: Url,
624    headers: RequestHeaders,
625    kind: ResourceKind,
626) -> Result<String> {
627    validate_url(&url)?;
628    let url_key = format!("{kind:?}:{url}");
629    if let Some(token) = state
630        .tokens_by_url
631        .lock()
632        .expect("relay URL registry poisoned")
633        .get(&url_key)
634        .cloned()
635    {
636        return Ok(token);
637    }
638    let mut resources = state.resources.lock().expect("relay registry poisoned");
639    if resources.len() >= MAX_RESOURCES {
640        return Err(AniError::ProviderInvalidResponse {
641            provider: "HLS relay upstream".to_string(),
642            message: "provider playlist contains too many resources".to_string(),
643            source: None,
644        });
645    }
646    let count = state.counter.fetch_add(1, Ordering::Relaxed);
647    let token =
648        hex::encode(Sha256::digest(format!("{}:{count}:{url}", state.secret)))[..32].to_owned();
649    resources.insert(
650        token.clone(),
651        Registered {
652            url,
653            headers,
654            kind,
655            subtitles: Vec::new(),
656            duration_seconds: None,
657        },
658    );
659    state
660        .tokens_by_url
661        .lock()
662        .expect("relay URL registry poisoned")
663        .insert(url_key, token.clone());
664    Ok(token)
665}
666
667fn corrected_content_type(registered: &Registered, upstream: &str, stripped: bool) -> String {
668    if matches!(
669        registered.kind,
670        ResourceKind::EntryPlaylist | ResourceKind::Playlist
671    ) {
672        return "application/vnd.apple.mpegurl".into();
673    }
674    if registered.kind == ResourceKind::Subtitle {
675        let path = registered.url.path().to_ascii_lowercase();
676        return if path.ends_with(".vtt") {
677            "text/vtt".into()
678        } else if path.ends_with(".srt") {
679            "application/x-subrip".into()
680        } else if path.ends_with(".ass") || path.ends_with(".ssa") {
681            "text/x-ssa".into()
682        } else if upstream.is_empty() {
683            "text/plain; charset=utf-8".into()
684        } else {
685            upstream.into()
686        };
687    }
688    if registered.kind != ResourceKind::Segment {
689        return if upstream.is_empty() {
690            "application/octet-stream".into()
691        } else {
692            upstream.into()
693        };
694    }
695    let path = registered.url.path().to_ascii_lowercase();
696    if path.ends_with(".m4s") || path.ends_with(".mp4") || path.ends_with(".m4v") {
697        "video/mp4".into()
698    } else if path.ends_with(".aac") {
699        "audio/aac".into()
700    } else if stripped || upstream.is_empty() || upstream.to_ascii_lowercase().starts_with("image/")
701    {
702        "video/mp2t".into()
703    } else {
704        upstream.into()
705    }
706}
707
708fn validate_upstream(value: &str) -> Result<Url> {
709    let url = Url::parse(value).map_err(|error| AniError::InvalidInput {
710        input: value.to_string(),
711        reason: format!("invalid HLS URL: {error}"),
712    })?;
713    validate_url(&url)?;
714    Ok(url)
715}
716
717fn validate_url(url: &Url) -> Result<()> {
718    if !url.username().is_empty() || url.password().is_some() {
719        return Err(AniError::ProviderUrlValidationFailed {
720            reason: "credential-bearing playlist URLs are not allowed".to_string(),
721        });
722    }
723    let loopback = url
724        .host_str()
725        .is_some_and(|host| host == "127.0.0.1" || host == "localhost" || host == "::1");
726    if url.scheme() != "https" && !(cfg!(test) && loopback && url.scheme() == "http") {
727        return Err(AniError::ProviderUrlValidationFailed {
728            reason: "relay resources must use HTTPS".to_string(),
729        });
730    }
731    Ok(())
732}
733
734fn strip_png_wrapper(bytes: Vec<u8>) -> (Vec<u8>, bool) {
735    if !bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
736        return (bytes, false);
737    }
738    let limit = bytes.len().min(SCAN_LIMIT);
739    for offset in 8..limit.saturating_sub(376) {
740        if bytes[offset] == 0x47 && bytes[offset + 188] == 0x47 && bytes[offset + 376] == 0x47 {
741            return (bytes[offset..].to_vec(), true);
742        }
743    }
744    (bytes, false)
745}
746
747fn response(status: StatusCode, content_type: &str, bytes: Vec<u8>) -> Response<Full<Bytes>> {
748    Response::builder()
749        .status(status)
750        .header(header::CONTENT_TYPE, content_type)
751        .header(header::CONTENT_LENGTH, bytes.len())
752        .header(header::ACCESS_CONTROL_ALLOW_ORIGIN, "*")
753        .header(header::ACCESS_CONTROL_ALLOW_METHODS, "GET, HEAD, OPTIONS")
754        .header(header::ACCESS_CONTROL_ALLOW_HEADERS, "Range, Content-Type")
755        .header(
756            header::ACCESS_CONTROL_EXPOSE_HEADERS,
757            "Content-Length, Content-Range, Accept-Ranges",
758        )
759        .body(Full::new(Bytes::from(bytes)))
760        .expect("valid relay response")
761}
762
763fn copy_upstream_headers(
764    source: &hyper::HeaderMap,
765    destination: &mut hyper::HeaderMap,
766    include_length: bool,
767) {
768    for name in [
769        header::CONTENT_RANGE,
770        header::ACCEPT_RANGES,
771        header::ETAG,
772        header::LAST_MODIFIED,
773        header::CACHE_CONTROL,
774    ] {
775        if let Some(value) = source.get(&name) {
776            destination.insert(name, value.clone());
777        }
778    }
779    if include_length && let Some(value) = source.get(header::CONTENT_LENGTH) {
780        destination.insert(header::CONTENT_LENGTH, value.clone());
781    }
782}
783
784fn local_url(address: SocketAddr, token: &str) -> String {
785    format!("http://{address}/r/{token}")
786}
787fn unix_nanos() -> u128 {
788    std::time::SystemTime::now()
789        .duration_since(std::time::UNIX_EPOCH)
790        .unwrap_or_default()
791        .as_nanos()
792}
793
794#[cfg(test)]
795mod tests {
796    use super::*;
797    use wiremock::{
798        Mock, MockServer, ResponseTemplate,
799        matchers::{method, path},
800    };
801
802    #[test]
803    fn strips_only_confirmed_png_wrapped_transport_streams() {
804        let mut wrapped = b"\x89PNG\r\n\x1a\nnot-really-png".to_vec();
805        let offset = wrapped.len();
806        wrapped.resize(offset + 377, 0);
807        wrapped[offset] = 0x47;
808        wrapped[offset + 188] = 0x47;
809        wrapped[offset + 376] = 0x47;
810        let (result, stripped) = strip_png_wrapper(wrapped);
811        assert!(stripped);
812        assert_eq!(result[0], 0x47);
813    }
814
815    #[test]
816    fn leaves_real_png_data_untouched() {
817        let value = b"\x89PNG\r\n\x1a\nordinary image".to_vec();
818        assert_eq!(strip_png_wrapper(value.clone()), (value, false));
819    }
820
821    #[test]
822    fn adds_external_subtitles_to_master_variants() {
823        let subtitles = vec![
824            SubtitleTrack {
825                label: "English".into(),
826                url: "http://127.0.0.1:1234/r/subtitle-one".into(),
827                default: true,
828            },
829            SubtitleTrack {
830                label: "Signs \"and\" Songs".into(),
831                url: "http://127.0.0.1:1234/r/subtitle-two".into(),
832                default: false,
833            },
834        ];
835        let master = add_subtitles_to_master(
836            "#EXTM3U\n#EXT-X-STREAM-INF:BANDWIDTH=1000\nvideo.m3u8\n",
837            &subtitles,
838        );
839
840        assert!(master.contains("TYPE=SUBTITLES,GROUP-ID=\"aniplay-subs\""));
841        assert!(master.contains("NAME=\"English\",DEFAULT=YES"));
842        assert!(master.contains("NAME=\"Signs 'and' Songs\",DEFAULT=NO"));
843        assert!(master.contains("#EXT-X-STREAM-INF:BANDWIDTH=1000,SUBTITLES=\"aniplay-subs\""));
844    }
845
846    #[test]
847    fn subtitle_media_playlist_wraps_the_actual_track() {
848        let playlist = subtitle_playlist("http://127.0.0.1:1234/r/subtitle", 1_421.25);
849        assert!(playlist.contains("#EXT-X-PLAYLIST-TYPE:VOD"));
850        assert!(playlist.contains("#EXT-X-TARGETDURATION:1422"));
851        assert!(playlist.contains("#EXTINF:1421.250,"));
852        assert!(playlist.contains("http://127.0.0.1:1234/r/subtitle"));
853        assert!(playlist.ends_with("#EXT-X-ENDLIST\n"));
854    }
855
856    #[test]
857    fn derives_subtitle_duration_from_webvtt_and_srt_cues() {
858        let webvtt = "WEBVTT\n\n00:00:01.000 --> 00:00:04.500\nHello\n\n00:23:40.250 --> 00:23:43.750\nEnd\n";
859        let srt =
860            "1\n00:00:01,000 --> 00:00:03,000\nHello\n\n2\n01:02:03,250 --> 01:02:05,500\nEnd\n";
861
862        assert_eq!(subtitle_duration(webvtt), Some(1_424.75));
863        assert_eq!(subtitle_duration(srt), Some(3_726.5));
864        assert_eq!(subtitle_duration("WEBVTT\n\nNo cues"), None);
865    }
866
867    #[tokio::test]
868    async fn rewrites_nested_resources_and_unwraps_segments() {
869        let server = MockServer::start().await;
870        Mock::given(method("GET"))
871            .and(path("/master.m3u8"))
872            .respond_with(
873                ResponseTemplate::new(200)
874                    .insert_header("content-type", "application/vnd.apple.mpegurl")
875                    .set_body_string("#EXTM3U\n#EXT-X-STREAM-INF:BANDWIDTH=1\nmedia.m3u8\n"),
876            )
877            .mount(&server)
878            .await;
879        Mock::given(method("GET"))
880            .and(path("/media.m3u8"))
881            .respond_with(ResponseTemplate::new(200).set_body_string(
882                "#EXTM3U\n#EXT-X-KEY:METHOD=AES-128,URI=\"key.bin\"\nsegment.png\n",
883            ))
884            .mount(&server)
885            .await;
886        Mock::given(method("GET"))
887            .and(path("/key.bin"))
888            .respond_with(ResponseTemplate::new(200).set_body_bytes(vec![1, 2, 3]))
889            .mount(&server)
890            .await;
891        let mut wrapped = b"\x89PNG\r\n\x1a\nwrapper".to_vec();
892        let offset = wrapped.len();
893        wrapped.resize(offset + 377, 0);
894        wrapped[offset] = 0x47;
895        wrapped[offset + 188] = 0x47;
896        wrapped[offset + 376] = 0x47;
897        Mock::given(method("GET"))
898            .and(path("/segment.png"))
899            .respond_with(ResponseTemplate::new(200).set_body_bytes(wrapped))
900            .mount(&server)
901            .await;
902        Mock::given(method("HEAD"))
903            .and(path("/segment.png"))
904            .respond_with(
905                ResponseTemplate::new(200)
906                    .insert_header("content-type", "image/png")
907                    .insert_header("content-length", "999"),
908            )
909            .mount(&server)
910            .await;
911
912        let stream = StreamLink {
913            url: format!("{}/master.m3u8", server.uri()),
914            resolution: "Auto".into(),
915            hls: true,
916            provider: "MegaPlay".into(),
917            downloadable: true,
918            headers: RequestHeaders::default(),
919            subtitles: vec![],
920        };
921        let (_relay, local) = relay_stream(&stream).await.unwrap();
922        let client = reqwest::Client::new();
923        let master = client
924            .get(&local.url)
925            .send()
926            .await
927            .unwrap()
928            .text()
929            .await
930            .unwrap();
931        let media_url = master
932            .lines()
933            .find(|line| line.starts_with("http://"))
934            .unwrap();
935        let media = client
936            .get(media_url)
937            .send()
938            .await
939            .unwrap()
940            .text()
941            .await
942            .unwrap();
943        let segment_url = media
944            .lines()
945            .find(|line| !line.is_empty() && !line.starts_with('#'))
946            .unwrap();
947        let segment_head = client.head(segment_url).send().await.unwrap();
948        assert_eq!(
949            segment_head.headers().get("content-type").unwrap(),
950            "video/mp2t"
951        );
952        let segment = client.get(segment_url).send().await.unwrap();
953        assert_eq!(segment.headers().get("content-type").unwrap(), "video/mp2t");
954        assert_eq!(segment.bytes().await.unwrap()[0], 0x47);
955        assert!(media.contains("URI=\"http://127.0.0.1:"));
956    }
957
958    #[tokio::test]
959    async fn wraps_media_playlists_with_external_subtitle_renditions() {
960        let server = MockServer::start().await;
961        Mock::given(method("GET"))
962            .and(path("/media.m3u8"))
963            .respond_with(
964                ResponseTemplate::new(200)
965                    .insert_header("content-type", "application/vnd.apple.mpegurl")
966                    .set_body_string("#EXTM3U\n#EXTINF:10,\nsegment.ts\n"),
967            )
968            .mount(&server)
969            .await;
970        Mock::given(method("GET"))
971            .and(path("/subtitles.vtt"))
972            .respond_with(
973                ResponseTemplate::new(200)
974                    .insert_header("content-type", "text/vtt")
975                    .set_body_string("WEBVTT\n"),
976            )
977            .mount(&server)
978            .await;
979
980        let stream = StreamLink {
981            url: format!("{}/media.m3u8", server.uri()),
982            resolution: "Auto".into(),
983            hls: true,
984            provider: "MegaPlay".into(),
985            downloadable: true,
986            headers: RequestHeaders::default(),
987            subtitles: vec![SubtitleTrack {
988                label: "English".into(),
989                url: format!("{}/subtitles.vtt", server.uri()),
990                default: true,
991            }],
992        };
993        let (_relay, local) = relay_stream(&stream).await.unwrap();
994        let client = reqwest::Client::new();
995        let master = client
996            .get(&local.url)
997            .send()
998            .await
999            .unwrap()
1000            .text()
1001            .await
1002            .unwrap();
1003
1004        assert!(master.contains("#EXT-X-MEDIA:TYPE=SUBTITLES"));
1005        let subtitle_playlist_url = master
1006            .lines()
1007            .find(|line| line.starts_with("#EXT-X-MEDIA:TYPE=SUBTITLES"))
1008            .and_then(|line| line.split("URI=\"").nth(1))
1009            .and_then(|value| value.split('"').next())
1010            .unwrap();
1011        let subtitle_playlist = client
1012            .get(subtitle_playlist_url)
1013            .send()
1014            .await
1015            .unwrap()
1016            .text()
1017            .await
1018            .unwrap();
1019        assert!(subtitle_playlist.contains(&local.subtitles[0].url));
1020        let media_url = master
1021            .lines()
1022            .find(|line| line.starts_with("http://"))
1023            .unwrap();
1024        let media = client
1025            .get(media_url)
1026            .send()
1027            .await
1028            .unwrap()
1029            .text()
1030            .await
1031            .unwrap();
1032        assert!(media.contains("#EXTINF:10,"));
1033        let subtitle = client.get(&local.subtitles[0].url).send().await.unwrap();
1034        assert_eq!(
1035            subtitle
1036                .headers()
1037                .get(header::CONTENT_TYPE)
1038                .unwrap()
1039                .to_str()
1040                .unwrap(),
1041            "text/vtt"
1042        );
1043    }
1044
1045    #[tokio::test]
1046    async fn desktop_relay_omits_hls_subtitle_renditions_but_keeps_the_sidecar() {
1047        let server = MockServer::start().await;
1048        Mock::given(method("GET"))
1049            .and(path("/media.m3u8"))
1050            .respond_with(
1051                ResponseTemplate::new(200)
1052                    .insert_header("content-type", "application/vnd.apple.mpegurl")
1053                    .set_body_string("#EXTM3U\n#EXTINF:10,\nsegment.ts\n"),
1054            )
1055            .mount(&server)
1056            .await;
1057        Mock::given(method("GET"))
1058            .and(path("/subtitles.vtt"))
1059            .respond_with(
1060                ResponseTemplate::new(200)
1061                    .insert_header("content-type", "text/vtt")
1062                    .set_body_string("WEBVTT\n\n00:00:01.000 --> 00:00:04.500\nHello\n"),
1063            )
1064            .mount(&server)
1065            .await;
1066
1067        let stream = StreamLink {
1068            url: format!("{}/media.m3u8", server.uri()),
1069            resolution: "Auto".into(),
1070            hls: true,
1071            provider: "MegaPlay".into(),
1072            downloadable: true,
1073            headers: RequestHeaders::default(),
1074            subtitles: vec![SubtitleTrack {
1075                label: "English".into(),
1076                url: format!("{}/subtitles.vtt", server.uri()),
1077                default: true,
1078            }],
1079        };
1080        let (_relay, local) = relay_stream_without_hls_subtitles(&stream).await.unwrap();
1081        let client = reqwest::Client::new();
1082        let master = client
1083            .get(&local.url)
1084            .send()
1085            .await
1086            .unwrap()
1087            .text()
1088            .await
1089            .unwrap();
1090
1091        assert!(!master.contains("#EXT-X-MEDIA:TYPE=SUBTITLES"));
1092        assert!(!master.contains("SUBTITLES=\"aniplay-subs\""));
1093
1094        // The direct sidecar subtitle URL is still relayed and reachable, so
1095        // desktop players can load it via `--sub-file`.
1096        assert_eq!(local.subtitles.len(), 1);
1097        let subtitle = client.get(&local.subtitles[0].url).send().await.unwrap();
1098        assert_eq!(subtitle.status(), reqwest::StatusCode::OK);
1099        assert_eq!(
1100            subtitle
1101                .headers()
1102                .get(header::CONTENT_TYPE)
1103                .unwrap()
1104                .to_str()
1105                .unwrap(),
1106            "text/vtt"
1107        );
1108        let body = subtitle.text().await.unwrap();
1109        assert!(body.contains("Hello"));
1110    }
1111
1112    #[tokio::test]
1113    async fn suppresses_segment_ranges_and_rejects_unknown_tokens() {
1114        let server = MockServer::start().await;
1115        Mock::given(method("GET"))
1116            .and(path("/master.m3u8"))
1117            .respond_with(ResponseTemplate::new(200).set_body_string("#EXTM3U\nsegment.ts\n"))
1118            .mount(&server)
1119            .await;
1120        Mock::given(method("GET"))
1121            .and(path("/segment.ts"))
1122            .respond_with(ResponseTemplate::new(200).set_body_bytes(vec![0x47, 1, 2]))
1123            .expect(1)
1124            .mount(&server)
1125            .await;
1126        let stream = StreamLink {
1127            url: format!("{}/master.m3u8", server.uri()),
1128            resolution: "Auto".into(),
1129            hls: true,
1130            provider: "MegaPlay".into(),
1131            downloadable: true,
1132            headers: RequestHeaders::default(),
1133            subtitles: vec![],
1134        };
1135        let (_relay, local) = relay_stream(&stream).await.unwrap();
1136        let client = reqwest::Client::new();
1137        let playlist = client
1138            .get(&local.url)
1139            .send()
1140            .await
1141            .unwrap()
1142            .text()
1143            .await
1144            .unwrap();
1145        let segment_url = playlist
1146            .lines()
1147            .find(|line| line.starts_with("http://"))
1148            .unwrap();
1149        let ranged = client
1150            .get(segment_url)
1151            .header("range", "bytes=0-2")
1152            .send()
1153            .await
1154            .unwrap();
1155        assert_eq!(ranged.status(), reqwest::StatusCode::OK);
1156        let requests = server.received_requests().await.unwrap();
1157        let segment_request = requests
1158            .iter()
1159            .find(|request| request.url.path() == "/segment.ts")
1160            .unwrap();
1161        assert!(!segment_request.headers.contains_key("range"));
1162        let base = Url::parse(&local.url).unwrap();
1163        let rejected = client
1164            .get(base.join("/r/unknown").unwrap())
1165            .send()
1166            .await
1167            .unwrap();
1168        assert_eq!(rejected.status(), reqwest::StatusCode::FORBIDDEN);
1169    }
1170
1171    #[test]
1172    fn rejects_non_https_remote_resources() {
1173        let error = validate_upstream("http://example.com/master.m3u8").unwrap_err();
1174        assert!(matches!(
1175            error,
1176            AniError::ProviderUrlValidationFailed { .. }
1177        ));
1178    }
1179
1180    #[test]
1181    fn caps_registered_playlist_resources() {
1182        let state = Arc::new(State {
1183            client: Client::new(),
1184            base: "127.0.0.1:1".parse().unwrap(),
1185            resources: Mutex::new(HashMap::new()),
1186            tokens_by_url: Mutex::new(HashMap::new()),
1187            counter: AtomicUsize::new(0),
1188            secret: "test".into(),
1189        });
1190        for index in 0..MAX_RESOURCES {
1191            register(
1192                &state,
1193                Url::parse(&format!("https://kotocdn.site/{index}")).unwrap(),
1194                RequestHeaders::default(),
1195                ResourceKind::Segment,
1196            )
1197            .unwrap();
1198        }
1199        let error = register(
1200            &state,
1201            Url::parse("https://kotocdn.site/overflow").unwrap(),
1202            RequestHeaders::default(),
1203            ResourceKind::Segment,
1204        )
1205        .unwrap_err();
1206        assert!(matches!(error, AniError::ProviderInvalidResponse { .. }));
1207    }
1208}
ani_lib/hls_relay.rs

ani_lib/
hls_relay.rs
