Skip to content
This repository was archived by the owner on Jul 9, 2023. It is now read-only.

Commit 02ea4ae

Browse files
justcoding121justcoding121
authored
Merge pull request #445 from justcoding121/master
Use polly for retries
2 parents 17a1aab + fc3a930 commit 02ea4ae

File tree

11 files changed

+152
-118
lines changed

11 files changed

+152
-118
lines changed

‎Examples/Titanium.Web.Proxy.Examples.Basic/ProxyTestController.cs‎

Lines changed: 41 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
using System.Collections.Generic;
44
using System.Net;
55
using System.Net.Security;
6+
using System.Threading;
67
using System.Threading.Tasks;
78
using Titanium.Web.Proxy.EventArguments;
89
using Titanium.Web.Proxy.Exceptions;
@@ -14,7 +15,7 @@ namespace Titanium.Web.Proxy.Examples.Basic
1415
{
1516
public class ProxyTestController
1617
{
17-
private readonly object lockObj = new object();
18+
private readonly SemaphoreSlim @lock = new SemaphoreSlim(1);
1819

1920
private readonly ProxyServer proxyServer;
2021

@@ -23,34 +24,41 @@ public class ProxyTestController
2324
public ProxyTestController()
2425
{
2526
proxyServer = new ProxyServer();
26-
//proxyServer.EnableConnectionPool = false;
27+
proxyServer.EnableConnectionPool = true;
2728
// generate root certificate without storing it in file system
2829
//proxyServer.CertificateManager.CreateRootCertificate(false);
2930

3031
//proxyServer.CertificateManager.TrustRootCertificate();
3132
//proxyServer.CertificateManager.TrustRootCertificateAsAdmin();
3233

33-
proxyServer.ExceptionFunc = exception =>
34+
proxyServer.ExceptionFunc = async exception =>
3435
{
35-
lock (lockObj)
36+
await @lock.WaitAsync();
37+
38+
try
39+
{
40+
var color = Console.ForegroundColor;
41+
Console.ForegroundColor = ConsoleColor.Red;
42+
if (exception is ProxyHttpException phex)
43+
{
44+
Console.WriteLine(exception.Message + ": " + phex.InnerException?.Message);
45+
}
46+
else
47+
{
48+
Console.WriteLine(exception.Message);
49+
}
50+
51+
Console.ForegroundColor = color;
52+
53+
}
54+
finally
3655
{
37-
var color = Console.ForegroundColor;
38-
Console.ForegroundColor = ConsoleColor.Red;
39-
if (exception is ProxyHttpException phex)
40-
{
41-
Console.WriteLine(exception.Message + ": " + phex.InnerException?.Message);
42-
}
43-
else
44-
{
45-
Console.WriteLine(exception.Message);
46-
}
47-
48-
Console.ForegroundColor = color;
56+
@lock.Release();
4957
}
5058
};
5159
proxyServer.ForwardToUpstreamGateway = true;
5260
proxyServer.CertificateManager.SaveFakeCertificates = true;
53-
61+
5462
// optionally set the Certificate Engine
5563
// Under Mono or Non-Windows runtimes only BouncyCastle will be supported
5664
//proxyServer.CertificateManager.CertificateEngine = Network.CertificateEngine.BouncyCastle;
@@ -122,15 +130,15 @@ public void Stop()
122130
proxyServer.ClientCertificateSelectionCallback -= OnCertificateSelection;
123131

124132
proxyServer.Stop();
125-
133+
126134
// remove the generated certificates
127135
//proxyServer.CertificateManager.RemoveTrustedRootCertificates();
128136
}
129137

130138
private async Task OnBeforeTunnelConnectRequest(object sender, TunnelConnectSessionEventArgs e)
131139
{
132140
string hostname = e.WebSession.Request.RequestUri.Host;
133-
WriteToConsole("Tunnel to: " + hostname);
141+
await WriteToConsole("Tunnel to: " + hostname);
134142

135143
if (hostname.Contains("dropbox.com"))
136144
{
@@ -148,8 +156,8 @@ private async Task OnBeforeTunnelConnectResponse(object sender, TunnelConnectSes
148156
// intecept & cancel redirect or update requests
149157
private async Task OnRequest(object sender, SessionEventArgs e)
150158
{
151-
WriteToConsole("Active Client Connections:" + ((ProxyServer)sender).ClientConnectionCount);
152-
WriteToConsole(e.WebSession.Request.Url);
159+
await WriteToConsole("Active Client Connections:" + ((ProxyServer)sender).ClientConnectionCount);
160+
await WriteToConsole(e.WebSession.Request.Url);
153161

154162
// store it in the UserData property
155163
// It can be a simple integer, Guid, or any type
@@ -187,19 +195,19 @@ private async Task OnRequest(object sender, SessionEventArgs e)
187195
}
188196

189197
// Modify response
190-
private void MultipartRequestPartSent(object sender, MultipartRequestPartSentEventArgs e)
198+
private async Task MultipartRequestPartSent(object sender, MultipartRequestPartSentEventArgs e)
191199
{
192200
var session = (SessionEventArgs)sender;
193-
WriteToConsole("Multipart form data headers:");
201+
await WriteToConsole("Multipart form data headers:");
194202
foreach (var header in e.Headers)
195203
{
196-
WriteToConsole(header.ToString());
204+
await WriteToConsole(header.ToString());
197205
}
198206
}
199207

200208
private async Task OnResponse(object sender, SessionEventArgs e)
201209
{
202-
WriteToConsole("Active Server Connections:" + ((ProxyServer)sender).ServerConnectionCount);
210+
await WriteToConsole("Active Server Connections:" + ((ProxyServer)sender).ServerConnectionCount);
203211

204212
string ext = System.IO.Path.GetExtension(e.WebSession.Request.RequestUri.AbsolutePath);
205213

@@ -271,12 +279,18 @@ public Task OnCertificateSelection(object sender, CertificateSelectionEventArgs
271279
return Task.FromResult(0);
272280
}
273281

274-
private void WriteToConsole(string message)
282+
private async Task WriteToConsole(string message)
275283
{
276-
lock (lockObj)
284+
await @lock.WaitAsync();
285+
286+
try
277287
{
278288
Console.WriteLine(message);
279289
}
290+
finally
291+
{
292+
@lock.Release();
293+
}
280294
}
281295

282296
///// <summary>

‎Titanium.Web.Proxy/ExplicitClientHandler.cs‎

Lines changed: 24 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using System;
2+
using System.Collections.Generic;
23
using System.IO;
34
using System.Net;
45
using System.Net.Security;
@@ -140,17 +141,17 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.WebSession.Response,
140141
// test server HTTP/2 support
141142
// todo: this is a hack, because Titanium does not support HTTP protocol changing currently
142143
var connection = await getServerConnection(connectArgs, true,
143-
SslExtensions.Http2ProtocolAsList, cancellationToken);
144+
SslExtensions.Http2ProtocolAsList, true, cancellationToken);
144145

145146
http2Supported = connection.NegotiatedApplicationProtocol == SslApplicationProtocol.Http2;
146147

147148
//release connection back to pool intead of closing when connection pool is enabled.
148-
await tcpConnectionFactory.Release(connection);
149+
await tcpConnectionFactory.Release(connection, true);
149150
}
150151

151152
SslStream sslStream = null;
152153
prefetchConnectionTask = getServerConnection(connectArgs, true,
153-
null, cancellationToken);
154+
null, false, cancellationToken);
154155
try
155156
{
156157
sslStream = new SslStream(clientStream);
@@ -209,8 +210,8 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.WebSession.Response,
209210
// create new connection to server.
210211
// If we detected that client tunnel CONNECTs without SSL by checking for empty client hello then
211212
// this connection should not be HTTPS.
212-
var connection = await getServerConnection(connectArgs, true,
213-
null, cancellationToken);
213+
var connection = await getServerConnection(connectArgs,
214+
true, SslExtensions.Http2ProtocolAsList, true, cancellationToken);
214215

215216
try
216217
{
@@ -224,31 +225,32 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.WebSession.Response,
224225

225226
try
226227
{
227-
// clientStream.Available sbould be at most BufferSize because it is using the same buffer size
228228
await clientStream.ReadAsync(data, 0, available, cancellationToken);
229-
await connection.StreamWriter.WriteAsync(data, 0, available, true,
230-
cancellationToken);
229+
// clientStream.Available sbould be at most BufferSize because it is using the same buffer size
230+
await connection.StreamWriter.WriteAsync(data, 0, available, true, cancellationToken);
231231
}
232232
finally
233233
{
234234
BufferPool.ReturnBuffer(data);
235235
}
236236
}
237237

238-
var serverHelloInfo =
239-
await SslTools.PeekServerHello(connection.Stream, BufferPool, cancellationToken);
238+
var serverHelloInfo = await SslTools.PeekServerHello(connection.Stream, BufferPool, cancellationToken);
240239
((ConnectResponse)connectArgs.WebSession.Response).ServerHelloInfo = serverHelloInfo;
241240
}
242241

243242
await TcpHelper.SendRaw(clientStream, connection.Stream, BufferPool, BufferSize,
244243
(buffer, offset, count) => { connectArgs.OnDataSent(buffer, offset, count); },
245244
(buffer, offset, count) => { connectArgs.OnDataReceived(buffer, offset, count); },
246245
connectArgs.CancellationTokenSource, ExceptionFunc);
246+
247+
247248
}
248249
finally
249250
{
250251
await tcpConnectionFactory.Release(connection, true);
251252
}
253+
calledRequestHandler = true;
252254
return;
253255
}
254256
}
@@ -278,29 +280,30 @@ await TcpHelper.SendRaw(clientStream, connection.Stream, BufferPool, BufferSize,
278280
throw new Exception($"HTTP/2 Protocol violation. Empty string expected, '{line}' received");
279281
}
280282

