Skip to content

Commit

Permalink
Merge pull request #652 from adrianodennanni/main
Browse files Browse the repository at this point in the history
Add Zstd support (optional dependency, just like Brotli)
  • Loading branch information
flavorjones authored Aug 21, 2024
2 parents 752e4d2 + 1956295 commit 3bace72
Show file tree
Hide file tree
Showing 3 changed files with 79 additions and 2 deletions.
5 changes: 4 additions & 1 deletion Gemfile
Original file line number Diff line number Diff line change
Expand Up @@ -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
31 changes: 31 additions & 0 deletions lib/mechanize/http/agent.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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']}"
Expand Down
45 changes: 44 additions & 1 deletion test/test_mechanize_http_agent.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 3bace72

Please sign in to comment.