Commit 2bbf9bc7 authored by justcoding121's avatar justcoding121

#410 implement task cancellation

parent de8d70bc
using System;
using System.Threading;
namespace Titanium.Web.Proxy.EventArguments
{
public class BeforeSslAuthenticateEventArgs : EventArgs
{
internal CancellationTokenSource TaskCancellationSource;
internal BeforeSslAuthenticateEventArgs(CancellationTokenSource taskCancellationSource)
{
this.TaskCancellationSource = taskCancellationSource;
}
public string SniHostName { get; internal set; }
public bool DecryptSsl { get; set; } = true;
public bool TerminateSession { get; set; }
public void TerminateSession()
{
TaskCancellationSource.Cancel();
}
}
}
......@@ -11,6 +11,7 @@ using Titanium.Web.Proxy.Http;
using Titanium.Web.Proxy.Http.Responses;
using Titanium.Web.Proxy.Models;
using Titanium.Web.Proxy.Network;
using System.Threading;
namespace Titanium.Web.Proxy.EventArguments
{
......@@ -56,13 +57,13 @@ namespace Titanium.Web.Proxy.EventArguments
/// <summary>
/// Constructor to initialize the proxy
/// </summary>
internal SessionEventArgs(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc)
: this(bufferSize, endPoint, exceptionFunc, null)
internal SessionEventArgs(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc, CancellationTokenSource cancellationTokenSource)
: this(bufferSize, endPoint, exceptionFunc, null, cancellationTokenSource)
{
}
protected SessionEventArgs(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc, Request request)
: base(bufferSize, endPoint, exceptionFunc, request)
protected SessionEventArgs(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc, Request request, CancellationTokenSource cancellationTokenSource)
: base(bufferSize, endPoint, exceptionFunc, request, cancellationTokenSource)
{
}
......
using System;
using System.Net;
using System.Threading;
using Titanium.Web.Proxy.Helpers;
using Titanium.Web.Proxy.Http;
using Titanium.Web.Proxy.Models;
......@@ -22,6 +23,8 @@ namespace Titanium.Web.Proxy.EventArguments
protected readonly ExceptionHandler ExceptionFunc;
internal CancellationTokenSource cancellationTokenSource;
/// <summary>
/// Holds a reference to client
/// </summary>
......@@ -60,11 +63,6 @@ namespace Titanium.Web.Proxy.EventArguments
public ProxyEndPoint LocalEndPoint { get; }
/// <summary>
/// Terminates the session abruptly by terminating client/server connections
/// </summary>
public bool TerminateSession { get; set; }
public bool IsTransparent => LocalEndPoint is TransparentProxyEndPoint;
public Exception Exception { get; internal set; }
......@@ -72,15 +70,17 @@ namespace Titanium.Web.Proxy.EventArguments
/// <summary>
/// Constructor to initialize the proxy
/// </summary>
internal SessionEventArgsBase(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc)
: this(bufferSize, endPoint, exceptionFunc, null)
internal SessionEventArgsBase(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc, CancellationTokenSource cancellationTokenSource)
: this(bufferSize, endPoint, exceptionFunc, null, cancellationTokenSource)
{
}
protected SessionEventArgsBase(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc, Request request)
protected SessionEventArgsBase(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc,
Request request, CancellationTokenSource cancellationTokenSource)
{
this.BufferSize = bufferSize;
this.ExceptionFunc = exceptionFunc;
this.cancellationTokenSource = cancellationTokenSource;
ProxyClient = new ProxyClient();
WebSession = new HttpWebClient(bufferSize, request);
......@@ -130,6 +130,13 @@ namespace Titanium.Web.Proxy.EventArguments
}
}
/// <summary>
/// Terminates the session abruptly by terminating client/server connections
/// </summary>
public void TerminateSession()
{
cancellationTokenSource.Cancel();
}
/// <summary>
/// implement any cleanup here
/// </summary>
......
using System;
using System.Threading;
using Titanium.Web.Proxy.Http;
using Titanium.Web.Proxy.Models;
......@@ -15,8 +16,8 @@ namespace Titanium.Web.Proxy.EventArguments
public bool IsHttpsConnect { get; internal set; }
internal TunnelConnectSessionEventArgs(int bufferSize, ProxyEndPoint endPoint, ConnectRequest connectRequest, ExceptionHandler exceptionFunc)
: base(bufferSize, endPoint, exceptionFunc, connectRequest)
internal TunnelConnectSessionEventArgs(int bufferSize, ProxyEndPoint endPoint, ConnectRequest connectRequest, ExceptionHandler exceptionFunc, CancellationTokenSource cancellationTokenSource)
: base(bufferSize, endPoint, exceptionFunc, connectRequest, cancellationTokenSource)
{
}
}
......
......@@ -12,6 +12,7 @@ using Titanium.Web.Proxy.Exceptions;
using Titanium.Web.Proxy.Helpers;
using Titanium.Web.Proxy.Http;
using Titanium.Web.Proxy.Models;
using System.Threading;
namespace Titanium.Web.Proxy
{
......@@ -26,6 +27,7 @@ namespace Titanium.Web.Proxy
/// <returns></returns>
private async Task HandleClient(ExplicitProxyEndPoint endPoint, TcpClient tcpClient)
{
var cancellationTokenSource = new CancellationTokenSource();
var clientStream = new CustomBufferedStream(tcpClient.GetStream(), BufferSize);
var clientStreamReader = new CustomBinaryReader(clientStream, BufferSize);
......@@ -60,7 +62,7 @@ namespace Titanium.Web.Proxy
await HeaderParser.ReadHeaders(clientStreamReader, connectRequest.Headers);
var connectArgs = new TunnelConnectSessionEventArgs(BufferSize, endPoint, connectRequest, ExceptionFunc);
var connectArgs = new TunnelConnectSessionEventArgs(BufferSize, endPoint, connectRequest, ExceptionFunc, cancellationTokenSource);
connectArgs.ProxyClient.TcpClient = tcpClient;
connectArgs.ProxyClient.ClientStream = clientStream;
......@@ -152,7 +154,7 @@ namespace Titanium.Web.Proxy
}
}
if (connectArgs.TerminateSession)
if (connectArgs.cancellationTokenSource.IsCancellationRequested)
{
throw new Exception("Session was terminated by user.");
}
......@@ -190,7 +192,7 @@ namespace Titanium.Web.Proxy
await TcpHelper.SendRaw(clientStream, connection.Stream, BufferSize,
(buffer, offset, count) => { connectArgs.OnDataSent(buffer, offset, count); },
(buffer, offset, count) => { connectArgs.OnDataReceived(buffer, offset, count); },
ExceptionFunc);
ExceptionFunc, connectArgs.cancellationTokenSource);
}
return;
......@@ -198,7 +200,7 @@ namespace Titanium.Web.Proxy
}
//Now create the request
await HandleHttpSessionRequest(tcpClient, clientStream, clientStreamReader, clientStreamWriter, connectHostname, endPoint, connectRequest);
await HandleHttpSessionRequest(tcpClient, clientStream, clientStreamReader, clientStreamWriter, connectHostname, endPoint, connectRequest, cancellationTokenSource);
}
catch (ProxyHttpException e)
{
......@@ -220,6 +222,10 @@ namespace Titanium.Web.Proxy
{
clientStreamReader.Dispose();
clientStream.Dispose();
if (!cancellationTokenSource.IsCancellationRequested)
{
cancellationTokenSource.Cancel();
}
}
}
......
......@@ -122,10 +122,9 @@ namespace Titanium.Web.Proxy.Helpers
/// <param name="exceptionFunc"></param>
/// <returns></returns>
internal static async Task SendRawApm(Stream clientStream, Stream serverStream, int bufferSize,
Action<byte[], int, int> onDataSend, Action<byte[], int, int> onDataReceive, ExceptionHandler exceptionFunc)
Action<byte[], int, int> onDataSend, Action<byte[], int, int> onDataReceive, ExceptionHandler exceptionFunc, CancellationTokenSource cts)
{
var tcs = new TaskCompletionSource<bool>();
var cts = new CancellationTokenSource();
cts.Token.Register(() => tcs.TrySetResult(true));
//Now async relay all server=>client & client=>server data
......@@ -144,10 +143,10 @@ namespace Titanium.Web.Proxy.Helpers
}
}
private static void BeginRead(Stream inputStream, Stream outputStream, byte[] buffer, CancellationTokenSource cts, Action<byte[], int, int> onCopy,
private static void BeginRead(Stream inputStream, Stream outputStream, byte[] buffer, CancellationTokenSource cancellationTokenSource, Action<byte[], int, int> onCopy,
ExceptionHandler exceptionFunc)
{
if (cts.IsCancellationRequested)
if (cancellationTokenSource.IsCancellationRequested)
{
return;
}
......@@ -155,7 +154,7 @@ namespace Titanium.Web.Proxy.Helpers
bool readFlag = false;
var readCallback = (AsyncCallback)(ar =>
{
if (cts.IsCancellationRequested || readFlag)
if (cancellationTokenSource.IsCancellationRequested || readFlag)
{
return;
}
......@@ -167,7 +166,7 @@ namespace Titanium.Web.Proxy.Helpers
int read = inputStream.EndRead(ar);
if (read <= 0)
{
cts.Cancel();
cancellationTokenSource.Cancel();
return;
}
......@@ -175,7 +174,7 @@ namespace Titanium.Web.Proxy.Helpers
var writeCallback = (AsyncCallback)(ar2 =>
{
if (cts.IsCancellationRequested)
if (cancellationTokenSource.IsCancellationRequested)
{
return;
}
......@@ -183,11 +182,11 @@ namespace Titanium.Web.Proxy.Helpers
try
{
outputStream.EndWrite(ar2);
BeginRead(inputStream, outputStream, buffer, cts, onCopy, exceptionFunc);
BeginRead(inputStream, outputStream, buffer, cancellationTokenSource, onCopy, exceptionFunc);
}
catch (IOException ex)
{
cts.Cancel();
cancellationTokenSource.Cancel();
exceptionFunc(ex);
}
});
......@@ -196,7 +195,7 @@ namespace Titanium.Web.Proxy.Helpers
}
catch (IOException ex)
{
cts.Cancel();
cancellationTokenSource.Cancel();
exceptionFunc(ex);
}
});
......@@ -220,17 +219,15 @@ namespace Titanium.Web.Proxy.Helpers
/// <param name="onDataReceive"></param>
/// <param name="exceptionFunc"></param>
/// <returns></returns>
internal static async Task SendRawTap(Stream clientStream, Stream serverStream, int bufferSize,
Action<byte[], int, int> onDataSend, Action<byte[], int, int> onDataReceive, ExceptionHandler exceptionFunc)
private static async Task SendRawTap(Stream clientStream, Stream serverStream, int bufferSize,
Action<byte[], int, int> onDataSend, Action<byte[], int, int> onDataReceive, ExceptionHandler exceptionFunc, CancellationTokenSource cancellationTokenSource)
{
var cts = new CancellationTokenSource();
//Now async relay all server=>client & client=>server data
var sendRelay = clientStream.CopyToAsync(serverStream, onDataSend, bufferSize, cts.Token);
var receiveRelay = serverStream.CopyToAsync(clientStream, onDataReceive, bufferSize, cts.Token);
var sendRelay = clientStream.CopyToAsync(serverStream, onDataSend, bufferSize, cancellationTokenSource.Token);
var receiveRelay = serverStream.CopyToAsync(clientStream, onDataReceive, bufferSize, cancellationTokenSource.Token);
await Task.WhenAny(sendRelay, receiveRelay);
cts.Cancel();
cancellationTokenSource.Cancel();
await Task.WhenAll(sendRelay, receiveRelay);
}
......@@ -247,10 +244,10 @@ namespace Titanium.Web.Proxy.Helpers
/// <param name="exceptionFunc"></param>
/// <returns></returns>
internal static Task SendRaw(Stream clientStream, Stream serverStream, int bufferSize,
Action<byte[], int, int> onDataSend, Action<byte[], int, int> onDataReceive, ExceptionHandler exceptionFunc)
Action<byte[], int, int> onDataSend, Action<byte[], int, int> onDataReceive, ExceptionHandler exceptionFunc, CancellationTokenSource cancellationTokenSource)
{
// todo: fix APM mode
return SendRawTap(clientStream, serverStream, bufferSize, onDataSend, onDataReceive, exceptionFunc);
return SendRawTap(clientStream, serverStream, bufferSize, onDataSend, onDataReceive, exceptionFunc, cancellationTokenSource);
}
}
}
\ No newline at end of file
......@@ -11,6 +11,7 @@ using Titanium.Web.Proxy.Helpers;
using Titanium.Web.Proxy.Http;
using Titanium.Web.Proxy.Models;
using Titanium.Web.Proxy.Network.Tcp;
using System.Threading;
namespace Titanium.Web.Proxy
{
......@@ -39,7 +40,7 @@ namespace Titanium.Web.Proxy
/// <returns></returns>
private async Task HandleHttpSessionRequest(TcpClient client, CustomBufferedStream clientStream,
CustomBinaryReader clientStreamReader, HttpResponseWriter clientStreamWriter, string httpsConnectHostname,
ProxyEndPoint endPoint, ConnectRequest connectRequest, bool isTransparentEndPoint = false)
ProxyEndPoint endPoint, ConnectRequest connectRequest, CancellationTokenSource cancellationTokenSource, bool isTransparentEndPoint = false)
{
TcpConnection connection = null;
......@@ -56,7 +57,7 @@ namespace Titanium.Web.Proxy
return;
}
var args = new SessionEventArgs(BufferSize, endPoint, ExceptionFunc)
var args = new SessionEventArgs(BufferSize, endPoint, ExceptionFunc, cancellationTokenSource)
{
ProxyClient = { TcpClient = client },
WebSession = { ConnectRequest = connectRequest }
......@@ -204,7 +205,7 @@ namespace Titanium.Web.Proxy
await TcpHelper.SendRaw(clientStream, connection.Stream, BufferSize,
(buffer, offset, count) => { args.OnDataSent(buffer, offset, count); },
(buffer, offset, count) => { args.OnDataReceived(buffer, offset, count); },
ExceptionFunc);
ExceptionFunc, args.cancellationTokenSource);
return;
}
......@@ -224,7 +225,7 @@ namespace Titanium.Web.Proxy
return;
}
if (args.TerminateSession)
if (args.cancellationTokenSource.IsCancellationRequested)
{
throw new Exception("Session was terminated by user.");
}
......
......@@ -10,6 +10,7 @@ using Titanium.Web.Proxy.EventArguments;
using Titanium.Web.Proxy.Extensions;
using Titanium.Web.Proxy.Helpers;
using Titanium.Web.Proxy.Models;
using System.Threading;
namespace Titanium.Web.Proxy
{
......@@ -25,6 +26,8 @@ namespace Titanium.Web.Proxy
/// <returns></returns>
private async Task HandleClient(TransparentProxyEndPoint endPoint, TcpClient tcpClient)
{
var cancellationTokenSource = new CancellationTokenSource();
var clientStream = new CustomBufferedStream(tcpClient.GetStream(), BufferSize);
var clientStreamReader = new CustomBinaryReader(clientStream, BufferSize);
......@@ -41,11 +44,11 @@ namespace Titanium.Web.Proxy
{
httpsHostName = clientHelloInfo.GetServerName() ?? endPoint.GenericCertificateName;
var args = new BeforeSslAuthenticateEventArgs();
var args = new BeforeSslAuthenticateEventArgs(cancellationTokenSource);
args.SniHostName = httpsHostName;
await endPoint.InvokeBeforeSslAuthenticate(this, args, ExceptionFunc);
if (args.TerminateSession)
if (args.TaskCancellationSource.IsCancellationRequested)
{
throw new Exception("Session was terminated by user.");
}
......@@ -112,7 +115,7 @@ namespace Titanium.Web.Proxy
//var serverHelloInfo = await SslTools.PeekServerHello(serverStream);
await TcpHelper.SendRaw(clientStream, serverStream, BufferSize,
null, null, ExceptionFunc);
null, null, ExceptionFunc, args.TaskCancellationSource);
}
}
}
......@@ -120,12 +123,16 @@ namespace Titanium.Web.Proxy
//HTTPS server created - we can now decrypt the client's traffic
//Now create the request
await HandleHttpSessionRequest(tcpClient, clientStream, clientStreamReader, clientStreamWriter,
isHttps ? httpsHostName : null, endPoint, null, true);
isHttps ? httpsHostName : null, endPoint, null, cancellationTokenSource, true);
}
finally
{
clientStreamReader.Dispose();
clientStream.Dispose();
if (!cancellationTokenSource.IsCancellationRequested)
{
cancellationTokenSource.Cancel();
}
}
}
}
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment