unit Sparkle.Middleware.Compress;

{$I Sparkle.Inc}

interface

uses
  System.Classes,
  Sparkle.HttpServer.Module,
  Sparkle.HttpServer.Context;

type
  ICompressMiddleware = interface(IHttpServerMiddleware)
  ['{71DB112A-5332-4D87-9454-4CE2278E721B}']
    function GetThreshold: Integer;
    procedure SetThreshold(const Value: Integer);
    property Threshold: Integer read GetThreshold write SetThreshold;
  end;

  TCompressMiddleware = class(THttpServerMiddleware, ICompressMiddleware)
  private
    FThreshold: Integer;
    function GetThreshold: Integer;
    procedure SetThreshold(const Value: Integer);
    function ShouldCompressBasedOnResponse(Response: THttpServerResponse): boolean;
  protected
    procedure ProcessRequest(Context: THttpServerContext; Next: THttpServerProc); override;
  public
    constructor Create;
    property Threshold: Integer read GetThreshold write SetThreshold;
  end;

implementation

uses
  System.SysUtils, System.ZLib,
  Sparkle.HttpServer.Response,
  Sparkle.Utils;

type
  TXDataBaseCompressionStream = class(TStream)
  private
    FStream: TStream;
    FStreamStartPos: Int64;
    FStreamPos: Int64;
    FZStream: TZStreamRec;
    FBuffer: TBytes;
  protected
    constructor Create(stream: TStream);
  end;

  TXDataCompressionStream = class(TXDataBaseCompressionStream)
  public
    constructor Create(dest: TStream); overload;
    constructor Create(dest: TStream; compressionLevel: TZCompressionLevel; windowBits: Integer); overload;
    destructor Destroy; override;
    procedure Flush;
    function Read(var buffer; count: Longint): Longint; override;
    function Write(const buffer; count: Longint): Longint; override;
    function Seek(const Offset: Int64; Origin: TSeekOrigin): Int64; override;
  end;

  TCompressEncoding = (Gzip, Deflate);

  TCompressRawWriter = class(TInterfacedObject, IHttpResponseBodyWriter)
  private
    FInnerWriter: IHttpResponseBodyWriter;
    FZStream: TXDataCompressionStream;
    FEnabled: boolean;
    FEncoding: TCompressEncoding;
    function Stream: TStream;
    procedure Flush;
  public
    constructor Create(AInnerWriter: IHttpResponseBodyWriter);
    destructor Destroy; override;
    property Enabled: boolean read FEnabled write FEnabled;
    property Encoding: TCompressEncoding read FEncoding write FEncoding;
  end;

{ TCompressRawWriter }

constructor TCompressRawWriter.Create(AInnerWriter: IHttpResponseBodyWriter);
begin
  FInnerWriter := AInnerWriter;
  FEnabled := true;
end;

destructor TCompressRawWriter.Destroy;
begin
  FZStream.Free;
  inherited;
end;

procedure TCompressRawWriter.Flush;
begin
  if FZStream <> nil then
    FZStream.Flush;
  FInnerWriter.Flush;
  FreeAndNil(FZStream);
end;

function TCompressRawWriter.Stream: TStream;
begin
  if Enabled then
  begin
    if FZStream = nil then
    begin
      case Encoding of
        TCompressEncoding.Deflate:
          FZStream := TXDataCompressionStream.Create(FInnerWriter.Stream);
      else
        // use gzip as default
        // TCompressEncoding.Gzip:
        FZStream := TXDataCompressionStream.Create(FInnerWriter.Stream,
          TZCompressionLevel.zcDefault, 15 or 16);
      end;
    end;
    Result := FZStream;
  end
  else
    Result := FInnerWriter.Stream;
end;

{ TCompressMiddleware }

constructor TCompressMiddleware.Create;
begin
  FThreshold := 1024;
end;

function TCompressMiddleware.GetThreshold: Integer;
begin
  Result := FThreshold;
end;

function TCompressMiddleware.ShouldCompressBasedOnResponse(Response: THttpServerResponse): boolean;
begin
  // Only compress if the response content type can be compressed
//  if not IsCompressible(C.Response.Headers.Get('content-type')) then Exit(false);

  // check for cache-control: no transform
//  if not IsCompressible(C.Response.Headers.Get('content-type')) then Exit(false);

  // if content length is specified and is lower then minimum size, do not compress
  if (Response.ContentLength >= 0) and (Response.ContentLength < Threshold) then
    Exit(false);

  // if response already has a content encoding and it's not identity, then it's already encoded
  if (Response.ContentEncoding <> '') and not SameText(Response.ContentEncoding, 'identity') then
    Exit(false);

  Result := True;
end;

procedure TCompressMiddleware.ProcessRequest(Context: THttpServerContext; Next: THttpServerProc);
var
  DoCompress: boolean;
  Encoding: string;
  CompressWriter: TCompressRawWriter;
  Writer: IHttpResponseBodyWriter;
