fix: Properly dispose stream

GZTimeWalker · Apr 14, 2024 · 4dfd1da · 4dfd1da
1 parent e24950e
commit 4dfd1da
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 35 deletions.
diff --git a/src/GZCTF/Controllers/ProxyController.cs b/src/GZCTF/Controllers/ProxyController.cs
@@ -31,13 +31,13 @@ public class ProxyController(
     readonly bool _enablePlatformProxy = provider.Value.PortMappingType == ContainerPortMappingType.PlatformProxy;
     readonly bool _enableTrafficCapture = provider.Value.EnableTrafficCapture;
 
-    readonly JsonSerializerOptions _jsonOptions =
+    static readonly JsonSerializerOptions _jsonOptions =
         new() { Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping, WriteIndented = true };
 
-    readonly DistributedCacheEntryOptions _storeOption =
+    static readonly DistributedCacheEntryOptions _storeOption =
         new() { AbsoluteExpirationRelativeToNow = TimeSpan.FromHours(10) };
 
-    readonly DistributedCacheEntryOptions _validOption =
+    static readonly DistributedCacheEntryOptions _validOption =
         new() { AbsoluteExpirationRelativeToNow = TimeSpan.FromMinutes(10) };
 
     /// <summary>
@@ -155,47 +155,54 @@ public async Task<IActionResult> ProxyForNoInstance(Guid id, CancellationToken t
     {
         using var socket = new Socket(target.AddressFamily, SocketType.Stream, ProtocolType.Tcp);
 
-        RecordableNetworkStream? stream;
+        RecordableNetworkStream? stream = null;
+
         try
         {
-            await socket.ConnectAsync(target, token);
+            try
+            {
+                await socket.ConnectAsync(target, token);
 
-            if (!socket.Connected)
-                throw new SocketException((int)SocketError.NotConnected);
+                if (!socket.Connected)
+                    throw new SocketException((int)SocketError.NotConnected);
 
-            stream = new RecordableNetworkStream(socket, metadata, options);
-        }
-        catch (SocketException e)
-        {
-            logger.SystemLog(
-                Program.StaticLocalizer[nameof(Resources.Program.Proxy_ContainerConnectionFailedLog), e.SocketErrorCode,
-                    $"{target.Address}:{target.Port}"],
-                TaskStatus.Failed, LogLevel.Debug);
-            return new JsonResult(new RequestResponse(
-                localizer[nameof(Resources.Program.Proxy_ContainerConnectionFailed), e.SocketErrorCode],
-                StatusCodes.Status418ImATeapot))
-            { StatusCode = StatusCodes.Status418ImATeapot };
-        }
+                stream = new RecordableNetworkStream(socket, metadata, options);
+            }
+            catch (SocketException e)
+            {
+                logger.SystemLog(
+                    Program.StaticLocalizer[nameof(Resources.Program.Proxy_ContainerConnectionFailedLog), e.SocketErrorCode,
+                        $"{target.Address}:{target.Port}"],
+                    TaskStatus.Failed, LogLevel.Debug);
+                return new JsonResult(new RequestResponse(
+                    localizer[nameof(Resources.Program.Proxy_ContainerConnectionFailed), e.SocketErrorCode],
+                    StatusCodes.Status418ImATeapot))
+                { StatusCode = StatusCodes.Status418ImATeapot };
+            }
 
-        using WebSocket ws = await HttpContext.WebSockets.AcceptWebSocketAsync();
+            using WebSocket ws = await HttpContext.WebSockets.AcceptWebSocketAsync();
 
-        try
-        {
-            var (tx, rx) = await RunProxy(stream, ws, token);
-            logger.SystemLog($"[{id}] {client.Address} -> {target.Address}:{target.Port}, tx {tx}, rx {rx}",
-                TaskStatus.Success, LogLevel.Debug);
-        }
-        catch (Exception e)
-        {
-            logger.LogError(e, Program.StaticLocalizer[nameof(Resources.Program.Proxy_Error)]);
+            try
+            {
+                var (tx, rx) = await RunProxy(stream, ws, token);
+                logger.SystemLog($"[{id}] {client.Address} -> {target.Address}:{target.Port}, tx {tx}, rx {rx}",
+                    TaskStatus.Success, LogLevel.Debug);
+            }
+            catch (Exception e)
+            {
+                logger.LogError(e, Program.StaticLocalizer[nameof(Resources.Program.Proxy_Error)]);
+            }
+            finally
+            {
+                await DecrementConnectionCount(CacheKey.ConnectionCount(id));
+            }
+
+            return new EmptyResult();
         }
         finally
         {
-            await DecrementConnectionCount(CacheKey.ConnectionCount(id));
-            stream.Close();
+            stream?.Dispose();
         }
-
-        return new EmptyResult();
     }
 
     /// <summary>
@@ -208,7 +215,7 @@ public async Task<IActionResult> ProxyForNoInstance(Guid id, CancellationToken t
     static async Task<(ulong, ulong)> RunProxy(RecordableNetworkStream stream, WebSocket ws,
         CancellationToken token = default)
     {
-        var cts = CancellationTokenSource.CreateLinkedTokenSource(token);
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(token);
         cts.CancelAfter(TimeSpan.FromMinutes(30));
 
         CancellationToken ct = cts.Token;

diff --git a/src/GZCTF/Utils/RecordableNetworkStream.cs b/src/GZCTF/Utils/RecordableNetworkStream.cs
@@ -41,6 +41,8 @@ public sealed class RecordableNetworkStream : NetworkStream
     readonly IPEndPoint _host = new(0, 65535);
     readonly RecordableNetworkStreamOptions _options;
 
+    bool _disposed;
+
     public RecordableNetworkStream(Socket socket, byte[]? metadata, RecordableNetworkStreamOptions options) :
         base(socket)
     {
@@ -106,6 +108,19 @@ void WriteCapturedData(IPEndPoint source, IPEndPoint dest, ReadOnlyMemory<byte>
         _device?.Write(new RawCapture(LinkLayers.Ethernet, new(), packet.Bytes));
     }
 
+    protected override void Dispose(bool disposing)
+    {
+        if (!_disposed)
+        {
+            base.Dispose(disposing);
+
+            Close();
+            _device?.Dispose();
+        }
+
+        _disposed = true;
+    }
+
     public override void Close()
     {
         base.Close();