ani_lib/
hls_relay.rs

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

ani_lib/
hls_relay.rs
