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

Commit e6a201e

Browse files
honfikahonfika
committed
another race condition fix
1 parent 42ac058 commit e6a201e

File tree

5 files changed

+77
-44
lines changed

5 files changed

+77
-44
lines changed

‎src/Titanium.Web.Proxy/EventArguments/SessionEventArgs.cs

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,6 @@ public class SessionEventArgs : SessionEventArgsBase
2828
/// </summary>
2929
private bool reRequest;
3030

31-
internal TaskCompletionSource<bool> ReadHttp2BodyTaskCompletionSource;
32-
3331
/// <summary>
3432
/// Constructor to initialize the proxy
3533
/// </summary>
@@ -99,7 +97,7 @@ private async Task readRequestBodyAsync(CancellationToken cancellationToken)
9997
request.ReadHttp2BodyTaskCompletionSource = tcs;
10098

10199
// signal to HTTP/2 copy frame method to continue
102-
ReadHttp2BodyTaskCompletionSource.SetResult(true);
100+
request.ReadHttp2BeforeHandlerTaskCompletionSource.SetResult(true);
103101

104102
await tcs.Task;
105103
}
@@ -163,9 +161,9 @@ private async Task readResponseBodyAsync(CancellationToken cancellationToken)
163161

164162
var tcs = new TaskCompletionSource<bool>();
165163
response.ReadHttp2BodyTaskCompletionSource = tcs;
166-
164+
167165
// signal to HTTP/2 copy frame method to continue
168-
ReadHttp2BodyTaskCompletionSource.SetResult(true);
166+
response.ReadHttp2BeforeHandlerTaskCompletionSource.SetResult(true);
169167

170168
await tcs.Task;
171169
}

‎src/Titanium.Web.Proxy/EventArguments/SessionEventArgsBase.cs

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -55,25 +55,7 @@ protected SessionEventArgsBase(ProxyServer server, ProxyEndPoint endPoint,
5555
HttpClient = new HttpWebClient(request);
5656
LocalEndPoint = endPoint;
5757

58-
HttpClient.ProcessId = new Lazy<int>(() =>
59-
{
60-
if (RunTime.IsWindows)
61-
{
62-
var remoteEndPoint = ClientEndPoint;
63-
64-
// If client is localhost get the process id
65-
if (NetworkHelper.IsLocalIpAddress(remoteEndPoint.Address))
66-
{
67-
var ipVersion = endPoint.IpV6Enabled ? IpVersion.Ipv6 : IpVersion.Ipv4;
68-
return TcpHelper.GetProcessIdByLocalPort(ipVersion, remoteEndPoint.Port);
69-
}
70-
71-
// can't access process Id of remote request from remote machine
72-
return -1;
73-
}
74-
75-
throw new PlatformNotSupportedException();
76-
});
58+
HttpClient.ProcessId = new Lazy<int>(() => ProxyClient.Connection.GetProcessId(endPoint));
7759
}
7860

7961
/// <summary>

‎src/Titanium.Web.Proxy/Http/RequestResponseBase.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ public abstract class RequestResponseBase
5050
/// </summary>
5151
internal string OriginalContentEncoding { get; set; }
5252

53+
internal TaskCompletionSource<bool> ReadHttp2BeforeHandlerTaskCompletionSource;
54+
5355
internal TaskCompletionSource<bool> ReadHttp2BodyTaskCompletionSource;
5456

5557
internal MemoryStream Http2BodyData;

‎src/Titanium.Web.Proxy/Http2/Http2Helper.cs

Lines changed: 38 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,30 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
8989

9090
bool endStream = false;
9191

92-
//System.Diagnostics.Debug.WriteLine("CLIENT: " + isClient + ", STREAM: " + streamId + ", TYPE: " + type);
92+
SessionEventArgs args = null;
93+
RequestResponseBase rr = null;
94+
if (type == 0 || type == 1)
95+
{
96+
if (!sessions.TryGetValue(streamId, out args))
97+
{
98+
if (type == 0)
99+
{
100+
throw new ProxyHttpException("HTTP Body data received before any header frame.", null, args);
101+
}
102+
103+
if (!isClient)
104+
{
105+
throw new ProxyHttpException("HTTP Response received before any Request header frame.", null, args);
106+
}
107+
108+
args = sessionFactory();
109+
sessions.TryAdd(streamId, args);
110+
}
111+
112+
rr = isClient ? (RequestResponseBase)args.HttpClient.Request : args.HttpClient.Response;
113+
}
114+
115+
//System.Diagnostics.Debug.WriteLine("CONN: " + connectionId + ", CLIENT: " + isClient + ", STREAM: " + streamId + ", TYPE: " + type);
93116
if (type == 0 /* data */)
94117
{
95118
bool endStreamFlag = (flags & (int)Http2FrameFlag.EndStream) != 0;
@@ -98,12 +121,6 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
98121
endStream = true;
99122
}
100123

