diff --git a/lib/httpx/plugins/compression.rb b/lib/httpx/plugins/compression.rb index 621eedd1..e6769048 100644 --- a/lib/httpx/plugins/compression.rb +++ b/lib/httpx/plugins/compression.rb @@ -19,16 +19,24 @@ module HTTPX def initialize(*) super @_decoders = @headers.get("content-encoding").map do |encoding| - Transcoder.registry(encoding) + Transcoder.registry(encoding).decoder end + @_compressed_length = if @headers.key?("content-length") + @headers["content-length"].to_i + else + Float::INFINITY + end end - def write(*) + def write(chunk) + @_compressed_length -= chunk.bytesize + chunk = decompress(chunk) + super(chunk) + end + + def close super - if @length == @headers["content-length"].to_i - @buffer.rewind - @buffer = decompress(@buffer) - end + @_decoders.each(&:close) end private @@ -36,10 +44,27 @@ module HTTPX def decompress(buffer) @_decoders.reverse_each do |decoder| buffer = decoder.decode(buffer) + buffer << decoder.finish if @_compressed_length <= 0 end buffer end end + + class Decoder + extend Forwardable + + def_delegator :@inflater, :finish + + def_delegator :@inflater, :close + + def initialize(inflater) + @inflater = inflater + end + + def decode(chunk) + @inflater.inflate(chunk) + end + end class CompressEncoder attr_reader :content_type @@ -57,7 +82,7 @@ module HTTPX @buffer.rewind return @buffer.each(&blk) end - @encoder.compress(@raw, @buffer, &blk) + @encoder.compress(@raw, @buffer, chunk_size: 16_384, &blk) end def to_s @@ -67,7 +92,7 @@ module HTTPX end def bytesize - @encoder.compress(@raw, @buffer) + @encoder.compress(@raw, @buffer, chunk_size: 16_384) @buffer.size end diff --git a/lib/httpx/plugins/compression/brotli.rb b/lib/httpx/plugins/compression/brotli.rb index 1c62e03d..99bb3d36 100644 --- a/lib/httpx/plugins/compression/brotli.rb +++ b/lib/httpx/plugins/compression/brotli.rb @@ -15,17 +15,10 @@ module HTTPX Compression.register "br", self end - module ResponseBodyMethods - def write(chunk) - chunk = decompress(chunk) - super(chunk) - end - end - module Encoder module_function - def compress(raw, buffer, chunk_size: 16_384) + def compress(raw, buffer, chunk_size: ) return unless buffer.size.zero? raw.rewind begin @@ -38,15 +31,27 @@ module HTTPX end end + module BrotliWrapper + module_function + def inflate(text) + ::Brotli.inflate(text) + end + def close + end + def finish + "" + end + end + module BrotliTranscoder module_function def encode(payload) CompressEncoder.new(payload, Encoder) end - - def decode(io) - ::Brotli.inflate(io) + + def decoder + Decoder.new(BrotliWrapper) end end end diff --git a/lib/httpx/plugins/compression/deflate.rb b/lib/httpx/plugins/compression/deflate.rb index 0a9d2e5c..1d869996 100644 --- a/lib/httpx/plugins/compression/deflate.rb +++ b/lib/httpx/plugins/compression/deflate.rb @@ -18,7 +18,7 @@ module HTTPX module Encoder module_function - def compress(raw, buffer, chunk_size: 16_384) + def compress(raw, buffer, chunk_size: ) return unless buffer.size.zero? raw.rewind begin @@ -42,39 +42,12 @@ module HTTPX module_function - class Decoder - def initialize(io) - @io = io - @inflater = Zlib::Inflate.new(32 + Zlib::MAX_WBITS) - @buffer = StringIO.new - end - - def rewind - @buffer.rewind - end - - def read(*args) - return @buffer.read(*args) if @io.eof? - chunk = @io.read(*args) - inflated_chunk = @inflater.inflate(chunk) - inflated_chunk << @inflater.finish if @io.eof? - @buffer << chunk - inflated_chunk - end - - def close - @io.close - @io.unlink if @io.respond_to?(:unlink) - @inflater.close - end - end - def encode(payload) CompressEncoder.new(payload, Encoder) end - def decode(io) - Decoder.new(io) + def decoder + Decoder.new(Zlib::Inflate.new(32 + Zlib::MAX_WBITS)) end end end diff --git a/lib/httpx/plugins/compression/gzip.rb b/lib/httpx/plugins/compression/gzip.rb index 2fadcab9..7c4ebf72 100644 --- a/lib/httpx/plugins/compression/gzip.rb +++ b/lib/httpx/plugins/compression/gzip.rb @@ -1,5 +1,7 @@ # frozen_string_literal: true +require "forwardable" + module HTTPX module Plugins module Compression @@ -15,7 +17,7 @@ module HTTPX module GZIPTranscoder class Encoder - def compress(raw, buffer, chunk_size: 16_384) + def compress(raw, buffer, chunk_size: ) return unless buffer.size.zero? raw.rewind begin @@ -46,15 +48,16 @@ module HTTPX @compressed_chunk = nil end end - + + module_function def encode(payload) CompressEncoder.new(payload, Encoder.new) end - - def decode(io) - Zlib::GzipReader.new(io, window_size: 32 + Zlib::MAX_WBITS) + + def decoder + Decoder.new(Zlib::Inflate.new(32 + Zlib::MAX_WBITS)) end end end