File: auto_deflate.rb

package info (click to toggle)
ruby-http 4.4.1-6
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 704 kB
  • sloc: ruby: 5,388; makefile: 9
file content (123 lines) | stat: -rw-r--r-- 2,926 bytes parent folder | download | duplicates (3)
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
123
# 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),
          :uri_normalizer => request.uri_normalizer
        )
      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