File: handler03.rb

package info (click to toggle)
ruby-websocket 1.2.9-3
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, forky, sid, trixie
  • size: 464 kB
  • sloc: ruby: 2,669; makefile: 4
file content (224 lines) | stat: -rw-r--r-- 7,673 bytes parent folder | download
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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
# encoding: binary
# frozen_string_literal: true

require 'securerandom'

module WebSocket
  module Frame
    module Handler
      class Handler03 < Base
        # Hash of frame names and it's opcodes
        FRAME_TYPES = {
          continuation: 0,
          close: 1,
          ping: 2,
          pong: 3,
          text: 4,
          binary: 5
        }.freeze

        # Hash of frame opcodes and it's names
        FRAME_TYPES_INVERSE = FRAME_TYPES.invert.freeze

        def initialize(frame)
          super
          @application_data_buffer = nil
        end

        # @see WebSocket::Frame::Base#supported_frames
        def supported_frames
          %i[text binary close ping pong]
        end

        # @see WebSocket::Frame::Handler::Base#encode_frame
        def encode_frame
          frame = if @frame.outgoing_masking?
                    masking_key = SecureRandom.random_bytes(4)
                    tmp_data = Data.new(masking_key + @frame.data)
                    tmp_data.set_mask
                    masking_key + tmp_data.getbytes(4, tmp_data.size)
                  else
                    @frame.data
                  end

          encode_header + frame
        end

        # @see WebSocket::Frame::Handler::Base#decode_frame
        def decode_frame
          while @frame.data.size > 1
            valid_header, more, frame_type, mask, payload_length = decode_header
            return unless valid_header

            application_data = decode_payload(payload_length, mask)

            if more
              decode_continuation_frame(application_data, frame_type)
            elsif frame_type == :continuation
              return decode_finish_continuation_frame(application_data)
            else
              raise(WebSocket::Error::Frame::InvalidPayloadEncoding) if frame_type == :text && !application_data.valid_encoding?
              return @frame.class.new(version: @frame.version, type: frame_type, data: application_data, decoded: true)
            end
          end
          nil
        end

        # Allow turning on or off masking
        def masking?
          false
        end

        private

        # This allows flipping the more bit to fin for draft 04
        def fin
          false
        end

        # Convert frame type name to opcode
        # @param [Symbol] frame_type Frame type name
        # @return [Integer] opcode or nil
        # @raise [WebSocket::Error] if frame opcode is not known
        def type_to_opcode(frame_type)
          FRAME_TYPES[frame_type] || raise(WebSocket::Error::Frame::UnknownFrameType)
        end

        # Convert frame opcode to type name
        # @param [Integer] opcode Opcode
        # @return [Symbol] Frame type name or nil
        # @raise [WebSocket::Error] if frame type name is not known
        def opcode_to_type(opcode)
          FRAME_TYPES_INVERSE[opcode] || raise(WebSocket::Error::Frame::UnknownOpcode)
        end

        def encode_header
          mask = @frame.outgoing_masking? ? 0b10000000 : 0b00000000

          output = String.new('')
          output << (type_to_opcode(@frame.type) | (fin ? 0b10000000 : 0b00000000)) # since more, rsv1-3 are 0 and 0x80 for Draft 4
          output << encode_payload_length(@frame.data.size, mask)
          output
        end

        def encode_payload_length(length, mask)
          output = String.new('')
          if length <= 125
            output << (length | mask) # since rsv4 is 0
          elsif length < 65_536 # write 2 byte length
            output << (126 | mask)
            output << [length].pack('n')
          else # write 8 byte length
            output << (127 | mask)
            output << [length >> 32, length & 0xFFFFFFFF].pack('NN')
          end
          output
        end

        def decode_header
          more, frame_type = decode_first_byte
          header_length, payload_length, mask = decode_second_byte(frame_type)
          return unless header_length

          # Compute the expected frame length
          frame_length = header_length + payload_length
          frame_length += 4 if mask

          raise(WebSocket::Error::Frame::TooLong) if frame_length > WebSocket.max_frame_size

          # Check buffer size
          return unless buffer_exists?(frame_length) # Buffer incomplete

          # Remove frame header
          @frame.data.slice!(0...header_length)

          [true, more, frame_type, mask, payload_length]
        end

        def buffer_exists?(buffer_number)
          !@frame.data.getbyte(buffer_number - 1).nil?
        end

        def decode_first_byte
          first_byte = @frame.data.getbyte(0)

          raise(WebSocket::Error::Frame::ReservedBitUsed) if first_byte & 0b01110000 != 0b00000000

          more = ((first_byte & 0b10000000) == 0b10000000) ^ fin
          frame_type = opcode_to_type first_byte & 0b00001111

          raise(WebSocket::Error::Frame::FragmentedControlFrame) if more && control_frame?(frame_type)
          raise(WebSocket::Error::Frame::DataFrameInsteadContinuation) if data_frame?(frame_type) && !@application_data_buffer.nil?

          [more, frame_type]
        end

        def decode_second_byte(frame_type)
          second_byte = @frame.data.getbyte(1)

          mask = @frame.incoming_masking? && (second_byte & 0b10000000) == 0b10000000
          length = second_byte & 0b01111111

          raise(WebSocket::Error::Frame::ControlFramePayloadTooLong) if length > 125 && control_frame?(frame_type)

          header_length, payload_length = decode_payload_length(length)

          [header_length, payload_length, mask]
        end

        def decode_payload_length(length)
          case length
          when 127 # Length defined by 8 bytes
            # Check buffer size
            return unless buffer_exists?(10) # Buffer incomplete

            # Only using the last 4 bytes for now, till I work out how to
            # unpack 8 bytes. I'm sure 4GB frames will do for now :)
            [10, @frame.data.getbytes(6, 4).unpack('N').first]
          when 126 # Length defined by 2 bytes
            # Check buffer size
            return unless buffer_exists?(4) # Buffer incomplete

            [4, @frame.data.getbytes(2, 2).unpack('n').first]
          else
            [2, length]
          end
        end

        def decode_payload(payload_length, mask)
          pointer = 0

          # Read application data (unmasked if required)
          @frame.data.set_mask if mask
          pointer += 4 if mask
          payload = @frame.data.getbytes(pointer, payload_length)
          payload.force_encoding('UTF-8')
          pointer += payload_length
          @frame.data.unset_mask if mask

          # Throw away data up to pointer
          @frame.data.slice!(0...pointer)

          payload
        end

        def decode_continuation_frame(application_data, frame_type)
          @application_data_buffer ||= String.new('')
          @application_data_buffer << application_data
          @frame_type ||= frame_type
        end

        def decode_finish_continuation_frame(application_data)
          raise(WebSocket::Error::Frame::UnexpectedContinuationFrame) unless @frame_type
          @application_data_buffer << application_data
          # Test valid UTF-8 encoding
          raise(WebSocket::Error::Frame::InvalidPayloadEncoding) if @frame_type == :text && !@application_data_buffer.valid_encoding?
          message = @frame.class.new(version: @frame.version, type: @frame_type, data: @application_data_buffer, decoded: true)
          @application_data_buffer = nil
          @frame_type = nil
          message
        end
      end
    end
  end
end