Commit dff59f07 authored by Honfika's avatar Honfika

allow to terminate the server connection in BeforeResponse (not ready yet)

parent ae80fbe1
...@@ -168,7 +168,7 @@ namespace Titanium.Web.Proxy.EventArguments ...@@ -168,7 +168,7 @@ namespace Titanium.Web.Proxy.EventArguments
//If not already read (not cached yet) //If not already read (not cached yet)
if (!request.IsBodyRead) if (!request.IsBodyRead)
{ {
var body = await ReadBodyAsync(ProxyClient.ClientStreamReader, true); var body = await ReadBodyAsync(true);
request.Body = body; request.Body = body;
//Now set the flag to true //Now set the flag to true
...@@ -243,7 +243,7 @@ namespace Titanium.Web.Proxy.EventArguments ...@@ -243,7 +243,7 @@ namespace Titanium.Web.Proxy.EventArguments
//If not already read (not cached yet) //If not already read (not cached yet)
if (!response.IsBodyRead) if (!response.IsBodyRead)
{ {
var body = await ReadBodyAsync(WebSession.ServerConnection.StreamReader, false); var body = await ReadBodyAsync(false);
response.Body = body; response.Body = body;
//Now set the flag to true //Now set the flag to true
...@@ -253,12 +253,21 @@ namespace Titanium.Web.Proxy.EventArguments ...@@ -253,12 +253,21 @@ namespace Titanium.Web.Proxy.EventArguments
} }
} }
private async Task<byte[]> ReadBodyAsync(CustomBinaryReader reader, bool isRequest) private async Task<byte[]> ReadBodyAsync(bool isRequest)
{ {
using (var bodyStream = new MemoryStream()) using (var bodyStream = new MemoryStream())
{ {
var writer = new HttpWriter(bodyStream, bufferSize); var writer = new HttpWriter(bodyStream, bufferSize);
await ReadBodyAsync(writer, reader, isRequest);
if (isRequest)
{
await CopyRequestBodyAsync(writer, TransformationMode.Uncompress);
}
else
{
await CopyResponseBodyAsync(writer, TransformationMode.Uncompress);
}
return bodyStream.ToArray(); return bodyStream.ToArray();
} }
} }
...@@ -271,22 +280,11 @@ namespace Titanium.Web.Proxy.EventArguments ...@@ -271,22 +280,11 @@ namespace Titanium.Web.Proxy.EventArguments
return; return;
} }
var reader = GetStreamReader(isRequest);
using (var bodyStream = new MemoryStream()) using (var bodyStream = new MemoryStream())
{ {
var writer = new HttpWriter(bodyStream, bufferSize); var writer = new HttpWriter(bodyStream, bufferSize);
await ReadBodyAsync(writer, reader, isRequest); await CopyBodyAsync(isRequest, writer, TransformationMode.None, null);
}
}
private Task ReadBodyAsync(HttpWriter writer, CustomBinaryReader reader, bool isRequest)
{
if (isRequest)
{
return CopyRequestBodyAsync(writer, TransformationMode.Uncompress);
} }
return CopyResponseBodyAsync(writer, TransformationMode.Uncompress);
} }
/// <summary> /// <summary>
...@@ -549,19 +547,13 @@ namespace Titanium.Web.Proxy.EventArguments ...@@ -549,19 +547,13 @@ namespace Titanium.Web.Proxy.EventArguments
/// Replace the response body with the specified string /// Replace the response body with the specified string
/// </summary> /// </summary>
/// <param name="body"></param> /// <param name="body"></param>
public async Task SetResponseBodyString(string body) public void SetResponseBodyString(string body)
{ {
if (!WebSession.Request.Locked) if (!WebSession.Request.Locked)
{ {
throw new Exception("You cannot call this function before request is made to server."); throw new Exception("You cannot call this function before request is made to server.");
} }
//syphon out the response body from server before setting the new body
if (!WebSession.Response.IsBodyRead)
{
await GetResponseBody();
}
var bodyBytes = WebSession.Response.Encoding.GetBytes(body); var bodyBytes = WebSession.Response.Encoding.GetBytes(body);
SetResponseBody(bodyBytes); SetResponseBody(bodyBytes);
...@@ -670,6 +662,8 @@ namespace Titanium.Web.Proxy.EventArguments ...@@ -670,6 +662,8 @@ namespace Titanium.Web.Proxy.EventArguments
throw new Exception("You cannot call this function after response is sent to the client."); throw new Exception("You cannot call this function after response is sent to the client.");
} }
response.Locked = true;
response.TerminateResponse = WebSession.Response.TerminateResponse;
WebSession.Response = response; WebSession.Response = response;
} }
else else
...@@ -677,14 +671,17 @@ namespace Titanium.Web.Proxy.EventArguments ...@@ -677,14 +671,17 @@ namespace Titanium.Web.Proxy.EventArguments
WebSession.Request.Locked = true; WebSession.Request.Locked = true;
response.Locked = true; response.Locked = true;
response.IsBodyRead = true;
WebSession.Response = response; WebSession.Response = response;
WebSession.Request.CancelRequest = true; WebSession.Request.CancelRequest = true;
} }
} }
public void TerminateServerConnection()
{
WebSession.Response.TerminateResponse = true;
}
/// <summary> /// <summary>
/// implement any cleanup here /// implement any cleanup here
/// </summary> /// </summary>
......
...@@ -20,7 +20,7 @@ namespace Titanium.Web.Proxy.Helpers ...@@ -20,7 +20,7 @@ namespace Titanium.Web.Proxy.Helpers
public async Task WriteResponseAsync(Response response, bool flush = true) public async Task WriteResponseAsync(Response response, bool flush = true)
{ {
await WriteResponseStatusAsync(response.HttpVersion, response.StatusCode, response.StatusDescription); await WriteResponseStatusAsync(response.HttpVersion, response.StatusCode, response.StatusDescription);
await WriteHeadersAsync(response.Headers, flush); await WriteAsync(response, flush);
} }
/// <summary> /// <summary>
......
...@@ -5,6 +5,7 @@ using System.Text; ...@@ -5,6 +5,7 @@ using System.Text;
using System.Threading.Tasks; using System.Threading.Tasks;
using StreamExtended.Helpers; using StreamExtended.Helpers;
using StreamExtended.Network; using StreamExtended.Network;
using Titanium.Web.Proxy.Compression;
using Titanium.Web.Proxy.Http; using Titanium.Web.Proxy.Http;
using Titanium.Web.Proxy.Shared; using Titanium.Web.Proxy.Shared;
...@@ -264,5 +265,62 @@ namespace Titanium.Web.Proxy.Helpers ...@@ -264,5 +265,62 @@ namespace Titanium.Web.Proxy.Helpers
onCopy?.Invoke(buffer, 0, bytesRead); onCopy?.Invoke(buffer, 0, bytesRead);
} }
} }
/// <summary>
/// Writes the request/response headers and body.
/// </summary>
/// <param name="requestResponse"></param>
/// <param name="flush"></param>
/// <returns></returns>
protected async Task WriteAsync(RequestResponseBase requestResponse, bool flush = true)
{
if (requestResponse.HasBody)
{
bool isChunked = requestResponse.IsChunked;
string contentEncoding = requestResponse.ContentEncoding;
var body = requestResponse.Body;
if (contentEncoding != null && body != null)
{
body = GetCompressedBody(contentEncoding, body);
if (isChunked == false)
{
requestResponse.ContentLength = body.Length;
}
else
{
requestResponse.ContentLength = -1;
}
}
await WriteHeadersAsync(requestResponse.Headers, flush);
await WriteBodyAsync(body, isChunked);
}
else
{
await WriteHeadersAsync(requestResponse.Headers, flush);
}
}
/// <summary>
/// get the compressed body from given bytes
/// </summary>
/// <param name="encodingType"></param>
/// <param name="body"></param>
/// <returns></returns>
internal byte[] GetCompressedBody(string encodingType, byte[] body)
{
var compressor = CompressionFactory.GetCompression(encodingType);
using (var ms = new MemoryStream())
{
using (var zip = compressor.GetStream(ms))
{
zip.Write(body, 0, body.Length);
}
return ms.ToArray();
}
}
} }
} }
...@@ -537,13 +537,14 @@ namespace Titanium.Web.Proxy ...@@ -537,13 +537,14 @@ namespace Titanium.Web.Proxy
//If request was modified by user //If request was modified by user
if (request.IsBodyRead) if (request.IsBodyRead)
{ {
var writer = args.WebSession.ServerConnection.StreamWriter;
bool isChunked = request.IsChunked; bool isChunked = request.IsChunked;
string contentEncoding = request.ContentEncoding; string contentEncoding = request.ContentEncoding;
var body = request.Body; var body = request.Body;
if (contentEncoding != null && body != null) if (contentEncoding != null && body != null)
{ {
body = GetCompressedBody(contentEncoding, body); body = writer.GetCompressedBody(contentEncoding, body);
if (isChunked == false) if (isChunked == false)
{ {
...@@ -555,7 +556,7 @@ namespace Titanium.Web.Proxy ...@@ -555,7 +556,7 @@ namespace Titanium.Web.Proxy
} }
} }
await args.WebSession.ServerConnection.StreamWriter.WriteBodyAsync(body, isChunked); await writer.WriteBodyAsync(body, isChunked);
} }
else else
{ {
......
...@@ -43,12 +43,24 @@ namespace Titanium.Web.Proxy ...@@ -43,12 +43,24 @@ namespace Titanium.Web.Proxy
await InvokeBeforeResponse(args); await InvokeBeforeResponse(args);
} }
if (response.TerminateResponse) // it may changed in the user event
response = args.WebSession.Response;
var clientStreamWriter = args.ProxyClient.ClientStreamWriter;
if (response.TerminateResponse || response.Locked)
{
await clientStreamWriter.WriteResponseAsync(response);
if (!response.TerminateResponse)
{ {
//syphon out the response body from server before setting the new body //syphon out the response body from server before setting the new body
await args.SyphonOutBodyAsync(false); await args.SyphonOutBodyAsync(false);
} }
return;
}
//if user requested to send request again //if user requested to send request again
//likely after making modifications from User Response Handler //likely after making modifications from User Response Handler
if (args.ReRequest) if (args.ReRequest)
...@@ -61,8 +73,6 @@ namespace Titanium.Web.Proxy ...@@ -61,8 +73,6 @@ namespace Titanium.Web.Proxy
response.Locked = true; response.Locked = true;
var clientStreamWriter = args.ProxyClient.ClientStreamWriter;
//Write back to client 100-conitinue response if that's what server returned //Write back to client 100-conitinue response if that's what server returned
if (response.Is100Continue) if (response.Is100Continue)
{ {
...@@ -75,9 +85,6 @@ namespace Titanium.Web.Proxy ...@@ -75,9 +85,6 @@ namespace Titanium.Web.Proxy
await clientStreamWriter.WriteLineAsync(); await clientStreamWriter.WriteLineAsync();
} }
//Write back response status to client
await clientStreamWriter.WriteResponseStatusAsync(response.HttpVersion, response.StatusCode, response.StatusDescription);
if (!args.IsTransparent) if (!args.IsTransparent)
{ {
response.Headers.FixProxyHeaders(); response.Headers.FixProxyHeaders();
...@@ -85,29 +92,12 @@ namespace Titanium.Web.Proxy ...@@ -85,29 +92,12 @@ namespace Titanium.Web.Proxy
if (response.IsBodyRead) if (response.IsBodyRead)
{ {
bool isChunked = response.IsChunked; await clientStreamWriter.WriteResponseAsync(response);
string contentEncoding = response.ContentEncoding;
var body = response.Body;
if (contentEncoding != null && body != null)
{
body = GetCompressedBody(contentEncoding, body);
if (isChunked == false)
{
response.ContentLength = body.Length;
}
else
{
response.ContentLength = -1;
}
}
await clientStreamWriter.WriteHeadersAsync(response.Headers);
await clientStreamWriter.WriteBodyAsync(body, isChunked);
} }
else else
{ {
//Write back response status to client
await clientStreamWriter.WriteResponseStatusAsync(response.HttpVersion, response.StatusCode, response.StatusDescription);
await clientStreamWriter.WriteHeadersAsync(response.Headers); await clientStreamWriter.WriteHeadersAsync(response.Headers);
//Write body if exists //Write body if exists
...@@ -123,27 +113,6 @@ namespace Titanium.Web.Proxy ...@@ -123,27 +113,6 @@ namespace Titanium.Web.Proxy
} }
} }
/// <summary>
/// get the compressed body from given bytes
/// </summary>
/// <param name="encodingType"></param>
/// <param name="body"></param>
/// <returns></returns>
private byte[] GetCompressedBody(string encodingType, byte[] body)
{
var compressor = CompressionFactory.GetCompression(encodingType);
using (var ms = new MemoryStream())
{
using (var zip = compressor.GetStream(ms))
{
zip.Write(body, 0, body.Length);
}
return ms.ToArray();
}
}
private async Task InvokeBeforeResponse(SessionEventArgs args) private async Task InvokeBeforeResponse(SessionEventArgs args)
{ {
if (BeforeResponse != null) if (BeforeResponse != null)
......
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