1

FastTunnel.Api/Controllers/BaseController.cs

@@ -4,6 +4,7 @@
 //     https://github.com/FastTunnel/FastTunnel/edit/v2/LICENSE
 // Copyright (c) 2019 Gui.H
 
+using FastTunnel.Api.Filters;
 using FastTunnel.Server.Models;
 using Microsoft.AspNetCore.Authorization;
 using Microsoft.AspNetCore.Mvc;
@@ -13,6 +14,7 @@ namespace FastTunnel.Api.Controllers
     [Authorize]
     [Route("api/[controller]/[action]")]
     [ApiController]
+    [ServiceFilter(typeof(CustomExceptionFilterAttribute))]
     public class BaseController : ControllerBase
     {
         protected ApiResponse ApiResponse = new ApiResponse();

FastTunnel.Core/Client/FastTunnelServer.cs

@@ -1,4 +1,4 @@
-// Licensed under the Apache License, Version 2.0 (the "License").
+// Licensed under the Apache License, Version 2.0 (the "License").
 // You may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //     https://github.com/FastTunnel/FastTunnel/edit/v2/LICENSE
@@ -25,7 +25,7 @@ namespace FastTunnel.Core.Client
         public IProxyConfigProvider proxyConfig;
         readonly ILogger<FastTunnelServer> logger;
 
-        public ConcurrentDictionary<string, TaskCompletionSource<Stream>> ResponseTasks { get; } = new();
+        public ConcurrentDictionary<string, (TaskCompletionSource<Stream>, CancellationToken)> ResponseTasks { get; } = new();
 
         public ConcurrentDictionary<string, WebInfo> WebList { get; private set; } = new();
 

FastTunnel.Core/Extensions/ServicesExtensions.cs

@@ -70,6 +70,7 @@ namespace FastTunnel.Core.Extensions
         {
             services.AddReverseProxy().LoadFromMemory();
             services.AddSingleton<IForwarderHttpClientFactory, FastTunnelForwarderHttpClientFactory>();
+            services.AddHttpContextAccessor();
 
             services.Configure<DefaultServerConfig>(configurationSection)
                 .AddSingleton<IExceptionFilter, FastTunnelExceptionFilter>()

FastTunnel.Core/FastTunnel.Core.csproj

@@ -21,12 +21,13 @@
   <ItemGroup>
     <FrameworkReference Include="Microsoft.AspNetCore.App" />
   </ItemGroup>
+  
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Logging.Abstractions" Version="5.0.0" />
     <PackageReference Include="Microsoft.Extensions.Options.ConfigurationExtensions" Version="5.0.0" />
     <PackageReference Include="Yarp.ReverseProxy" Version="1.1.0" />
   </ItemGroup>
-
+  
   <ItemGroup>
     <Compile Update="TunnelResource.Designer.cs">
       <DesignTime>True</DesignTime>

FastTunnel.Core/Forwarder/FastTunnelForwarderHttpClientFactory.cs

@@ -1,4 +1,4 @@
-// Copyright (c) 2019-2022 Gui.H. https://github.com/FastTunnel/FastTunnel
+// Copyright (c) 2019-2022 Gui.H. https://github.com/FastTunnel/FastTunnel
 // The FastTunnel licenses this file to you under the Apache License Version 2.0.
 // For more details,You may obtain License file at: https://github.com/FastTunnel/FastTunnel/blob/v2/LICENSE
 
@@ -6,6 +6,8 @@ using FastTunnel.Core.Client;
 using FastTunnel.Core.Extensions;
 using FastTunnel.Core.Models;
 using FastTunnel.Core.Sockets;
+using Microsoft.AspNetCore.Connections.Features;
+using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Logging;
 using System;
 using System.Collections.Concurrent;
@@ -26,11 +28,15 @@ namespace FastTunnel.Core.Forwarder
     {
         readonly ILogger<FastTunnelForwarderHttpClientFactory> logger;
         readonly FastTunnelServer fastTunnelServer;
+        private readonly IHttpContextAccessor _httpContextAccessor;
 
-        public FastTunnelForwarderHttpClientFactory(ILogger<FastTunnelForwarderHttpClientFactory> logger, FastTunnelServer fastTunnelServer)
+        public FastTunnelForwarderHttpClientFactory(
+            ILogger<FastTunnelForwarderHttpClientFactory> logger,
+            IHttpContextAccessor httpContextAccessor, FastTunnelServer fastTunnelServer)
         {
             this.fastTunnelServer = fastTunnelServer;
             this.logger = logger;
+            _httpContextAccessor = httpContextAccessor;
         }
 
         protected override void ConfigureHandler(ForwarderHttpClientContext context, SocketsHttpHandler handler)
@@ -43,9 +49,17 @@ namespace FastTunnel.Core.Forwarder
         {
             var host = context.InitialRequestMessage.RequestUri.Host;
 
+            var contextRequest = _httpContextAccessor.HttpContext;
+            //var lifetime = contextRequest.Features.Get<IConnectionLifetimeFeature>()!;
+
+            contextRequest.RequestAborted.Register(() =>
+            {
+                logger.LogDebug($"[ConnectionClosed]");
+            });
+
             try
             {
-                var res = await proxyAsync(host, context, cancellationToken);
+                var res = await proxyAsync(host, context, contextRequest.RequestAborted);
                 return res;
             }
             catch (Exception)
@@ -65,17 +79,22 @@ namespace FastTunnel.Core.Forwarder
 
             var msgId = Guid.NewGuid().ToString().Replace("-", "");
 
-            TaskCompletionSource<Stream> tcs = new(cancellation);
+            TaskCompletionSource<Stream> tcs = new();
             logger.LogDebug($"[Http]Swap开始 {msgId}|{host}=>{web.WebConfig.LocalIp}:{web.WebConfig.LocalPort}");
-            tcs.SetTimeOut(10000, () => { logger.LogDebug($"[Proxy TimeOut]:{msgId}"); });
 
-            fastTunnelServer.ResponseTasks.TryAdd(msgId, tcs);
+            cancellation.Register(() =>
+            {
+                logger.LogDebug($"[Proxy TimeOut]:{msgId}");
+                tcs.TrySetCanceled();
+            });
+
+            fastTunnelServer.ResponseTasks.TryAdd(msgId, (tcs, cancellation));
 
             try
             {
                 // 发送指令给客户端，等待建立隧道
                 await web.Socket.SendCmdAsync(MessageType.SwapMsg, $"{msgId}|{web.WebConfig.LocalIp}:{web.WebConfig.LocalPort}", cancellation);
-                var res = await tcs.Task;
+                var res = await tcs.Task.WaitAsync(cancellation);
 
                 logger.LogDebug($"[Http]Swap OK {msgId}");
                 return res;

FastTunnel.Core/Forwarder/MiddleWare/FastTunnelSwapHandler.cs

@@ -1,4 +1,4 @@
-using FastTunnel.Core.Client;
+using FastTunnel.Core.Client;
 using FastTunnel.Core.Extensions;
 using Microsoft.AspNetCore.Connections.Features;
 using Microsoft.AspNetCore.Http;
@@ -7,6 +7,7 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace FastTunnel.Core.Forwarder.MiddleWare
@@ -50,16 +51,26 @@ namespace FastTunnel.Core.Forwarder.MiddleWare
                 }
 
                 using var reverseConnection = new WebSocketStream(lifetime, transport);
-                responseAwaiter.TrySetResult(reverseConnection);
+                responseAwaiter.Item1.TrySetResult(reverseConnection);
+
+                CancellationTokenSource cts;
+                if (responseAwaiter.Item2 != CancellationToken.None)
+                {
+                    cts = CancellationTokenSource.CreateLinkedTokenSource(lifetime.ConnectionClosed, responseAwaiter.Item2);
+                }
+                else
+                {
+                    cts = CancellationTokenSource.CreateLinkedTokenSource(lifetime.ConnectionClosed);
+                }
 
                 var closedAwaiter = new TaskCompletionSource<object>();
 
-                lifetime.ConnectionClosed.Register((task) =>
-                {
-                    (task as TaskCompletionSource<object>).SetResult(null);
-                }, closedAwaiter);
+                //lifetime.ConnectionClosed.Register((task) =>
+                //{
+                //    (task as TaskCompletionSource<object>).SetResult(null);
+                //}, closedAwaiter);
 
-                await closedAwaiter.Task;
+                await closedAwaiter.Task.WaitAsync(cts.Token);
                 logger.LogDebug($"[PROXY]:Closed {requestId}");
             }
             catch (Exception ex)

FastTunnel.Core/Handlers/Server/ForwardDispatcher.cs

@@ -53,7 +53,7 @@ namespace FastTunnel.Core.Handlers.Server
                 var tcs = new TaskCompletionSource<Stream>();
                 tcs.SetTimeOut(10000, () => { logger.LogDebug($"[Dispatch TimeOut]:{msgId}"); });
 
-                _server.ResponseTasks.TryAdd(msgId, tcs);
+                _server.ResponseTasks.TryAdd(msgId, (tcs, CancellationToken.None));
 
                 try
                 {

FastTunnel.Server/Startup.cs

@@ -16,6 +16,7 @@ using Microsoft.IdentityModel.Tokens;
 using System;
 using FastTunnel.Core.Config;
 using System.Text;
+using FastTunnel.Api.Filters;
 
 #if DEBUG
 using Microsoft.OpenApi.Models;
@@ -80,7 +81,7 @@ public class Startup
             c.SwaggerDoc("v2", new OpenApiInfo { Title = "FastTunel.Api", Version = "v2" });
         });
 #endif
-
+        services.AddSingleton<CustomExceptionFilterAttribute>();
         // -------------------FastTunnel STEP1 OF 3------------------
         services.AddFastTunnelServer(Configuration.GetSection("FastTunnel"));
         // -------------------FastTunnel STEP1 END-------------------
@@ -98,12 +99,12 @@ public class Startup
 #endif
         }
 
+        app.UseRouting();
+
         // -------------------FastTunnel STEP2 OF 3------------------
         app.UseFastTunnelServer();
         // -------------------FastTunnel STEP2 END-------------------
 
-        app.UseRouting();
-
         // --------------------- Custom UI ----------------
         app.UseStaticFiles();
         app.UseAuthentication();