diff --git a/Gemfile b/Gemfile index 8e10e273..79538f60 100644 --- a/Gemfile +++ b/Gemfile @@ -6,4 +6,7 @@ gem "minitest", "~> 5.14" gem "rake", "~> 13.0" gem "rdoc", "~> 6.3" gem "rubocop", "~> 1.12" -gem "brotli", ">= 0.5" unless RUBY_PLATFORM == "java" +unless RUBY_PLATFORM == 'java' + gem 'brotli', '>= 0.5' + gem 'zstd-ruby', '~> 1.5' +end diff --git a/lib/mechanize/http/agent.rb b/lib/mechanize/http/agent.rb index 199eaa0f..d3f318a6 100644 --- a/lib/mechanize/http/agent.rb +++ b/lib/mechanize/http/agent.rb @@ -523,6 +523,35 @@ def content_encoding_brotli(body_io) body_io.close end + ## + # Decodes a Zstd-encoded +body_io+ + # + # (Experimental, CRuby only) Although Mechanize will never request a zstd-encoded response via + # `accept-encoding`, buggy servers may return zstd-encoded responses, or you might need to + # inform the zstd keyword on your Accept-Encoding headers. Let's try to handle those cases if + # the Zstd gem is loaded. + # + # If you need to handle Zstd-encoded responses, install the 'zstd-ruby' gem and require it in your + # application. If the `Zstd` constant is defined, Mechanize will attempt to use it to inflate + # the response. + # + def content_encoding_zstd(body_io) + log.debug('deflate zstd body') if log + + unless defined?(::Zstd) + raise Mechanize::Error, "cannot deflate zstd-encoded response. Please install and require the 'zstd-ruby' gem." + end + + begin + return StringIO.new(Zstd.decompress(body_io.read)) + rescue StandardError + log.error("unable to zstd#decompress response") if log + raise Mechanize::Error, "error decompressing zstd-encoded response." + end + ensure + body_io.close + end + def disable_keep_alive request request['connection'] = 'close' unless @keep_alive end @@ -861,6 +890,8 @@ def response_content_encoding response, body_io content_encoding_gunzip body_io when 'br' then content_encoding_brotli body_io + when 'zstd' then + content_encoding_zstd body_io else raise Mechanize::Error, "unsupported content-encoding: #{response['Content-Encoding']}" diff --git a/test/test_mechanize_http_agent.rb b/test/test_mechanize_http_agent.rb index ea807962..4ab00ebe 100644 --- a/test/test_mechanize_http_agent.rb +++ b/test/test_mechanize_http_agent.rb @@ -2,7 +2,10 @@ # frozen_string_literal: true require 'mechanize/test_case' -require "brotli" unless RUBY_PLATFORM == "java" +unless RUBY_PLATFORM == 'java' + require 'brotli' + require 'zstd-ruby' +end class TestMechanizeHttpAgent < Mechanize::TestCase @@ -965,6 +968,46 @@ def test_response_content_encoding_brotli_corrupt assert(body_io.closed?) end + def test_response_content_encoding_zstd_when_zstd_not_loaded + skip("only test this on jruby which doesn't have zstd support") unless RUBY_ENGINE == 'jruby' + + @res.instance_variable_set :@header, 'content-encoding' => %w[zstd] + body_io = StringIO.new("content doesn't matter for this test") + + e = assert_raises(Mechanize::Error) do + @agent.response_content_encoding(@res, body_io) + end + assert_includes(e.message, 'cannot deflate zstd-encoded response') + + assert(body_io.closed?) + end + + def test_response_content_encoding_zstd + skip('jruby does not have zstd support') if RUBY_ENGINE == 'jruby' + + @res.instance_variable_set :@header, 'content-encoding' => %w[zstd] + body_io = StringIO.new(Zstd.compress('this is compressed by zstd')) + + body = @agent.response_content_encoding(@res, body_io) + + assert_equal('this is compressed by zstd', body.read) + assert(body_io.closed?) + end + + def test_response_content_encoding_zstd_corrupt + skip('jruby does not have zstd support') if RUBY_ENGINE == 'jruby' + + @res.instance_variable_set :@header, 'content-encoding' => %w[zstd] + body_io = StringIO.new('not a zstd payload') + + e = assert_raises(Mechanize::Error) do + @agent.response_content_encoding(@res, body_io) + end + assert_includes(e.message, 'error decompressing zstd-encoded response') + assert_kind_of(RuntimeError, e.cause) + assert(body_io.closed?) + end + def test_response_content_encoding_gzip_corrupt log = StringIO.new logger = Logger.new log