From 80a8a1020caad562578af3899510ad397d0ce804 Mon Sep 17 00:00:00 2001 From: MentalBlank Date: Tue, 20 May 2025 03:27:41 +1000 Subject: [PATCH] Expand enrichment to torrent files --- .../Models/Internal/DbSettings.cs | 4 +- .../Services/EnricherTest.cs | 134 ++++++++++++++++++ .../Services/TorrentsTest.cs | 12 +- server/RdtClient.Service/DiConfig.cs | 1 + server/RdtClient.Service/Services/Enricher.cs | 111 +++++++++++++++ server/RdtClient.Service/Services/Torrents.cs | 49 ++----- .../Services/TrackerListGrabber.cs | 8 +- 7 files changed, 266 insertions(+), 53 deletions(-) create mode 100644 server/RdtClient.Service.Test/Services/EnricherTest.cs create mode 100644 server/RdtClient.Service/Services/Enricher.cs diff --git a/server/RdtClient.Data/Models/Internal/DbSettings.cs b/server/RdtClient.Data/Models/Internal/DbSettings.cs index 0d6e693..24db43f 100644 --- a/server/RdtClient.Data/Models/Internal/DbSettings.cs +++ b/server/RdtClient.Data/Models/Internal/DbSettings.cs @@ -75,8 +75,8 @@ Supports the following parameters: [Description("When a torrent file or magnet is added, create a copy in this directory.")] public String? CopyAddedTorrents { get; set; } = null; - [DisplayName("Magnet enrichment list")] - [Description("Optional. Specify the URL of a tracker list file to be appended to magnet links.")] + [DisplayName("Tracker enrichment list")] + [Description("Optional. Specify the URL of a tracker list file to be appended to magnet links and torrent files.")] public String? MagnetTrackerEnrichment { get; set; } = null; [DisplayName("Disable update notifications")] diff --git a/server/RdtClient.Service.Test/Services/EnricherTest.cs b/server/RdtClient.Service.Test/Services/EnricherTest.cs new file mode 100644 index 0000000..4c7ceb7 --- /dev/null +++ b/server/RdtClient.Service.Test/Services/EnricherTest.cs @@ -0,0 +1,134 @@ +using Microsoft.Extensions.Logging; +using Moq; +using RdtClient.Service.Services; +using MonoTorrent.BEncoding; + +namespace RdtClient.Service.Test.Services; + +public class EnricherTest : IDisposable +{ + private readonly MockRepository _mockRepository; + private readonly Mock> _loggerMock; + private readonly Mock _trackerListGrabberMock; + + public EnricherTest() + { + _mockRepository = new MockRepository(MockBehavior.Strict); + _loggerMock = _mockRepository.Create>(MockBehavior.Loose); + _trackerListGrabberMock = _mockRepository.Create(); + } + + public void Dispose() + { + _mockRepository.VerifyAll(); + } + + private const String TestMagnetLink = "magnet:?xt=urn:btih:1234567890123456789012345678901234567890&dn=TestFile&tr=http%3A%2F%2Ftracker1.com%2Fannounce&tr=http%3A%2F%2Ftracker2.com%2Fannounce"; + + [Fact] + public async Task EnrichMagnetLink_AddsNoTrackers_WhenNoTrackersFromTrackerGrabber() + { + // Arrange + SetupTrackerListGrabber([]); + + var enricher = new Enricher(_loggerMock.Object, _trackerListGrabberMock.Object); + + // Act + var enriched = await enricher.EnrichMagnetLink(TestMagnetLink); + + // Assert + Assert.Equal(TestMagnetLink, enriched); + } + + [Fact] + public async Task EnrichMagnetLink_AddsTrackers_WhenTrackersFromTrackerGrabber() + { + // Arrange + SetupTrackerListGrabber(["http://new-tracker.com/announce"]); + + var Enricher = new Enricher(_loggerMock.Object, _trackerListGrabberMock.Object); + + // Act + var enriched = await Enricher.EnrichMagnetLink(TestMagnetLink); + + // Assert + Assert.Equal(TestMagnetLink + $"&tr={Uri.EscapeDataString("http://new-tracker.com/announce")}", enriched); + } + + [Fact] + public async Task EnrichMagnetLink_DoesNotAddDuplicateTrackers_WhenTrackersFromTrackerGrabberAlreadyPresent() + { + // Arrange + SetupTrackerListGrabber(["http://new-tracker.com/announce", "http://tracker1.com/announce"]); + + var enricher = new Enricher(_loggerMock.Object, _trackerListGrabberMock.Object); + + // Act + var enriched = await enricher.EnrichMagnetLink(TestMagnetLink); + + // Assert + Assert.Equal(TestMagnetLink + $"&tr={Uri.EscapeDataString("http://new-tracker.com/announce")}", enriched); + } + + [Fact] + public async Task EnrichMagnetLink_Throws_WhenTrackerGrabberThrows() + { + // Arrange + _trackerListGrabberMock + .Setup(t => t.GetTrackers()) + .ThrowsAsync(new InvalidOperationException("Unable to fetch tracker list for enrichment.")); + + var enricher = new Enricher(_loggerMock.Object, _trackerListGrabberMock.Object); + + // Act & Assert + await Assert.ThrowsAsync(() => enricher.EnrichMagnetLink(TestMagnetLink)); + } + + + [Fact] + public async Task EnrichTorrentBytes_AddsTrackers_WhenTrackersFromTrackerGrabber() + { + // Arrange + var originalTracker = "http://tracker1.com/announce"; + var newTracker = "http://new-tracker.com/announce"; + + var torrentDict = new BEncodedDictionary + { + ["announce"] = new BEncodedString(originalTracker), + ["announce-list"] = new BEncodedList + { + new BEncodedList { new BEncodedString(originalTracker) } + } + }; + + var originalTorrentBytes = torrentDict.Encode(); + + SetupTrackerListGrabber(new[] { newTracker }); + var enricher = new Enricher(_loggerMock.Object, _trackerListGrabberMock.Object); + + // Act + var enrichedBytes = await enricher.EnrichTorrentBytes(originalTorrentBytes); + var enrichedDict = BEncodedValue.Decode(enrichedBytes); + + // Assert + Assert.True(enrichedDict.ContainsKey("announce")); + Assert.True(enrichedDict.ContainsKey("announce-list")); + + var announceList = (BEncodedList)enrichedDict["announce-list"]; + var flattened = announceList.Cast().SelectMany(l => l.Cast().Select(s => s.Text)).ToList(); + + Assert.Contains(originalTracker, flattened); + Assert.Contains(newTracker, flattened); + + var announce = ((BEncodedString)enrichedDict["announce"]).Text; + Assert.Equal(flattened.First(), announce); + } + + private void SetupTrackerListGrabber(String[] trackerList) + { + _trackerListGrabberMock + .Setup(t => t.GetTrackers()) + .ReturnsAsync(trackerList) + .Verifiable(); + } +} \ No newline at end of file diff --git a/server/RdtClient.Service.Test/Services/TorrentsTest.cs b/server/RdtClient.Service.Test/Services/TorrentsTest.cs index 90fe94d..054a62f 100644 --- a/server/RdtClient.Service.Test/Services/TorrentsTest.cs +++ b/server/RdtClient.Service.Test/Services/TorrentsTest.cs @@ -20,13 +20,13 @@ class Mocks public readonly Mock> TorrentsLoggerMock; public readonly Mock DownloadsMock; public readonly Mock TorrentDataMock; - public readonly Mock TrackerListGrabberMock; + public readonly Mock EnricherMock; public Mocks() { TorrentDataMock = new(); DownloadsMock = new(); - TrackerListGrabberMock = new(); + EnricherMock = new(); TorrentsLoggerMock = new(); @@ -104,7 +104,7 @@ public class TorrentsTest mocks.DownloadsMock.Object, mocks.ProcessFactoryMock.Object, fileSystemMock, - mocks.TrackerListGrabberMock.Object, + mocks.EnricherMock.Object, null!, // Torrent Clients are not used by `RunTorrentComplete`, this is fine null!, null!, @@ -170,7 +170,7 @@ public class TorrentsTest mocks.DownloadsMock.Object, mocks.ProcessFactoryMock.Object, fileSystemMock, - mocks.TrackerListGrabberMock.Object, + mocks.EnricherMock.Object, null!, // Torrent Clients are not used by `RunTorrentComplete`, this is fine null!, null!, @@ -218,7 +218,7 @@ public class TorrentsTest mocks.DownloadsMock.Object, mocks.ProcessFactoryMock.Object, fileSystemMock, - mocks.TrackerListGrabberMock.Object, + mocks.EnricherMock.Object, null!, // Torrent Clients are not used by `RunTorrentComplete`, this is fine null!, null!, @@ -285,7 +285,7 @@ public class TorrentsTest mocks.DownloadsMock.Object, mocks.ProcessFactoryMock.Object, fileSystemMock, - mocks.TrackerListGrabberMock.Object, + mocks.EnricherMock.Object, null!, // Torrent Clients are not used by `RunTorrentComplete`, this is fine null!, null!, diff --git a/server/RdtClient.Service/DiConfig.cs b/server/RdtClient.Service/DiConfig.cs index a62986d..fe3edbf 100644 --- a/server/RdtClient.Service/DiConfig.cs +++ b/server/RdtClient.Service/DiConfig.cs @@ -39,6 +39,7 @@ public static class DiConfig services.AddSingleton(); services.AddSingleton(); + services.AddSingleton(); services.AddSingleton(); diff --git a/server/RdtClient.Service/Services/Enricher.cs b/server/RdtClient.Service/Services/Enricher.cs new file mode 100644 index 0000000..140d07a --- /dev/null +++ b/server/RdtClient.Service/Services/Enricher.cs @@ -0,0 +1,111 @@ +using System.Web; +using Microsoft.Extensions.Logging; +using MonoTorrent.BEncoding; + +namespace RdtClient.Service.Services; + +public interface IEnricher +{ + Task EnrichMagnetLink(String magnetLink); + Task EnrichTorrentBytes(Byte[] torrentBytes); +} + +/// +/// Enriches magnet links and torrents by adding trackers from the tracker list grabber. +/// +public class Enricher(ILogger logger, ITrackerListGrabber trackerListGrabber) : IEnricher +{ + /// + /// Add trackers from the tracker list grabber to the magnet link. + /// + /// Magnet link to add trackres to. Is not modified + /// Magnet link with additional trackers + public async Task EnrichMagnetLink(String magnetLink) + { + var newTrackers = await trackerListGrabber.GetTrackers(); + + var uri = new Uri(magnetLink); + var query = HttpUtility.ParseQueryString(uri.Query); + + var existingTrackers = query.GetValues("tr") ?? []; + var allTrackers = existingTrackers.Concat(newTrackers).Distinct(StringComparer.OrdinalIgnoreCase); + + var trackerQuery = String.Join("&tr=", allTrackers.Select(Uri.EscapeDataString)); + + if (!String.IsNullOrEmpty(trackerQuery)) + { + trackerQuery = "&tr=" + trackerQuery; + } + + var baseWithoutTrackers = magnetLink.Split("&tr=")[0]; + + var separator = baseWithoutTrackers.Contains('?') ? "&" : "?"; + + var newUri = baseWithoutTrackers + separator + trackerQuery.TrimStart('&'); + try + { + _ = new Uri(newUri); + } + catch (UriFormatException ex) + { + logger.LogWarning(ex, "Failed to enrich magnet link: {newUri}", newUri); + throw new InvalidOperationException($"Failed to enrich magnet link: {newUri}", ex); + } + return newUri; + } + + /// + /// Add trackers from the tracker list grabber to the .torrent file bytes. + /// + /// Torrent file bytes to add trackers to. Is not modified + /// Torrent file bytes with additional trackers + public async Task EnrichTorrentBytes(Byte[] torrentBytes) + { + var newTrackers = await trackerListGrabber.GetTrackers(); + + if (torrentBytes == null) throw new ArgumentNullException(nameof(torrentBytes)); + var torrentDict = BEncodedValue.Decode(torrentBytes); + + + if (!torrentDict.TryGetValue("announce-list", out var announceListValue) || announceListValue is not BEncodedList announceList) + { + announceList = new BEncodedList(); + if (torrentDict.TryGetValue("announce", out var announceValue) && announceValue is BEncodedString announceStr) + { + announceList.Add(new BEncodedList { announceStr }); + } + } + + var existingTrackers = new HashSet(StringComparer.OrdinalIgnoreCase); + foreach (var tier in announceList) + { + if (tier is BEncodedList tierList) + { + foreach (var tracker in tierList) + { + if (tracker is BEncodedString trackerStr) + { + existingTrackers.Add(trackerStr.Text); + } + } + } + } + + foreach (var tracker in newTrackers) + { + if (!existingTrackers.Contains(tracker)) + { + announceList.Add(new BEncodedList { new BEncodedString(tracker) }); + existingTrackers.Add(tracker); + } + } + + torrentDict["announce-list"] = announceList; + if (announceList.Count > 0 && announceList[0] is BEncodedList firstTier && firstTier.Count > 0 && firstTier[0] is BEncodedString firstTracker) + { + torrentDict["announce"] = firstTracker; + } + + return torrentDict.Encode(); + } +} \ No newline at end of file diff --git a/server/RdtClient.Service/Services/Torrents.cs b/server/RdtClient.Service/Services/Torrents.cs index aa59d59..3ff2e69 100644 --- a/server/RdtClient.Service/Services/Torrents.cs +++ b/server/RdtClient.Service/Services/Torrents.cs @@ -25,7 +25,7 @@ public class Torrents( IDownloads downloads, IProcessFactory processFactory, IFileSystem fileSystem, - ITrackerListGrabber trackerListGrabber, + IEnricher enricher, AllDebridTorrentClient allDebridTorrentClient, PremiumizeTorrentClient premiumizeTorrentClient, RealDebridTorrentClient realDebridTorrentClient, @@ -33,7 +33,7 @@ public class Torrents( TorBoxTorrentClient torBoxTorrentClient) { private static readonly SemaphoreSlim RealDebridUpdateLock = new(1, 1); - private readonly ITrackerListGrabber _trackerListGrabber = trackerListGrabber; + private static readonly JsonSerializerOptions JsonSerializerOptions = new() { ReferenceHandler = ReferenceHandler.IgnoreCycles @@ -109,45 +109,9 @@ public class Torrents( await torrentData.UpdateCategory(torrent.TorrentId, category); } - private async Task EnrichMagnetLink(String magnetLink) - { - if (String.IsNullOrWhiteSpace(Settings.Get.General.MagnetTrackerEnrichment)) - { - return magnetLink; - } - - try - { - var newTrackers = await _trackerListGrabber.GetTrackers(); - - var uri = new Uri(magnetLink); - var query = HttpUtility.ParseQueryString(uri.Query); - var existingTrackers = query.GetValues("tr") ?? []; - var allTrackers = existingTrackers.Concat(newTrackers).Distinct(StringComparer.OrdinalIgnoreCase); - - var trackerQuery = String.Join("&tr=", allTrackers.Select(Uri.EscapeDataString)); - - if (!String.IsNullOrEmpty(trackerQuery)) - { - trackerQuery = "&tr=" + trackerQuery; - } - - var baseWithoutTrackers = magnetLink.Split(["&tr="], StringSplitOptions.None)[0]; - var separator = baseWithoutTrackers.Contains('?') ? "&" : "?"; - - return baseWithoutTrackers + separator + trackerQuery.TrimStart('&'); - } - catch (Exception ex) - { - logger.LogError(ex, "{Message}, trying to enrich {Magnet}", ex.Message, magnetLink); - - return magnetLink; - } - } - public async Task AddMagnetToDebridQueue(String magnetLink, Torrent torrent) { - var enriched = await EnrichMagnetLink(magnetLink); + var enriched = await enricher.EnrichMagnetLink(magnetLink); MagnetLink magnet; try { @@ -166,7 +130,7 @@ public class Torrents( var newTorrent = await AddQueued(hash, enriched, false, torrent); Log($"Adding {hash} (magnet link) to queue", newTorrent); - await CopyAddedTorrent(magnet.Name!, enriched); + await CopyAddedTorrent(magnet.Name!, magnetLink); return newTorrent; } @@ -178,9 +142,12 @@ public class Torrents( var fileAsBase64 = Convert.ToBase64String(bytes); logger.LogDebug($"bytes {bytes}"); + var enriched = await enricher.EnrichTorrentBytes(bytes); + logger.LogDebug($"enriched bytes {enriched}"); + try { - monoTorrent = await MonoTorrent.Torrent.LoadAsync(bytes); + monoTorrent = await MonoTorrent.Torrent.LoadAsync(enriched); } catch (Exception ex) { diff --git a/server/RdtClient.Service/Services/TrackerListGrabber.cs b/server/RdtClient.Service/Services/TrackerListGrabber.cs index 5e034a9..d82fb07 100644 --- a/server/RdtClient.Service/Services/TrackerListGrabber.cs +++ b/server/RdtClient.Service/Services/TrackerListGrabber.cs @@ -12,9 +12,9 @@ public class TrackerListGrabber : ITrackerListGrabber public async Task GetTrackers() { var trackerUrlList = Settings.Get.General.MagnetTrackerEnrichment; - if (string.IsNullOrWhiteSpace(trackerUrlList)) + if (String.IsNullOrWhiteSpace(trackerUrlList)) { - return Array.Empty(); + return []; } try @@ -26,9 +26,9 @@ public class TrackerListGrabber : ITrackerListGrabber .Distinct(StringComparer.OrdinalIgnoreCase) .ToArray(); } - catch + catch (Exception ex) { - return Array.Empty(); + throw new InvalidOperationException("Unable to fetch tracker list for enrichment.", ex); } } } \ No newline at end of file