From a0db75bfdd09e74faed77e763d4b6fbc76e40418 Mon Sep 17 00:00:00 2001
From: uwaa <uwaa@noreply.git.rape.pet>
Date: Mon, 20 Jan 2025 22:28:00 +0000
Subject: [PATCH] http: longer timeouts for ws

---
 HTTP/HttpServer.cs           |  8 ++++++++
 HTTP/HttpStream.cs           | 35 +++++++++++++++++++----------------
 HTTP/Websockets/Websocket.cs |  9 +++++++++
 3 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/HTTP/HttpServer.cs b/HTTP/HttpServer.cs
index 12d5a82..6426bd9 100644
--- a/HTTP/HttpServer.cs
+++ b/HTTP/HttpServer.cs
@@ -60,6 +60,11 @@ public sealed class HttpServer
     /// </summary>
     public TimeSpan Timeout = TimeSpan.FromSeconds(20);
 
+    /// <summary>
+    /// The maximum time a websocket may be inactive before it is presumed dead and closed.
+    /// </summary>
+    public TimeSpan TimeoutWS = TimeSpan.FromSeconds(60);
+
     readonly Dictionary<IPAddress, int> IPCounts = new Dictionary<IPAddress, int>();
 
     readonly SemaphoreSlim IPCountsLock = new SemaphoreSlim(1, 1);
@@ -165,6 +170,9 @@ public sealed class HttpServer
 
                     if (response is SwitchingProtocols swp)
                     {
+                        httpStream.Timeout = TimeoutWS;
+                        httpStream.KeepAlive();
+
                         //Create and run websocket
                         WebsocketRemote ws = new WebsocketRemote(req, clientInfo, httpStream, swp.Fields.WebSocketProtocol);
                         CloseStatus closeStatus = await swp.Callback(ws);
diff --git a/HTTP/HttpStream.cs b/HTTP/HttpStream.cs
index d30253d..56c6e1c 100644
--- a/HTTP/HttpStream.cs
+++ b/HTTP/HttpStream.cs
@@ -1,8 +1,6 @@
 using System.Buffers;
-using System.Collections.Specialized;
 using System.Net.Sockets;
 using System.Text;
-using System.Web;
 
 namespace Uwaa.HTTP;
 
@@ -28,19 +26,27 @@ class HttpStream : IDisposable
     /// </summary>
     public TimeSpan Timeout;
 
+    internal readonly CancellationTokenSource CancelSrc = new CancellationTokenSource();
+
     public HttpStream(Stream stream, TimeSpan timeout) : base()
     {
         Stream = stream;
         Timeout = timeout;
         Buffer = new BufferedStream(stream);
         Decoder = Encoding.ASCII.GetDecoder();
-        
+    }
+
+    /// <summary>
+    /// Resets the timeout timer.
+    /// </summary>
+    public void KeepAlive()
+    {
+        CancelSrc.CancelAfter(Timeout);
     }
 
     public async ValueTask<string> ReadLine()
     {
-        CancellationTokenSource cancelSrc = new CancellationTokenSource();
-        cancelSrc.CancelAfter(Timeout);
+        CancelSrc.CancelAfter(Timeout);
 
         const int maxChars = 4096;
         byte[] dataBuffer = ArrayPool<byte>.Shared.Rent(1);
@@ -50,7 +56,7 @@ class HttpStream : IDisposable
             int charBufferIndex = 0;
             while (true)
             {
-                if (await Buffer.ReadAsync(dataBuffer.AsMemory(0, 1), cancelSrc.Token) == 0)
+                if (await Stream.ReadAsync(dataBuffer.AsMemory(0, 1), CancelSrc.Token) == 0)
                     if (charBufferIndex == 0)
                         throw new SocketException((int)SocketError.ConnectionReset);
                     else
@@ -87,15 +93,14 @@ class HttpStream : IDisposable
 
     public async ValueTask<int> Read(Memory<byte> buffer)
     {
-        CancellationTokenSource cancelSrc = new CancellationTokenSource();
-        cancelSrc.CancelAfter(Timeout);
+        CancelSrc.CancelAfter(Timeout);
 
         try
         {
             int index = 0;
             while (index < buffer.Length)
             {
-                int count = await Buffer.ReadAsync(buffer[index..], cancelSrc.Token);
+                int count = await Stream.ReadAsync(buffer[index..], CancelSrc.Token);
                 if (count == 0)
                     break;
 
@@ -134,11 +139,10 @@ class HttpStream : IDisposable
 
     public ValueTask Write(ReadOnlyMemory<byte> bytes)
     {
-        CancellationTokenSource cancelSrc = new CancellationTokenSource();
-        cancelSrc.CancelAfter(Timeout);
+        CancelSrc.CancelAfter(Timeout);
         try
         {
-            return Buffer.WriteAsync(bytes, cancelSrc.Token);
+            return Buffer.WriteAsync(bytes, CancelSrc.Token);
         }
         catch (IOException e)
         {
@@ -151,12 +155,11 @@ class HttpStream : IDisposable
 
     public async Task Flush()
     {
-        CancellationTokenSource cancelSrc = new CancellationTokenSource();
-        cancelSrc.CancelAfter(Timeout);
+        CancelSrc.CancelAfter(Timeout);
         try
         {
-            await Buffer.FlushAsync(cancelSrc.Token);
-            await Stream.FlushAsync(cancelSrc.Token);
+            await Buffer.FlushAsync(CancelSrc.Token);
+            await Stream.FlushAsync(CancelSrc.Token);
         }
         catch (IOException e)
         {
diff --git a/HTTP/Websockets/Websocket.cs b/HTTP/Websockets/Websocket.cs
index b08f855..7a34535 100644
--- a/HTTP/Websockets/Websocket.cs
+++ b/HTTP/Websockets/Websocket.cs
@@ -14,6 +14,15 @@ public class Websocket
     /// </summary>
     public readonly string? SubProtocol;
 
+    /// <summary>
+    /// The maximum time the websocket may be inactive before it is presumed dead and closed.
+    /// </summary>
+    public TimeSpan Timeout
+    {
+        get => Stream.Timeout;
+        set => Stream.Timeout = value;
+    }
+
     internal readonly HttpStream Stream;
 
     readonly List<byte> finalPayload = new List<byte>();