begin
  // Find encoding to be used based on Accept-Encoding
  Encoding := TSparkleUtils.AcceptsEncoding(
    Context.Request.Headers.Get('accept-encoding'),
    TArray<string>.Create('gzip', 'deflate', 'identity')
  );

  DoCompress := true;
  // if client doesn't accept gzip or deflate, do no compress
  if (Encoding = '') or (Encoding = 'identity') then
    DoCompress := false;

  // if it's a HEAD request, do not compress
  if Context.Request.MethodType = THttpMethod.Head then
    DoCompress := false;

  // if we should not compress, just leave, forwarding the execution to next handler
  if not DoCompress then
  begin
    Next(Context);
    Exit;
  end;

  // Replace the raw writer
  CompressWriter := TCompressRawWriter.Create(Context.Response.RawWriter);
  Writer := CompressWriter;
  if SameText(Encoding, 'deflate') then
    CompressWriter.Encoding := TCompressEncoding.Deflate
  else
    CompressWriter.Encoding := TCompressEncoding.Gzip;
  Context.Response.RawWriter := Writer;

  // add the code that will be execute after headers are sent and right before
  // response body is about to be sent
  Context.Response.OnHeaders(
    procedure(Response: THttpServerResponse)
    begin
      DoCompress := ShouldCompressBasedOnResponse(Response);
      CompressWriter.Enabled := DoCompress;
      if DoCompress then
      begin
        Response.Headers.SetValue('content-encoding', Encoding); // set contend-encoding response
        Response.ContentLength := -1; // remove content-length header
//        Response.Chunked := true;
//        Response.Headers.SetValue('transfer-encoding', 'chunked');
      end;
    end
  );

  Next(Context);
end;

procedure TCompressMiddleware.SetThreshold(const Value: Integer);
begin
  FThreshold := Value;
end;

{ TXDataBaseCompressionStream }

constructor TXDataBaseCompressionStream.Create(stream: TStream);
begin
  inherited Create;
  FStream := stream;
  FStreamStartPos := Stream.Position;
  FStreamPos := FStreamStartPos;
  SetLength(FBuffer, $10000);
end;

{ TXDataCompressionStream }

function ZCompressCheck(code: Integer): Integer; overload;
begin
  Result := code;

  if code < 0 then
     raise EZCompressionError.Create(string(_z_errmsg[2 - code]));
end;

function ZCompressCheckWithoutBufferError(code: Integer): Integer; overload;
begin
  Result := code;

  if code < 0 then
    if (code <> Z_BUF_ERROR) then
     raise EZCompressionError.Create(string(_z_errmsg[2 - code]));
end;

constructor TXDataCompressionStream.Create(dest: TStream);
begin
  Create(dest, zcDefault, 15);
end;

constructor TXDataCompressionStream.Create(dest: TStream;
  compressionLevel: TZCompressionLevel; windowBits: Integer);
begin
  inherited Create(dest);

  FZStream.next_out := @FBuffer[0];
  FZStream.avail_out := Length(FBuffer);

  ZCompressCheck(DeflateInit2(FZStream, ZLevels[compressionLevel],
    Z_DEFLATED, windowBits, 8, Z_DEFAULT_STRATEGY));
end;

destructor TXDataCompressionStream.Destroy;
begin
  inherited;
end;

procedure TXDataCompressionStream.Flush;
begin
  FZStream.next_in := nil;
  FZStream.avail_in := 0;

  try
    if FStream.Position <> FStreamPos then FStream.Position := FStreamPos;

    while ZCompressCheckWithoutBufferError(deflate(FZStream, Z_FINISH)) <> Z_STREAM_END do
    begin
      FStream.WriteBuffer(FBuffer, Length(FBuffer) - Integer(FZStream.avail_out));

      FZStream.next_out := @FBuffer[0];
      FZStream.avail_out := Length(FBuffer);
    end;

    if Integer(FZStream.avail_out) < Length(FBuffer) then
    begin
      FStream.WriteBuffer(FBuffer, Length(FBuffer) - Integer(FZStream.avail_out));
    end;
  finally
    deflateEnd(FZStream);
  end;
end;

function TXDataCompressionStream.Read(var buffer; count: Longint): Longint;
begin
  raise EStreamError.Create('TXDataCompressionStream stream cannot be read.');
end;

function TXDataCompressionStream.Seek(const Offset: Int64;
  Origin: TSeekOrigin): Int64;
begin
  if (offset = 0) and (origin = soCurrent) then
  begin
    result := FZStream.total_in;
  end
  else
    raise EStreamError.Create('XData compression stream invalid seek.');
end;

function TXDataCompressionStream.Write(const buffer; count: Longint): Longint;
begin
  FZStream.next_in := @buffer;
  FZStream.avail_in := count;

  if FStream.Position <> FStreamPos then FStream.Position := FStreamPos;

  while FZStream.avail_in > 0 do
  begin
    ZCompressCheckWithoutBufferError(deflate(FZStream, Z_NO_FLUSH));

    if FZStream.avail_out = 0 then
    begin
      FStream.WriteBuffer(FBuffer, Length(FBuffer));

      FZStream.next_out := @FBuffer[0];
      FZStream.avail_out := Length(FBuffer);

      FStreamPos := FStream.Position;
    end;
  end;

  result := Count;
end;

end.
