jellyfin-srfPlay/Jellyfin.Plugin.SRFPlay/Services/StreamProxyService.cs

using System;
using System.Collections.Concurrent;
using System.Net.Http;
using System.Text.RegularExpressions;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;

namespace Jellyfin.Plugin.SRFPlay.Services;

/// <summary>
/// Service for proxying SRF Play streams and managing authentication.
/// </summary>
public class StreamProxyService : IDisposable
{
    private readonly ILogger<StreamProxyService> _logger;
    private readonly StreamUrlResolver _streamResolver;
    private readonly HttpClient _httpClient;
    private readonly ConcurrentDictionary<string, StreamInfo> _streamMappings;
    private bool _disposed;

    /// <summary>
    /// Initializes a new instance of the <see cref="StreamProxyService"/> class.
    /// </summary>
    /// <param name="logger">The logger.</param>
    /// <param name="streamResolver">The stream URL resolver.</param>
    public StreamProxyService(ILogger<StreamProxyService> logger, StreamUrlResolver streamResolver)
    {
        _logger = logger;
        _streamResolver = streamResolver;
        _httpClient = new HttpClient
        {
            Timeout = TimeSpan.FromSeconds(30)
        };
        _streamMappings = new ConcurrentDictionary<string, StreamInfo>();
    }

    /// <summary>
    /// Registers a stream for proxying.
    /// </summary>
    /// <param name="itemId">The item ID.</param>
    /// <param name="authenticatedUrl">The authenticated stream URL.</param>
    public void RegisterStream(string itemId, string authenticatedUrl)
    {
        var streamInfo = new StreamInfo
        {
            AuthenticatedUrl = authenticatedUrl,
            RegisteredAt = DateTime.UtcNow
        };

        _streamMappings.AddOrUpdate(itemId, streamInfo, (key, old) => streamInfo);
        _logger.LogInformation("Registered stream for item {ItemId}: {Url}", itemId, authenticatedUrl);
    }

    /// <summary>
    /// Gets the authenticated URL for an item.
    /// </summary>
    /// <param name="itemId">The item ID.</param>
    /// <returns>The authenticated URL, or null if not found.</returns>
    public string? GetAuthenticatedUrl(string itemId)
    {
        if (_streamMappings.TryGetValue(itemId, out var streamInfo))
        {
            return streamInfo.AuthenticatedUrl;
        }

        _logger.LogWarning("No stream mapping found for item {ItemId}", itemId);
        return null;
    }

    /// <summary>
    /// Fetches and rewrites an HLS manifest to use proxy URLs.
    /// </summary>
    /// <param name="itemId">The item ID.</param>
    /// <param name="baseProxyUrl">The base proxy URL (e.g., https://jellyfin-server/Plugins/SRFPlay/Proxy/{itemId}).</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>The rewritten manifest content.</returns>
    public async Task<string?> GetRewrittenManifestAsync(
        string itemId,
        string baseProxyUrl,
        CancellationToken cancellationToken = default)
    {
        var authenticatedUrl = GetAuthenticatedUrl(itemId);
        if (authenticatedUrl == null)
        {
            return null;
        }

        try
        {
            _logger.LogDebug("Fetching manifest from: {Url}", authenticatedUrl);
            var manifestContent = await _httpClient.GetStringAsync(authenticatedUrl, cancellationToken).ConfigureAwait(false);

            // Rewrite the manifest to replace Akamai URLs with proxy URLs
            var rewrittenContent = RewriteManifestUrls(manifestContent, authenticatedUrl, baseProxyUrl);

            _logger.LogDebug("Successfully rewrote manifest for item {ItemId}", itemId);
            return rewrittenContent;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Failed to fetch manifest for item {ItemId} from {Url}", itemId, authenticatedUrl);
            return null;
        }
    }

