diff --git a/FastTunnel.Api/Controllers/BaseController.cs b/FastTunnel.Api/Controllers/BaseController.cs index 7a40cff..083af95 100644 --- a/FastTunnel.Api/Controllers/BaseController.cs +++ b/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(); diff --git a/FastTunnel.Core/Client/FastTunnelServer.cs b/FastTunnel.Core/Client/FastTunnelServer.cs index f794d73..6612a20 100644 --- a/FastTunnel.Core/Client/FastTunnelServer.cs +++ b/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 logger; - public ConcurrentDictionary> ResponseTasks { get; } = new(); + public ConcurrentDictionary, CancellationToken)> ResponseTasks { get; } = new(); public ConcurrentDictionary WebList { get; private set; } = new(); diff --git a/FastTunnel.Core/Extensions/ServicesExtensions.cs b/FastTunnel.Core/Extensions/ServicesExtensions.cs index f9e0a37..91824d1 100644 --- a/FastTunnel.Core/Extensions/ServicesExtensions.cs +++ b/FastTunnel.Core/Extensions/ServicesExtensions.cs @@ -70,6 +70,7 @@ namespace FastTunnel.Core.Extensions { services.AddReverseProxy().LoadFromMemory(); services.AddSingleton(); + services.AddHttpContextAccessor(); services.Configure(configurationSection) .AddSingleton() diff --git a/FastTunnel.Core/FastTunnel.Core.csproj b/FastTunnel.Core/FastTunnel.Core.csproj index 7b4eec4..c232960 100644 --- a/FastTunnel.Core/FastTunnel.Core.csproj +++ b/FastTunnel.Core/FastTunnel.Core.csproj @@ -21,12 +21,13 @@ + - + True diff --git a/FastTunnel.Core/Forwarder/FastTunnelForwarderHttpClientFactory.cs b/FastTunnel.Core/Forwarder/FastTunnelForwarderHttpClientFactory.cs index f1e2d17..816d612 100644 --- a/FastTunnel.Core/Forwarder/FastTunnelForwarderHttpClientFactory.cs +++ b/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 logger; readonly FastTunnelServer fastTunnelServer; + private readonly IHttpContextAccessor _httpContextAccessor; - public FastTunnelForwarderHttpClientFactory(ILogger logger, FastTunnelServer fastTunnelServer) + public FastTunnelForwarderHttpClientFactory( + ILogger 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()!; + + 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 tcs = new(cancellation); + TaskCompletionSource 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; diff --git a/FastTunnel.Core/Forwarder/MiddleWare/FastTunnelSwapHandler.cs b/FastTunnel.Core/Forwarder/MiddleWare/FastTunnelSwapHandler.cs index c7f35de..9920118 100644 --- a/FastTunnel.Core/Forwarder/MiddleWare/FastTunnelSwapHandler.cs +++ b/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(); - lifetime.ConnectionClosed.Register((task) => - { - (task as TaskCompletionSource).SetResult(null); - }, closedAwaiter); + //lifetime.ConnectionClosed.Register((task) => + //{ + // (task as TaskCompletionSource).SetResult(null); + //}, closedAwaiter); - await closedAwaiter.Task; + await closedAwaiter.Task.WaitAsync(cts.Token); logger.LogDebug($"[PROXY]:Closed {requestId}"); } catch (Exception ex) diff --git a/FastTunnel.Core/Handlers/Server/ForwardDispatcher.cs b/FastTunnel.Core/Handlers/Server/ForwardDispatcher.cs index 31bb54d..119ee56 100644 --- a/FastTunnel.Core/Handlers/Server/ForwardDispatcher.cs +++ b/FastTunnel.Core/Handlers/Server/ForwardDispatcher.cs @@ -53,7 +53,7 @@ namespace FastTunnel.Core.Handlers.Server var tcs = new TaskCompletionSource(); tcs.SetTimeOut(10000, () => { logger.LogDebug($"[Dispatch TimeOut]:{msgId}"); }); - _server.ResponseTasks.TryAdd(msgId, tcs); + _server.ResponseTasks.TryAdd(msgId, (tcs, CancellationToken.None)); try { diff --git a/FastTunnel.Server/Startup.cs b/FastTunnel.Server/Startup.cs index bc0942f..b4aea36 100644 --- a/FastTunnel.Server/Startup.cs +++ b/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(); // -------------------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();