Merge pull request #119 from /issues/111-CancellationToken

tpill90 · May 7, 2024 · a715d1c · a715d1c
2 parents ccb0254 + baa1fe1
commit a715d1c
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/BattleNetPrefill/Web/CdnRequestManager.cs b/BattleNetPrefill/Web/CdnRequestManager.cs
@@ -243,8 +243,9 @@ public async Task<byte[]> GetRequestAsBytesAsync(Request request, ProgressTask t
                 requestMessage.Headers.Range = new RangeHeaderValue(startBytes, endBytes);
             }
 
-            using var responseMessage = await _client.SendAsync(requestMessage, HttpCompletionOption.ResponseHeadersRead);
-            await using Stream responseStream = await responseMessage.Content.ReadAsStreamAsync();
+            using var cts = new CancellationTokenSource();
+            using var responseMessage = await _client.SendAsync(requestMessage, HttpCompletionOption.ResponseHeadersRead, cts.Token);
+            await using Stream responseStream = await responseMessage.Content.ReadAsStreamAsync(cts.Token);
 
             responseMessage.EnsureSuccessStatusCode();
             if (writeToDevNull)
@@ -256,7 +257,7 @@ public async Task<byte[]> GetRequestAsBytesAsync(Request request, ProgressTask t
                     while (true)
                     {
                         // Dump the received data, so we don't have to waste time writing it to disk.
-                        var read = await responseStream.ReadAsync(buffer, 0, buffer.Length);
+                        var read = await responseStream.ReadAsync(buffer, cts.Token);
                         if (read == 0)
                         {
                             return null;
@@ -274,7 +275,7 @@ public async Task<byte[]> GetRequestAsBytesAsync(Request request, ProgressTask t
             }
 
             await using var memoryStream = new MemoryStream();
-            await responseStream.CopyToAsync(memoryStream);
+            await responseStream.CopyToAsync(memoryStream, cts.Token);
 
             var byteArray = memoryStream.ToArray();
             if (SkipDiskCache)
@@ -285,7 +286,7 @@ public async Task<byte[]> GetRequestAsBytesAsync(Request request, ProgressTask t
             // Cache to disk
             FileInfo file = new FileInfo(Path.Combine(AppConfig.CacheDir + uri.AbsolutePath));
             file.Directory.Create();
-            await File.WriteAllBytesAsync(file.FullName, byteArray);
+            await File.WriteAllBytesAsync(file.FullName, byteArray, cts.Token);
 
             return await Task.FromResult(byteArray);
         }