From 768716aedf2bba6dabeb55f2d18f1c7ae50651ff Mon Sep 17 00:00:00 2001
From: GitHubProUser67 <127040195+GitHubProUser67@users.noreply.github.com>
Date: Tue, 30 Apr 2024 18:33:53 +0200
Subject: [PATCH] CRITICAL: Fixes a Watson/Lite crashing issue with Concurrent
 Streams.

CRITICAL: Fixes a Watson/Lite crashing issue with Concurrent Streams.
---
 src/WatsonWebserver.Lite/Webserver.cs | 15 ++++++++++++---
 src/WatsonWebserver/Webserver.cs      | 16 +++++++++++++---
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/src/WatsonWebserver.Lite/Webserver.cs b/src/WatsonWebserver.Lite/Webserver.cs
index d0f53a0..09c5972 100644
--- a/src/WatsonWebserver.Lite/Webserver.cs
+++ b/src/WatsonWebserver.Lite/Webserver.cs
@@ -618,7 +618,10 @@ private async void ClientConnected(object sender, ClientConnectedEventArgs args)
                 {
                     ctx.Response.StatusCode = 404;
                     ctx.Response.ContentType = DefaultPages.Pages[404].ContentType;
-                    await ctx.Response.Send(DefaultPages.Pages[404].Content, _Token).ConfigureAwait(false);
+                    if (ctx.Response.ChunkedTransfer)
+                        await ctx.Response.SendFinalChunk(Encoding.UTF8.GetBytes(DefaultPages.Pages[404].Content), _Token).ConfigureAwait(false);
+                    else
+                        await ctx.Response.Send(DefaultPages.Pages[404].Content, _Token).ConfigureAwait(false);
                     return;
                 }
 
@@ -633,7 +636,10 @@ private async void ClientConnected(object sender, ClientConnectedEventArgs args)
 
                     try
                     {
-                        await ctx.Response.Send(DefaultPages.Pages[500].Content, _Token).ConfigureAwait(false);
+                        if (ctx.Response.ChunkedTransfer)
+                            await ctx.Response.SendFinalChunk(Encoding.UTF8.GetBytes(DefaultPages.Pages[500].Content), _Token).ConfigureAwait(false);
+                        else
+                            await ctx.Response.Send(DefaultPages.Pages[500].Content, _Token).ConfigureAwait(false);
                     }
                     catch
                     {
@@ -657,7 +663,10 @@ private async void ClientConnected(object sender, ClientConnectedEventArgs args)
                     {
                         ctx.Response.StatusCode = 500;
                         ctx.Response.ContentType = DefaultPages.Pages[500].ContentType;
-                        await ctx.Response.Send(DefaultPages.Pages[500].Content).ConfigureAwait(false);
+                        if (ctx.Response.ChunkedTransfer)
+                            await ctx.Response.SendFinalChunk(Encoding.UTF8.GetBytes(DefaultPages.Pages[500].Content)).ConfigureAwait(false);
+                        else
+                            await ctx.Response.Send(DefaultPages.Pages[500].Content).ConfigureAwait(false);
                     }
 
                     ctx.Timestamp.End = DateTime.UtcNow;
diff --git a/src/WatsonWebserver/Webserver.cs b/src/WatsonWebserver/Webserver.cs
index 900597a..cabb94a 100644
--- a/src/WatsonWebserver/Webserver.cs
+++ b/src/WatsonWebserver/Webserver.cs
@@ -9,6 +9,7 @@
 using System.Text.Json.Serialization;
 using WatsonWebserver.Core;
 using System.Runtime.InteropServices;
+using System.Text;
 
 namespace WatsonWebserver
 {
@@ -554,7 +555,10 @@ private async Task AcceptConnections(CancellationToken token)
                             {
                                 ctx.Response.StatusCode = 404;
                                 ctx.Response.ContentType = DefaultPages.Pages[404].ContentType;
-                                await ctx.Response.Send(DefaultPages.Pages[404].Content).ConfigureAwait(false);
+                                if (ctx.Response.ChunkedTransfer)
+                                    await ctx.Response.SendFinalChunk(Encoding.UTF8.GetBytes(DefaultPages.Pages[404].Content)).ConfigureAwait(false);
+                                else
+                                    await ctx.Response.Send(DefaultPages.Pages[404].Content).ConfigureAwait(false);
                                 return;
                             }
 
@@ -564,7 +568,10 @@ private async Task AcceptConnections(CancellationToken token)
                         {
                             ctx.Response.StatusCode = 500;
                             ctx.Response.ContentType = DefaultPages.Pages[500].ContentType;
-                            await ctx.Response.Send(DefaultPages.Pages[500].Content).ConfigureAwait(false);
+                            if (ctx.Response.ChunkedTransfer)
+                                await ctx.Response.SendFinalChunk(Encoding.UTF8.GetBytes(DefaultPages.Pages[500].Content)).ConfigureAwait(false);
+                            else
+                                await ctx.Response.Send(DefaultPages.Pages[500].Content).ConfigureAwait(false);
                             Events.HandleExceptionEncountered(this, new ExceptionEventArgs(ctx, eInner));
                         }
                         finally
@@ -577,7 +584,10 @@ private async Task AcceptConnections(CancellationToken token)
                                 {
                                     ctx.Response.StatusCode = 500;
                                     ctx.Response.ContentType = DefaultPages.Pages[500].ContentType;
-                                    await ctx.Response.Send(DefaultPages.Pages[500].Content).ConfigureAwait(false);
+                                    if (ctx.Response.ChunkedTransfer)
+                                        await ctx.Response.SendFinalChunk(Encoding.UTF8.GetBytes(DefaultPages.Pages[500].Content)).ConfigureAwait(false);
+                                    else
+                                        await ctx.Response.Send(DefaultPages.Pages[500].Content).ConfigureAwait(false);
                                 }
 
                                 ctx.Timestamp.End = DateTime.UtcNow;