Merge pull request #652 from adrianodennanni/main

Add Zstd support (optional dependency, just like Brotli)

Author: flavorjones

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

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']}"

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