From c880d53bb5aa4b835e9ea0a32dff39f601b830ed Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Thu, 17 Apr 2025 23:07:14 +0400
Subject: [PATCH] feat: allow cancelation of download of a different URL

---
 Tests.Vpn.Service/DownloaderTest.cs | 26 ++++++++++++++++++++++++++
 Vpn.Service/Downloader.cs           | 18 +++++++++++++++++-
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/Tests.Vpn.Service/DownloaderTest.cs b/Tests.Vpn.Service/DownloaderTest.cs
index 8b55f50..b30e3e4 100644
--- a/Tests.Vpn.Service/DownloaderTest.cs
+++ b/Tests.Vpn.Service/DownloaderTest.cs
@@ -412,6 +412,32 @@ public async Task MismatchedETag(CancellationToken ct)
         Assert.That(ex.Message, Does.Contain("ETag does not match SHA1 hash of downloaded file").And.Contains("beef"));
     }
 
+    [Test(Description = "Timeout waiting for existing download")]
+    [CancelAfter(30_000)]
+    public async Task CancelledWaitingForOther(CancellationToken ct)
+    {
+        var testCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
+        using var httpServer = new TestHttpServer(async _ =>
+        {
+            await Task.Delay(TimeSpan.FromSeconds(5), testCts.Token);
+        });
+        var url0 = new Uri(httpServer.BaseUrl + "/test0");
+        var url1 = new Uri(httpServer.BaseUrl + "/test1");
+        var destPath = Path.Combine(_tempDir, "test");
+        var manager = new Downloader(NullLogger<Downloader>.Instance);
+
+        // first outer task succeeds, getting download started
+        var dlTask0 = await manager.StartDownloadAsync(new HttpRequestMessage(HttpMethod.Get, url0), destPath,
+            NullDownloadValidator.Instance, testCts.Token);
+
+        // The second request fails if the timeout is short
+        var smallerCt = new CancellationTokenSource(TimeSpan.FromSeconds(1)).Token;
+        Assert.ThrowsAsync<TaskCanceledException>(async () => await manager.StartDownloadAsync(
+            new HttpRequestMessage(HttpMethod.Get, url1), destPath,
+            NullDownloadValidator.Instance, smallerCt));
+        await testCts.CancelAsync();
+    }
+
     [Test(Description = "Timeout on response body")]
     [CancelAfter(30_000)]
     public async Task CancelledInner(CancellationToken ct)
diff --git a/Vpn.Service/Downloader.cs b/Vpn.Service/Downloader.cs
index 6a665ae..467c9af 100644
--- a/Vpn.Service/Downloader.cs
+++ b/Vpn.Service/Downloader.cs
@@ -287,6 +287,7 @@ public async Task<DownloadTask> StartDownloadAsync(HttpRequestMessage req, strin
     {
         while (true)
         {
+            ct.ThrowIfCancellationRequested();
             var task = _downloads.GetOrAdd(destinationPath,
                 _ => new DownloadTask(_logger, req, destinationPath, validator));
             // EnsureStarted is a no-op if we didn't create a new DownloadTask.
@@ -322,7 +323,22 @@ public async Task<DownloadTask> StartDownloadAsync(HttpRequestMessage req, strin
             _logger.LogWarning(
                 "Download for '{DestinationPath}' is already in progress, but is for a different Url - awaiting completion",
                 destinationPath);
-            await task.Task;
+            await TaskOrCancellation(task.Task, ct);
+        }
+    }
+
+    /// <summary>
+    /// TaskOrCancellation waits for either the task to complete, or the given token to be canceled.
+    /// </summary>
+    internal static async Task TaskOrCancellation(Task task, CancellationToken cancellationToken)
+    {
+        var cancellationTask = new TaskCompletionSource();
+        await using (cancellationToken.Register(() => cancellationTask.TrySetCanceled()))
+        {
+            // Wait for either the task or the cancellation
+            var completedTask = await Task.WhenAny(task, cancellationTask.Task);
+            // Await to propagate exceptions, if any
+            await completedTask;
         }
     }
 }