handle exceptions and fix UTs

spikecurtis · spikecurtis · commit 35abc6766095 · 2025-04-17T17:13:17.000+04:00
diff --git a/Tests.Vpn.Service/DownloaderTest.cs b/Tests.Vpn.Service/DownloaderTest.cs
@@ -375,17 +375,17 @@ public async Task DownloadExistingDifferentContent(CancellationToken ct)
 
     [Test(Description = "Unexpected response code from server")]
     [CancelAfter(30_000)]
-    public void UnexpectedResponseCode(CancellationToken ct)
+    public async Task UnexpectedResponseCode(CancellationToken ct)
     {
         using var httpServer = new TestHttpServer(ctx => { ctx.Response.StatusCode = 404; });
         var url = new Uri(httpServer.BaseUrl + "/test");
         var destPath = Path.Combine(_tempDir, "test");
 
         var manager = new Downloader(NullLogger<Downloader>.Instance);
-        // The "outer" Task should fail.
-        var ex = Assert.ThrowsAsync<HttpRequestException>(async () =>
-            await manager.StartDownloadAsync(new HttpRequestMessage(HttpMethod.Get, url), destPath,
-                NullDownloadValidator.Instance, ct));
+        // The "inner" Task should fail.
+        var dlTask = await manager.StartDownloadAsync(new HttpRequestMessage(HttpMethod.Get, url), destPath,
+            NullDownloadValidator.Instance, ct);
+        var ex = Assert.ThrowsAsync<HttpRequestException>(async () => await dlTask.Task);
         Assert.That(ex.Message, Does.Contain("404"));
     }
 
@@ -412,22 +412,6 @@ 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 on response headers")]
-    [CancelAfter(30_000)]
-    public void CancelledOuter(CancellationToken ct)
-    {
-        using var httpServer = new TestHttpServer(async _ => { await Task.Delay(TimeSpan.FromSeconds(5), ct); });
-        var url = new Uri(httpServer.BaseUrl + "/test");
-        var destPath = Path.Combine(_tempDir, "test");
-
-        var manager = new Downloader(NullLogger<Downloader>.Instance);
-        // The "outer" Task should fail.
-        var smallerCt = new CancellationTokenSource(TimeSpan.FromSeconds(1)).Token;
-        Assert.ThrowsAsync<TaskCanceledException>(async () => await manager.StartDownloadAsync(
-            new HttpRequestMessage(HttpMethod.Get, url), destPath,
-            NullDownloadValidator.Instance, smallerCt));
-    }
-
     [Test(Description = "Timeout on response body")]
     [CancelAfter(30_000)]
     public async Task CancelledInner(CancellationToken ct)
@@ -479,12 +463,10 @@ public async Task ValidationFailureExistingFile(CancellationToken ct)
         await File.WriteAllTextAsync(destPath, "test", ct);
 
         var manager = new Downloader(NullLogger<Downloader>.Instance);
+        var dlTask = await manager.StartDownloadAsync(new HttpRequestMessage(HttpMethod.Get, url), destPath,
+            new TestDownloadValidator(new Exception("test exception")), ct);
         // The "outer" Task should fail because the inner task never starts.
-        var ex = Assert.ThrowsAsync<Exception>(async () =>
-        {
-            await manager.StartDownloadAsync(new HttpRequestMessage(HttpMethod.Get, url), destPath,
-                new TestDownloadValidator(new Exception("test exception")), ct);
-        });
+        var ex = Assert.ThrowsAsync<Exception>(async () => { await dlTask.Task; });
         Assert.That(ex.Message, Does.Contain("Existing file failed validation"));
         Assert.That(ex.InnerException, Is.Not.Null);
         Assert.That(ex.InnerException!.Message, Is.EqualTo("test exception"));
diff --git a/Vpn.Service/Downloader.cs b/Vpn.Service/Downloader.cs
@@ -3,6 +3,7 @@
 using System.Formats.Asn1;
 using System.Net;
 using System.Runtime.CompilerServices;
+using System.Runtime.ExceptionServices;
 using System.Security.Cryptography;
 using System.Security.Cryptography.X509Certificates;
 using Coder.Desktop.Vpn.Utilities;
@@ -290,7 +291,24 @@ public async Task<DownloadTask> StartDownloadAsync(HttpRequestMessage req, strin
                 _ => new DownloadTask(_logger, req, destinationPath, validator));
             // EnsureStarted is a no-op if we didn't create a new DownloadTask.
             // So, we will only remove the destination once for each time we start a new task.
-            task.EnsureStarted(tsk => _downloads.TryRemove(destinationPath, out _), ct);
+            task.EnsureStarted(tsk =>
+            {
+                // remove the key first, before checking the exception, to ensure
+                // we still clean up.
+                _downloads.TryRemove(destinationPath, out _);
+                if (tsk.Exception == null)
+                {
+                    return;
+                }
+
+                if (tsk.Exception.InnerException != null)
+                {
+                    ExceptionDispatchInfo.Capture(tsk.Exception.InnerException).Throw();
+                }
+
+                // not sure if this is hittable, but just in case:
+                throw tsk.Exception;
+            }, ct);
 
             // If the existing (or new) task is for the same URL, return it.
             if (task.Request.RequestUri == req.RequestUri)
@@ -363,7 +381,7 @@ internal void EnsureStarted(Action<Task> continuation, CancellationToken ct = de
     {
         using var _ = _semaphore.Lock();
         if (Task == null!)
-            Task = Start(ct).ContinueWith(continuation, CancellationToken.None);
+            Task = Start(ct).ContinueWith(continuation, ct);
     }
 
     /// <summary>
