/
auto_deflate.rb
122 lines (97 loc) · 2.8 KB
/
auto_deflate.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
# frozen_string_literal: true
require "zlib"
require "tempfile"
require "http/request/body"
module HTTP
module Features
class AutoDeflate < Feature
attr_reader :method
def initialize(*)
super
@method = @opts.key?(:method) ? @opts[:method].to_s : "gzip"
raise Error, "Only gzip and deflate methods are supported" unless %w[gzip deflate].include?(@method)
end
def wrap_request(request)
return request unless method
return request if request.body.size.zero?
# We need to delete Content-Length header. It will be set automatically by HTTP::Request::Writer
request.headers.delete(Headers::CONTENT_LENGTH)
request.headers[Headers::CONTENT_ENCODING] = method
Request.new(
:version => request.version,
:verb => request.verb,
:uri => request.uri,
:headers => request.headers,
:proxy => request.proxy,
:body => deflated_body(request.body)
)
end
def deflated_body(body)
case method
when "gzip"
GzippedBody.new(body)
when "deflate"
DeflatedBody.new(body)
end
end
HTTP::Options.register_feature(:auto_deflate, self)
class CompressedBody < HTTP::Request::Body
def initialize(uncompressed_body)
@body = uncompressed_body
@compressed = nil
end
def size
compress_all! unless @compressed
@compressed.size
end
def each(&block)
return to_enum __method__ unless block
if @compressed
compressed_each(&block)
else
compress(&block)
end
self
end
private
def compressed_each
while (data = @compressed.read(Connection::BUFFER_SIZE))
yield data
end
ensure
@compressed.close!
end
def compress_all!
@compressed = Tempfile.new("http-compressed_body", :binmode => true)
compress { |data| @compressed.write(data) }
@compressed.rewind
end
end
class GzippedBody < CompressedBody
def compress(&block)
gzip = Zlib::GzipWriter.new(BlockIO.new(block))
@body.each { |chunk| gzip.write(chunk) }
ensure
gzip.finish
end
class BlockIO
def initialize(block)
@block = block
end
def write(data)
@block.call(data)
end
end
end
class DeflatedBody < CompressedBody
def compress
deflater = Zlib::Deflate.new
@body.each { |chunk| yield deflater.deflate(chunk) }
yield deflater.finish
ensure
deflater.close
end
end
end
end
end