281-
// create new connection
282-
var connection = await getServerConnection(connectArgs, true, SslExtensions.Http2ProtocolAsList,
283-
cancellationToken);
283+
var connection = await getServerConnection(connectArgs, true,
284+
SslExtensions.Http2ProtocolAsList, true,
285+
cancellationToken);
284286
try
285287
{
286288
await connection.StreamWriter.WriteLineAsync("PRI * HTTP/2.0", cancellationToken);
287289
await connection.StreamWriter.WriteLineAsync(cancellationToken);
288290
await connection.StreamWriter.WriteLineAsync("SM", cancellationToken);
289291
await connection.StreamWriter.WriteLineAsync(cancellationToken);
290-
291292
#if NETCOREAPP2_1
292293
await Http2Helper.SendHttp2(clientStream, connection.Stream, BufferSize,
293294
(buffer, offset, count) => { connectArgs.OnDataSent(buffer, offset, count); },
294295
(buffer, offset, count) => { connectArgs.OnDataReceived(buffer, offset, count); },
295296
connectArgs.CancellationTokenSource, clientConnection.Id, ExceptionFunc);
296297
#endif
298+
297299
}
298300
finally
299301
{
300302
await tcpConnectionFactory.Release(connection, true);
301303
}
302304
}
303305
}
306+
304307
calledRequestHandler = true;
305308
// Now create the request
306309
await handleHttpSessionRequest(endPoint, clientConnection, clientStream, clientStreamWriter,
@@ -328,19 +331,19 @@ await handleHttpSessionRequest(endPoint, clientConnection, clientStream, clientS
328331
}
329332
finally
330333
{
331-
if (!calledRequestHandler
332-
&& prefetchConnectionTask != null)
333-
{
334-
var connection = await prefetchConnectionTask;
335-
await tcpConnectionFactory.Release(connection, closeServerConnection);
336-
}
337-
338334
clientStream.Dispose();
339335

340336
if (!cancellationTokenSource.IsCancellationRequested)
341337
{
342338
cancellationTokenSource.Cancel();
343339
}
340+
341+
if (!calledRequestHandler
342+
&& prefetchConnectionTask != null)
343+
{
344+
var connection = await prefetchConnectionTask;
345+
await tcpConnectionFactory.Release(connection, closeServerConnection);
346+
}
344347
}
345348
}
346349
}

