Commit 428b18b8 authored by Honfika's avatar Honfika

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

parent 6a73fd7d
......@@ -168,7 +168,7 @@ namespace Titanium.Web.Proxy.EventArguments
//If not already read (not cached yet)
if (!request.IsBodyRead)
{
var body = await ReadBodyAsync(ProxyClient.ClientStreamReader, true);
var body = await ReadBodyAsync(true);
request.Body = body;
//Now set the flag to true
......@@ -243,7 +243,7 @@ namespace Titanium.Web.Proxy.EventArguments
//If not already read (not cached yet)
if (!response.IsBodyRead)
{
var body = await ReadBodyAsync(WebSession.ServerConnection.StreamReader, false);
var body = await ReadBodyAsync(false);
response.Body = body;
//Now set the flag to true
......@@ -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())
{
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();
}
}
......@@ -271,24 +280,13 @@ namespace Titanium.Web.Proxy.EventArguments
return;
}
var reader = GetStreamReader(isRequest);
using (var bodyStream = new MemoryStream())
{
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>
/// This is called when the request is PUT/POST/PATCH to read the body
/// </summary>
......@@ -549,19 +547,13 @@ namespace Titanium.Web.Proxy.EventArguments
/// Replace the response body with the specified string
/// </summary>
/// <param name="body"></param>
public async Task SetResponseBodyString(string body)
public void SetResponseBodyString(string body)
{
if (!WebSession.Request.Locked)
{
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);
SetResponseBody(bodyBytes);
......@@ -670,6 +662,8 @@ namespace Titanium.Web.Proxy.EventArguments
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;
}
else
......@@ -677,14 +671,17 @@ namespace Titanium.Web.Proxy.EventArguments
WebSession.Request.Locked = true;
response.Locked = true;
response.IsBodyRead = true;
WebSession.Response = response;
WebSession.Request.CancelRequest = true;
}
}
public void TerminateServerConnection()
{
WebSession.Response.TerminateResponse = true;
}
/// <summary>
/// implement any cleanup here
/// </summary>
......
......@@ -20,7 +20,7 @@ namespace Titanium.Web.Proxy.Helpers
public async Task WriteResponseAsync(Response response, bool flush = true)
{
await WriteResponseStatusAsync(response.HttpVersion, response.StatusCode, response.StatusDescription);
await WriteHeadersAsync(response.Headers, flush);
await WriteAsync(response, flush);
}
/// <summary>
......
......@@ -5,6 +5,7 @@ using System.Text;
using System.Threading.Tasks;
using StreamExtended.Helpers;
using StreamExtended.Network;
using Titanium.Web.Proxy.Compression;
using Titanium.Web.Proxy.Http;
using Titanium.Web.Proxy.Shared;
......@@ -264,5 +265,62 @@ namespace Titanium.Web.Proxy.Helpers
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
//If request was modified by user
if (request.IsBodyRead)
{
var writer = args.WebSession.ServerConnection.StreamWriter;
bool isChunked = request.IsChunked;
string contentEncoding = request.ContentEncoding;
var body = request.Body;
if (contentEncoding != null && body != null)
{
body = GetCompressedBody(contentEncoding, body);
body = writer.GetCompressedBody(contentEncoding, body);
if (isChunked == false)
{
......@@ -555,7 +556,7 @@ namespace Titanium.Web.Proxy
}
}
await args.WebSession.ServerConnection.StreamWriter.WriteBodyAsync(body, isChunked);
await writer.WriteBodyAsync(body, isChunked);
}
else
{
......
......@@ -43,10 +43,22 @@ namespace Titanium.Web.Proxy
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)
{
//syphon out the response body from server before setting the new body
await args.SyphonOutBodyAsync(false);
await clientStreamWriter.WriteResponseAsync(response);
if (!response.TerminateResponse)
{
//syphon out the response body from server before setting the new body
await args.SyphonOutBodyAsync(false);
}
return;
}
//if user requested to send request again
......@@ -61,8 +73,6 @@ namespace Titanium.Web.Proxy
response.Locked = true;
var clientStreamWriter = args.ProxyClient.ClientStreamWriter;
//Write back to client 100-conitinue response if that's what server returned
if (response.Is100Continue)
{
......@@ -75,9 +85,6 @@ namespace Titanium.Web.Proxy
await clientStreamWriter.WriteLineAsync();
}
//Write back response status to client
await clientStreamWriter.WriteResponseStatusAsync(response.HttpVersion, response.StatusCode, response.StatusDescription);
if (!args.IsTransparent)
{
response.Headers.FixProxyHeaders();
......@@ -85,29 +92,12 @@ namespace Titanium.Web.Proxy
if (response.IsBodyRead)
{
bool isChunked = response.IsChunked;
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);
await clientStreamWriter.WriteResponseAsync(response);
}
else
{
//Write back response status to client
await clientStreamWriter.WriteResponseStatusAsync(response.HttpVersion, response.StatusCode, response.StatusDescription);
await clientStreamWriter.WriteHeadersAsync(response.Headers);
//Write body if exists
......@@ -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)
{
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