    /// <summary>
    /// Fetches a segment from the original source.
    /// </summary>
    /// <param name="itemId">The item ID.</param>
    /// <param name="segmentPath">The segment path.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>The segment content as bytes.</returns>
    public async Task<byte[]?> GetSegmentAsync(
        string itemId,
        string segmentPath,
        CancellationToken cancellationToken = default)
    {
        var authenticatedUrl = GetAuthenticatedUrl(itemId);
        if (authenticatedUrl == null)
        {
            return null;
        }

        try
        {
            // Build the full segment URL by combining the base URL with the segment path
            var baseUri = new Uri(authenticatedUrl);
            var baseUrl = $"{baseUri.Scheme}://{baseUri.Host}{string.Join('/', baseUri.AbsolutePath.Split('/')[..^1])}";

            // Extract query parameters (auth tokens) from authenticated URL
            var queryParams = baseUri.Query;

            // Build full segment URL
            var segmentUrl = $"{baseUrl}/{segmentPath}{queryParams}";

            _logger.LogDebug("Fetching segment: {SegmentUrl}", segmentUrl);
            var segmentData = await _httpClient.GetByteArrayAsync(segmentUrl, cancellationToken).ConfigureAwait(false);

            _logger.LogDebug("Successfully fetched segment {SegmentPath} ({Size} bytes)", segmentPath, segmentData.Length);
            return segmentData;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Failed to fetch segment {SegmentPath} for item {ItemId}", segmentPath, itemId);
            return null;
        }
    }

    /// <summary>
    /// Rewrites URLs in HLS manifest to point to proxy.
    /// </summary>
    /// <param name="manifestContent">The original manifest content.</param>
    /// <param name="originalBaseUrl">The original base URL.</param>
    /// <param name="proxyBaseUrl">The proxy base URL.</param>
    /// <returns>The rewritten manifest.</returns>
    private string RewriteManifestUrls(string manifestContent, string originalBaseUrl, string proxyBaseUrl)
    {
        var baseUri = new Uri(originalBaseUrl);
        var baseUrl = $"{baseUri.Scheme}://{baseUri.Host}{string.Join('/', baseUri.AbsolutePath.Split('/')[..^1])}";

        // Pattern to match .m3u8 and .ts/.mp4 segment references
        var pattern = @"(?:^|\n)([^#\n][^\n]*\.(?:m3u8|ts|mp4|m4s|aac)[^\n]*)";

        var rewritten = Regex.Replace(manifestContent, pattern, match =>
        {
            var url = match.Groups[1].Value.Trim();

            // Skip if it's already an absolute URL
            if (url.StartsWith("http://", StringComparison.OrdinalIgnoreCase) ||
                url.StartsWith("https://", StringComparison.OrdinalIgnoreCase))
            {
                // Rewrite absolute URLs to proxy
                var relativePath = url.Replace(baseUrl + "/", string.Empty, StringComparison.Ordinal);
                return $"\n{proxyBaseUrl}/{relativePath}";
            }

            // Relative URL - rewrite to proxy
            return $"\n{proxyBaseUrl}/{url}";
        });

        return rewritten;
    }

    /// <summary>
    /// Cleans up old stream mappings.
    /// </summary>
    public void CleanupOldMappings()
    {
        var cutoff = DateTime.UtcNow.AddHours(-24);
        var keysToRemove = new System.Collections.Generic.List<string>();

        foreach (var kvp in _streamMappings)
        {
            if (kvp.Value.RegisteredAt < cutoff)
            {
                keysToRemove.Add(kvp.Key);
            }
        }

        foreach (var key in keysToRemove)
        {
            _streamMappings.TryRemove(key, out _);
            _logger.LogDebug("Removed old stream mapping for item {ItemId}", key);
        }

        if (keysToRemove.Count > 0)
        {
            _logger.LogInformation("Cleaned up {Count} old stream mappings", keysToRemove.Count);
        }
    }

    /// <summary>
    /// Disposes the service.
    /// </summary>
    public void Dispose()
    {
        Dispose(true);
        GC.SuppressFinalize(this);
    }

    /// <summary>
    /// Disposes the service.
    /// </summary>
    /// <param name="disposing">True if disposing.</param>
    protected virtual void Dispose(bool disposing)
    {
        if (_disposed)
        {
            return;
        }

        if (disposing)
        {
            _httpClient?.Dispose();
        }

        _disposed = true;
    }

    /// <summary>
    /// Stream information.
    /// </summary>
    private sealed class StreamInfo
    {
        public string AuthenticatedUrl { get; set; } = string.Empty;

        public DateTime RegisteredAt { get; set; }
    }
}