‎Titanium.Web.Proxy/Network/Tcp/TcpConnectionFactory.cs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,18 +81,19 @@ internal string GetConnectionCacheKey(string remoteHostName, int remotePort,
8181
/// <param name="proxyServer">The current ProxyServer instance.</param>
8282
/// <param name="upStreamEndPoint">The local upstream endpoint to make request via.</param>
8383
/// <param name="externalProxy">The external proxy to make request via.</param>
84+
/// <param name="noCache">Not from cache/create new connection.</param>
8485
/// <param name="cancellationToken">The cancellation token for this async task.</param>
8586
/// <returns></returns>
8687
internal async Task<TcpServerConnection> GetClient(string remoteHostName, int remotePort,
8788
Version httpVersion, bool isHttps, List<SslApplicationProtocol> applicationProtocols, bool isConnect,
8889
ProxyServer proxyServer, IPEndPoint upStreamEndPoint, ExternalProxy externalProxy,
89-
CancellationToken cancellationToken)
90+
bool noCache, CancellationToken cancellationToken)
9091
{
9192
var cacheKey = GetConnectionCacheKey(remoteHostName, remotePort,
9293
isHttps, applicationProtocols,
9394
proxyServer, upStreamEndPoint, externalProxy);
9495

95-
if (proxyServer.EnableConnectionPool)
96+
if (proxyServer.EnableConnectionPool && !noCache)
9697
{
9798
if (cache.TryGetValue(cacheKey, out var existingConnections))
9899
{

‎Titanium.Web.Proxy/ProxyServer.cs‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
using System.Security.Cryptography.X509Certificates;
88
using System.Threading;
99
using System.Threading.Tasks;
10+
using Polly;
1011
using StreamExtended;
1112
using StreamExtended.Network;
1213
using Titanium.Web.Proxy.EventArguments;
@@ -123,6 +124,9 @@ public ProxyServer(string rootCertificateName, string rootCertificateIssuerName,
123124
/// </summary>
124125
private SystemProxyManager systemProxySettingsManager { get; }
125126

127+
//Number of exception retries when connection pool is enabled.
128+
private int retries => EnableConnectionPool ? MaxCachedConnections + 1 : 0;
129+
126130
/// <summary>
127131
/// Is the proxy currently running?
128132
/// </summary>
@@ -197,7 +201,7 @@ public ProxyServer(string rootCertificateName, string rootCertificateIssuerName,
197201
/// Realm used during Proxy Basic Authentication.
198202
/// </summary>
199203
public string ProxyRealm { get; set; } = "TitaniumProxy";
200-
204+
201205
/// <summary>
202206
/// List of supported Ssl versions.
203207
/// </summary>
@@ -252,7 +256,8 @@ public ProxyServer(string rootCertificateName, string rootCertificateIssuerName,
252256
public ExceptionHandler ExceptionFunc
253257
{
254258
get => exceptionFunc ?? defaultExceptionFunc;
255-
set {
259+
set
260+
{
256261
exceptionFunc = value;
257262
CertificateManager.ExceptionFunc = value;
258263
}
@@ -274,7 +279,7 @@ public void Dispose()
274279
{
275280
Stop();
276281
}
277-
282+
278283
CertificateManager?.Dispose();
279284
BufferPool?.Dispose();
280285
}
@@ -800,5 +805,26 @@ internal async Task InvokeConnectionCreateEvent(TcpClient client, bool isClientC
800805
await OnServerConnectionCreate.InvokeAsync(this, client, ExceptionFunc);
801806
}
802807
}
808+
809+
/// <summary>
810+
/// Connection retry policy when using connection pool.
811+
/// </summary>
812+
private Policy retryPolicy<T>() where T : Exception
813+
{
814+
return Policy.Handle<T>()
815+
.RetryAsync(retries,
816+
onRetryAsync: async (ex, i, context) =>
817+
{
818+
if (context.ContainsKey("connection"))
819+
{
820+
//close connection on error
821+
var connection = (TcpServerConnection)context["connection"];
822+
await tcpConnectionFactory.Release(connection, true);
823+
context.Remove("connection");
824+
}
825+
826+
});
827+
828+
}
803829
}
804830
}

0 commit comments

Comments
 (0)