101-
if (!sessions.TryGetValue(streamId, out var args))
102-
{
103-
throw new ProxyHttpException("HTTP Body data received before any header frame.", null, args);
104-
}
105-
106-
var rr = isClient ? (RequestResponseBase)args.HttpClient.Request : args.HttpClient.Response;
107124
if (rr.ReadHttp2BodyTaskCompletionSource != null)
108125
{
109126
// Get body method was called in the "before" event handler
@@ -115,12 +132,18 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
115132
{
116133
rr.Body = data.ToArray();
117134
rr.IsBodyRead = true;
118-
rr.ReadHttp2BodyTaskCompletionSource.SetResult(true);
119135

136+
var tcs = rr.ReadHttp2BodyTaskCompletionSource;
120137
rr.ReadHttp2BodyTaskCompletionSource = null;
138+
139+
if (!tcs.Task.IsCompleted)
140+
{
141+
tcs.SetResult(true);
142+
}
143+
121144
rr.Http2BodyData = null;
122145
}
123-
}
146+
}
124147
}
125148
else if (type == 1 /*headers*/)
126149
{
@@ -150,12 +173,6 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
150173
dataLength -= buffer[0];
151174
}
152175

153-
if (!sessions.TryGetValue(streamId, out var args))
154-
{
155-
args = sessionFactory();
156-
sessions.TryAdd(streamId, args);
157-
}
158-
159176
var headerListener = new MyHeaderListener(
160177
(name, value) =>
161178
{
@@ -194,24 +211,25 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
194211

195212
if (endHeaders)
196213
{
197-
var handler = onBeforeRequestResponse(args);
198-
199214
var tcs = new TaskCompletionSource<bool>();
200-
args.ReadHttp2BodyTaskCompletionSource = tcs;
215+
rr.ReadHttp2BeforeHandlerTaskCompletionSource = tcs;
216+
217+
var handler = onBeforeRequestResponse(args);
201218

202-
if (handler == await Task.WhenAny(handler, tcs.Task))
219+
if (handler == await Task.WhenAny(tcs.Task, handler))
203220
{
221+
rr.ReadHttp2BeforeHandlerTaskCompletionSource = null;
204222
tcs.SetResult(true);
205223
}
206224

207-
var rr = isClient ? (RequestResponseBase)args.HttpClient.Request : args.HttpClient.Response;
208225
rr.Locked = true;
209226
}
210227
}
211228

212229
if (!isClient && endStream)
213230
{
214231
sessions.TryRemove(streamId, out _);
232+
//System.Diagnostics.Debug.WriteLine("REMOVED CONN: " + connectionId + ", CLIENT: " + isClient + ", STREAM: " + streamId + ", TYPE: " + type);
215233
}
216234

217235
// do not cancel the write operation

‎src/Titanium.Web.Proxy/Network/Tcp/TcpClientConnection.cs

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
using System.Net.Sockets;
88
using System.Threading.Tasks;
99
using Titanium.Web.Proxy.Extensions;
10+
using Titanium.Web.Proxy.Helpers;
11+
using Titanium.Web.Proxy.Models;
1012

1113
namespace Titanium.Web.Proxy.Network.Tcp
1214
{
@@ -34,11 +36,42 @@ internal TcpClientConnection(ProxyServer proxyServer, TcpClient tcpClient)
3436

3537
private readonly TcpClient tcpClient;
3638

39+
private int? processId;
40+
3741
public Stream GetStream()
3842
{
3943
return tcpClient.GetStream();
4044
}
4145

46+
public int GetProcessId(ProxyEndPoint endPoint)
47+
{
48+
if (processId.HasValue)
49+
{
50+
return processId.Value;
51+
}
52+
53+
if (RunTime.IsWindows)
54+
{
55+
var remoteEndPoint = (IPEndPoint)RemoteEndPoint;
56+
57+
// If client is localhost get the process id
58+
if (NetworkHelper.IsLocalIpAddress(remoteEndPoint.Address))
59+
{
60+
var ipVersion = endPoint.IpV6Enabled ? IpVersion.Ipv6 : IpVersion.Ipv4;
61+
processId = TcpHelper.GetProcessIdByLocalPort(ipVersion, remoteEndPoint.Port);
62+
}
63+
else
64+
{
65+
// can't access process Id of remote request from remote machine
66+
processId = -1;
67+
}
68+
69+
return processId.Value;
70+
}
71+
72+
throw new PlatformNotSupportedException();
73+
}
74+
4275
/// <summary>
4376
/// Dispose.
4477
/// </summary>

0 commit comments

Comments
 (0)