Commit 49d94058 authored by justcoding121's avatar justcoding121

refactor connection factory

parent db6a0291
......@@ -151,7 +151,7 @@ namespace Titanium.Web.Proxy.Examples.Basic
private Task OnBeforeTunnelConnectResponse(object sender, TunnelConnectSessionEventArgs e)
{
return default;
return Task.FromResult(false);
}
// intecept & cancel redirect or update requests
......
......@@ -140,7 +140,7 @@ namespace Titanium.Web.Proxy
{
// test server HTTP/2 support
// todo: this is a hack, because Titanium does not support HTTP protocol changing currently
var connection = await getServerConnection(connectArgs, true,
var connection = await tcpConnectionFactory.GetServerConnection(this, connectArgs, true,
SslExtensions.Http2ProtocolAsList, true, cancellationToken);
http2Supported = connection.NegotiatedApplicationProtocol == SslApplicationProtocol.Http2;
......@@ -153,7 +153,7 @@ namespace Titanium.Web.Proxy
//don't pass cancellation token here
//it could cause floating server connections when client exits
prefetchConnectionTask = getServerConnection(connectArgs, true,
prefetchConnectionTask = tcpConnectionFactory.GetServerConnection(this, connectArgs, true,
null, false, CancellationToken.None);
try
......@@ -214,7 +214,7 @@ namespace Titanium.Web.Proxy
// create new connection to server.
// If we detected that client tunnel CONNECTs without SSL by checking for empty client hello then
// this connection should not be HTTPS.
var connection = await getServerConnection(connectArgs,
var connection = await tcpConnectionFactory.GetServerConnection(this, connectArgs,
true, SslExtensions.Http2ProtocolAsList, true, cancellationToken);
try
......@@ -284,7 +284,7 @@ namespace Titanium.Web.Proxy
throw new Exception($"HTTP/2 Protocol violation. Empty string expected, '{line}' received");
}
var connection = await getServerConnection(connectArgs, true,
var connection = await tcpConnectionFactory.GetServerConnection(this, connectArgs, true,
SslExtensions.Http2ProtocolAsList, true,
cancellationToken);
try
......
......@@ -9,6 +9,7 @@ using System.Text;
using System.Threading;
using System.Threading.Tasks;
using StreamExtended.Network;
using Titanium.Web.Proxy.EventArguments;
using Titanium.Web.Proxy.Extensions;
using Titanium.Web.Proxy.Helpers;
using Titanium.Web.Proxy.Http;
......@@ -69,6 +70,91 @@ namespace Titanium.Web.Proxy.Network.Tcp
return cacheKeyBuilder.ToString();
}
/// <summary>
/// Gets the connection cache key.
/// </summary>
/// <param name="args">The session event arguments.</param>
/// <param name="applicationProtocol"></param>
/// <returns></returns>
internal async Task<string> GetConnectionCacheKey(ProxyServer server, SessionEventArgsBase args,
SslApplicationProtocol applicationProtocol)
{
List<SslApplicationProtocol> applicationProtocols = null;
if (applicationProtocol != default)
{
applicationProtocols = new List<SslApplicationProtocol> { applicationProtocol };
}
ExternalProxy customUpStreamProxy = null;
bool isHttps = args.IsHttps;
if (server.GetCustomUpStreamProxyFunc != null)
{
customUpStreamProxy = await server.GetCustomUpStreamProxyFunc(args);
}
args.CustomUpStreamProxyUsed = customUpStreamProxy;
return GetConnectionCacheKey(
args.WebSession.Request.RequestUri.Host,
args.WebSession.Request.RequestUri.Port,
isHttps, applicationProtocols,
server, args.WebSession.UpStreamEndPoint ?? server.UpStreamEndPoint,
customUpStreamProxy ?? (isHttps ? server.UpStreamHttpsProxy : server.UpStreamHttpProxy));
}
/// <summary>
/// Create a server connection.
/// </summary>
/// <param name="args">The session event arguments.</param>
/// <param name="isConnect">Is this a CONNECT request.</param>
/// <param name="applicationProtocol"></param>
/// <param name="cancellationToken">The cancellation token for this async task.</param>
/// <returns></returns>
internal Task<TcpServerConnection> GetServerConnection(ProxyServer server, SessionEventArgsBase args, bool isConnect,
SslApplicationProtocol applicationProtocol, bool noCache, CancellationToken cancellationToken)
{
List<SslApplicationProtocol> applicationProtocols = null;
if (applicationProtocol != default)
{
applicationProtocols = new List<SslApplicationProtocol> { applicationProtocol };
}
return GetServerConnection(server, args, isConnect, applicationProtocols, noCache, cancellationToken);
}
/// <summary>
/// Create a server connection.
/// </summary>
/// <param name="args">The session event arguments.</param>
/// <param name="isConnect">Is this a CONNECT request.</param>
/// <param name="applicationProtocols"></param>
/// <param name="cancellationToken">The cancellation token for this async task.</param>
/// <returns></returns>
internal async Task<TcpServerConnection> GetServerConnection(ProxyServer server, SessionEventArgsBase args, bool isConnect,
List<SslApplicationProtocol> applicationProtocols, bool noCache, CancellationToken cancellationToken)
{
ExternalProxy customUpStreamProxy = null;
bool isHttps = args.IsHttps;
if (server.GetCustomUpStreamProxyFunc != null)
{
customUpStreamProxy = await server.GetCustomUpStreamProxyFunc(args);
}
args.CustomUpStreamProxyUsed = customUpStreamProxy;
return await GetServerConnection(
args.WebSession.Request.RequestUri.Host,
args.WebSession.Request.RequestUri.Port,
args.WebSession.Request.HttpVersion,
isHttps, applicationProtocols, isConnect,
server, args.WebSession.UpStreamEndPoint ?? server.UpStreamEndPoint,
customUpStreamProxy ?? (isHttps ? server.UpStreamHttpsProxy : server.UpStreamHttpProxy),
noCache, cancellationToken);
}
/// <summary>
/// Gets a TCP connection to server from connection pool.
/// </summary>
......@@ -84,7 +170,7 @@ namespace Titanium.Web.Proxy.Network.Tcp
/// <param name="noCache">Not from cache/create new connection.</param>
/// <param name="cancellationToken">The cancellation token for this async task.</param>
/// <returns></returns>
internal async Task<TcpServerConnection> GetClient(string remoteHostName, int remotePort,
internal async Task<TcpServerConnection> GetServerConnection(string remoteHostName, int remotePort,
Version httpVersion, bool isHttps, List<SslApplicationProtocol> applicationProtocols, bool isConnect,
ProxyServer proxyServer, IPEndPoint upStreamEndPoint, ExternalProxy externalProxy,
bool noCache, CancellationToken cancellationToken)
......@@ -116,7 +202,7 @@ namespace Titanium.Web.Proxy.Network.Tcp
}
}
var connection = await createClient(remoteHostName, remotePort, httpVersion, isHttps,
var connection = await createServerConnection(remoteHostName, remotePort, httpVersion, isHttps,
applicationProtocols, isConnect, proxyServer, upStreamEndPoint, externalProxy, cancellationToken);
connection.CacheKey = cacheKey;
......@@ -124,115 +210,6 @@ namespace Titanium.Web.Proxy.Network.Tcp
return connection;
}
/// <summary>
/// Release connection back to cache.
/// </summary>
/// <param name="connection">The Tcp server connection to return.</param>
/// <param name="close">Should we just close the connection instead of reusing?</param>
internal async Task Release(TcpServerConnection connection, bool close = false)
{
if (connection == null)
{
return;
}
if (close || connection.IsWinAuthenticated || !server.EnableConnectionPool)
{
disposalBag.Add(connection);
return;
}
connection.LastAccess = DateTime.Now;
try
{
await @lock.WaitAsync();
while (true)
{
if (cache.TryGetValue(connection.CacheKey, out var existingConnections))
{
while (existingConnections.Count >= server.MaxCachedConnections)
{
if (existingConnections.TryDequeue(out var staleConnection))
{
disposalBag.Add(staleConnection);
}
}
existingConnections.Enqueue(connection);
break;
}
if (cache.TryAdd(connection.CacheKey,
new ConcurrentQueue<TcpServerConnection>(new[] { connection })))
{
break;
}
}
}
finally
{
@lock.Release();
}
}
private async Task clearOutdatedConnections()
{
while (runCleanUpTask)
{
foreach (var item in cache)
{
var queue = item.Value;
while (queue.Count > 0)
{
if (queue.TryDequeue(out var connection))
{
var cutOff = DateTime.Now.AddSeconds(-1 * server.ConnectionTimeOutSeconds);
if (!server.EnableConnectionPool
|| connection.LastAccess < cutOff)
{
disposalBag.Add(connection);
continue;
}
queue.Enqueue(connection);
break;
}
}
}
try
{
await @lock.WaitAsync();
//clear empty queues
var emptyKeys = cache.Where(x => x.Value.Count == 0).Select(x => x.Key).ToList();
foreach (string key in emptyKeys)
{
cache.TryRemove(key, out var _);
}
}
finally
{
@lock.Release();
}
while (!disposalBag.IsEmpty)
{
if (disposalBag.TryTake(out var connection))
{
connection?.Dispose();
}
}
//cleanup every 3 seconds by default
await Task.Delay(1000 * 3);
}
}
/// <summary>
/// Creates a TCP connection to server
/// </summary>
......@@ -247,7 +224,7 @@ namespace Titanium.Web.Proxy.Network.Tcp
/// <param name="externalProxy">The external proxy to make request via.</param>
/// <param name="cancellationToken">The cancellation token for this async task.</param>
/// <returns></returns>
private async Task<TcpServerConnection> createClient(string remoteHostName, int remotePort,
private async Task<TcpServerConnection> createServerConnection(string remoteHostName, int remotePort,
Version httpVersion, bool isHttps, List<SslApplicationProtocol> applicationProtocols, bool isConnect,
ProxyServer proxyServer, IPEndPoint upStreamEndPoint, ExternalProxy externalProxy,
CancellationToken cancellationToken)
......@@ -370,6 +347,117 @@ namespace Titanium.Web.Proxy.Network.Tcp
Version = httpVersion
};
}
/// <summary>
/// Release connection back to cache.
/// </summary>
/// <param name="connection">The Tcp server connection to return.</param>
/// <param name="close">Should we just close the connection instead of reusing?</param>
internal async Task Release(TcpServerConnection connection, bool close = false)
{
if (connection == null)
{
return;
}
if (close || connection.IsWinAuthenticated || !server.EnableConnectionPool)
{
disposalBag.Add(connection);
return;
}
connection.LastAccess = DateTime.Now;
try
{
await @lock.WaitAsync();
while (true)
{
if (cache.TryGetValue(connection.CacheKey, out var existingConnections))
{
while (existingConnections.Count >= server.MaxCachedConnections)
{
if (existingConnections.TryDequeue(out var staleConnection))
{
disposalBag.Add(staleConnection);
}
}
existingConnections.Enqueue(connection);
break;
}
if (cache.TryAdd(connection.CacheKey,
new ConcurrentQueue<TcpServerConnection>(new[] { connection })))
{
break;
}
}
}
finally
{
@lock.Release();
}
}
private async Task clearOutdatedConnections()
{
while (runCleanUpTask)
{
foreach (var item in cache)
{
var queue = item.Value;
while (queue.Count > 0)
{
if (queue.TryDequeue(out var connection))
{
var cutOff = DateTime.Now.AddSeconds(-1 * server.ConnectionTimeOutSeconds);
if (!server.EnableConnectionPool
|| connection.LastAccess < cutOff)
{
disposalBag.Add(connection);
continue;
}
queue.Enqueue(connection);
break;
}
}
}
try
{
await @lock.WaitAsync();
//clear empty queues
var emptyKeys = cache.Where(x => x.Value.Count == 0).Select(x => x.Key).ToList();
foreach (string key in emptyKeys)
{
cache.TryRemove(key, out var _);
}
}
finally
{
@lock.Release();
}
while (!disposalBag.IsEmpty)
{
if (disposalBag.TryTake(out var connection))
{
connection?.Dispose();
}
}
//cleanup every 3 seconds by default
await Task.Delay(1000 * 3);
}
}
/// <summary>
/// Check if a TcpClient is good to be used.
/// This only checks if send is working so local socket is still connected.
......
......@@ -196,7 +196,7 @@ namespace Titanium.Web.Proxy
// only gets hit when connection pool is disabled.
// or when prefetch task has a unexpectedly different connection.
if (connection != null
&& (await getConnectionCacheKey(args,
&& (await tcpConnectionFactory.GetConnectionCacheKey(this, args,
clientConnection.NegotiatedApplicationProtocol)
!= connection.CacheKey))
{
......@@ -205,7 +205,8 @@ namespace Titanium.Web.Proxy
}
//a connection generator task with captured parameters via closure.
Func<Task<TcpServerConnection>> generator = () => getServerConnection(args, false,
Func<Task<TcpServerConnection>> generator = () =>
tcpConnectionFactory.GetServerConnection(this, args, false,
clientConnection.NegotiatedApplicationProtocol,
false, cancellationToken);
......@@ -373,62 +374,6 @@ namespace Titanium.Web.Proxy
}
}
/// <summary>
/// Handle upgrade to websocket
/// </summary>
private async Task handleWebSocketUpgrade(string httpCmd,
SessionEventArgs args, Request request, Response response,
CustomBufferedStream clientStream, HttpResponseWriter clientStreamWriter,
TcpServerConnection serverConnection,
CancellationTokenSource cancellationTokenSource, CancellationToken cancellationToken)
{
// prepare the prefix content
await serverConnection.StreamWriter.WriteLineAsync(httpCmd, cancellationToken);
await serverConnection.StreamWriter.WriteHeadersAsync(request.Headers,
cancellationToken: cancellationToken);
string httpStatus;
try
{
httpStatus = await serverConnection.Stream.ReadLineAsync(cancellationToken);
if (httpStatus == null)
{
throw new ServerConnectionException("Server connection was closed.");
}
}
catch (Exception e) when (!(e is ServerConnectionException))
{
throw new ServerConnectionException("Server connection was closed.", e);
}
Response.ParseResponseLine(httpStatus, out var responseVersion,
out int responseStatusCode,
out string responseStatusDescription);
response.HttpVersion = responseVersion;
response.StatusCode = responseStatusCode;
response.StatusDescription = responseStatusDescription;
await HeaderParser.ReadHeaders(serverConnection.Stream, response.Headers,
cancellationToken);
if (!args.IsTransparent)
{
await clientStreamWriter.WriteResponseAsync(response,
cancellationToken: cancellationToken);
}
// If user requested call back then do it
if (!args.WebSession.Response.Locked)
{
await invokeBeforeResponse(args);
}
await TcpHelper.SendRaw(clientStream, serverConnection.Stream, BufferPool, BufferSize,
(buffer, offset, count) => { args.OnDataSent(buffer, offset, count); },
(buffer, offset, count) => { args.OnDataReceived(buffer, offset, count); },
cancellationTokenSource, ExceptionFunc);
}
/// <summary>
/// Prepare the request headers so that we can avoid encodings not parsable by this proxy
/// </summary>
......@@ -455,92 +400,6 @@ namespace Titanium.Web.Proxy
requestHeaders.FixProxyHeaders();
}
/// <summary>
/// Gets the connection cache key.
/// </summary>
/// <param name="args">The session event arguments.</param>
/// <param name="applicationProtocol"></param>
/// <returns></returns>
private async Task<string> getConnectionCacheKey(SessionEventArgsBase args,
SslApplicationProtocol applicationProtocol)
{
List<SslApplicationProtocol> applicationProtocols = null;
if (applicationProtocol != default)
{
applicationProtocols = new List<SslApplicationProtocol> { applicationProtocol };
}
ExternalProxy customUpStreamProxy = null;
bool isHttps = args.IsHttps;
if (GetCustomUpStreamProxyFunc != null)
{
customUpStreamProxy = await GetCustomUpStreamProxyFunc(args);
}
args.CustomUpStreamProxyUsed = customUpStreamProxy;
return tcpConnectionFactory.GetConnectionCacheKey(
args.WebSession.Request.RequestUri.Host,
args.WebSession.Request.RequestUri.Port,
isHttps, applicationProtocols,
this, args.WebSession.UpStreamEndPoint ?? UpStreamEndPoint,
customUpStreamProxy ?? (isHttps ? UpStreamHttpsProxy : UpStreamHttpProxy));
}
/// <summary>
/// Create a server connection.
/// </summary>
/// <param name="args">The session event arguments.</param>
/// <param name="isConnect">Is this a CONNECT request.</param>
/// <param name="applicationProtocol"></param>
/// <param name="cancellationToken">The cancellation token for this async task.</param>
/// <returns></returns>
private Task<TcpServerConnection> getServerConnection(SessionEventArgsBase args, bool isConnect,
SslApplicationProtocol applicationProtocol, bool noCache, CancellationToken cancellationToken)
{
List<SslApplicationProtocol> applicationProtocols = null;
if (applicationProtocol != default)
{
applicationProtocols = new List<SslApplicationProtocol> { applicationProtocol };
}
return getServerConnection(args, isConnect, applicationProtocols, noCache, cancellationToken);
}
/// <summary>
/// Create a server connection.
/// </summary>
/// <param name="args">The session event arguments.</param>
/// <param name="isConnect">Is this a CONNECT request.</param>
/// <param name="applicationProtocols"></param>
/// <param name="cancellationToken">The cancellation token for this async task.</param>
/// <returns></returns>
private async Task<TcpServerConnection> getServerConnection(SessionEventArgsBase args, bool isConnect,
List<SslApplicationProtocol> applicationProtocols, bool noCache, CancellationToken cancellationToken)
{
ExternalProxy customUpStreamProxy = null;
bool isHttps = args.IsHttps;
if (GetCustomUpStreamProxyFunc != null)
{
customUpStreamProxy = await GetCustomUpStreamProxyFunc(args);
}
args.CustomUpStreamProxyUsed = customUpStreamProxy;
return await tcpConnectionFactory.GetClient(
args.WebSession.Request.RequestUri.Host,
args.WebSession.Request.RequestUri.Port,
args.WebSession.Request.HttpVersion,
isHttps, applicationProtocols, isConnect,
this, args.WebSession.UpStreamEndPoint ?? UpStreamEndPoint,
customUpStreamProxy ?? (isHttps ? UpStreamHttpsProxy : UpStreamHttpProxy),
noCache, cancellationToken);
}
/// <summary>
/// Invoke before request handler if it is set.
/// </summary>
......
......@@ -65,7 +65,7 @@ namespace Titanium.Web.Proxy
{
//don't pass cancellation token here
//it could cause floating server connections when client exits
prefetchConnectionTask = tcpConnectionFactory.GetClient(httpsHostName, endPoint.Port,
prefetchConnectionTask = tcpConnectionFactory.GetServerConnection(httpsHostName, endPoint.Port,
null, true, null, false, this,
UpStreamEndPoint, UpStreamHttpsProxy, false, CancellationToken.None);
......@@ -96,7 +96,7 @@ namespace Titanium.Web.Proxy
}
else
{
var connection = await tcpConnectionFactory.GetClient(httpsHostName, endPoint.Port,
var connection = await tcpConnectionFactory.GetServerConnection(httpsHostName, endPoint.Port,
null, false, null,
true, this, UpStreamEndPoint, UpStreamHttpsProxy, true, cancellationToken);
......
using StreamExtended.Network;
using System;
using System.Threading;
using System.Threading.Tasks;
using Titanium.Web.Proxy.EventArguments;
using Titanium.Web.Proxy.Exceptions;
using Titanium.Web.Proxy.Helpers;
using Titanium.Web.Proxy.Http;
using Titanium.Web.Proxy.Network.Tcp;
namespace Titanium.Web.Proxy
{
public partial class ProxyServer
{
/// <summary>
/// Handle upgrade to websocket
/// </summary>
private async Task handleWebSocketUpgrade(string httpCmd,
SessionEventArgs args, Request request, Response response,
CustomBufferedStream clientStream, HttpResponseWriter clientStreamWriter,
TcpServerConnection serverConnection,
CancellationTokenSource cancellationTokenSource, CancellationToken cancellationToken)
{
// prepare the prefix content
await serverConnection.StreamWriter.WriteLineAsync(httpCmd, cancellationToken);
await serverConnection.StreamWriter.WriteHeadersAsync(request.Headers,
cancellationToken: cancellationToken);
string httpStatus;
try
{
httpStatus = await serverConnection.Stream.ReadLineAsync(cancellationToken);
if (httpStatus == null)
{
throw new ServerConnectionException("Server connection was closed.");
}
}
catch (Exception e) when (!(e is ServerConnectionException))
{
throw new ServerConnectionException("Server connection was closed.", e);
}
Response.ParseResponseLine(httpStatus, out var responseVersion,
out int responseStatusCode,
out string responseStatusDescription);
response.HttpVersion = responseVersion;
response.StatusCode = responseStatusCode;
response.StatusDescription = responseStatusDescription;
await HeaderParser.ReadHeaders(serverConnection.Stream, response.Headers,
cancellationToken);
if (!args.IsTransparent)
{
await clientStreamWriter.WriteResponseAsync(response,
cancellationToken: cancellationToken);
}
// If user requested call back then do it
if (!args.WebSession.Response.Locked)
{
await invokeBeforeResponse(args);
}
await TcpHelper.SendRaw(clientStream, serverConnection.Stream, BufferPool, BufferSize,
(buffer, offset, count) => { args.OnDataSent(buffer, offset, count); },
(buffer, offset, count) => { args.OnDataReceived(buffer, offset, count); },
cancellationTokenSource, ExceptionFunc);
}
}
}
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