Module: Rex::Proto::Http::WebSocket::Interface

Defined in:
lib/rex/proto/http/web_socket.rb

Overview

This defines the interface that the standard socket is extended with to provide WebSocket functionality. It should be used on a socket when the server has already successfully handled a WebSocket upgrade request.

Defined Under Namespace

Classes: Channel

Instance Method Summary collapse

Instance Method Details

#closeObject


305
306
307
308
309
310
311
312
313
# File 'lib/rex/proto/http/web_socket.rb', line 305

def close
  # if #wsloop was ever called, a synchronization lock will have been initialized
  @wsstream_lock.lock_write unless @wsstream_lock.nil?
  begin
    super
  ensure
    @wsstream_lock.unlock_write unless @wsstream_lock.nil?
  end
end

#get_wsframe(_opts = {}) ⇒ WebSocket::Frame

Read a WebSocket::Frame from the peer.

Returns:


191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
# File 'lib/rex/proto/http/web_socket.rb', line 191

def get_wsframe(_opts = {})
  frame = Frame.new
  frame.header.read(self)
  payload_data = ''
  while payload_data.length < frame.payload_len
    chunk = read(frame.payload_len - payload_data.length)
    if chunk.empty? # no partial reads!
      elog('WebSocket::Interface#get_wsframe: received an empty websocket payload data chunk')
      return nil
    end

    payload_data << chunk
  end
  frame.payload_data.assign(payload_data)
  frame
rescue ::IOError
  wlog('WebSocket::Interface#get_wsframe: encountered an IOError while reading a websocket frame')
  nil
end

#put_wsbinary(value, opts = {}) ⇒ Object

Build a WebSocket::Frame representing the binary data and send it to the peer.

Parameters:

  • value (String)

    the binary value to use as the frame payload.


175
176
177
# File 'lib/rex/proto/http/web_socket.rb', line 175

def put_wsbinary(value, opts = {})
  put_wsframe(Frame.from_binary(value), opts = opts)
end

#put_wsframe(frame, opts = {}) ⇒ Object

Send a WebSocket::Frame to the peer.

Parameters:


167
168
169
# File 'lib/rex/proto/http/web_socket.rb', line 167

def put_wsframe(frame, opts = {})
  put(frame.to_binary_s, opts = opts)
end

#put_wstext(value, opts = {}) ⇒ Object

Build a WebSocket::Frame representing the text data and send it to the peer.

Parameters:

  • value (String)

    the binary value to use as the frame payload.


183
184
185
# File 'lib/rex/proto/http/web_socket.rb', line 183

def put_wstext(value, opts = {})
  put_wsframe(Frame.from_text(value), opts = opts)
end

#to_wschannel(**kwargs) ⇒ WebSocket::Interface::Channel

Build a channel to allow reading and writing from the WebSocket. This provides high level functionality so the caller needn't worry about individual frames.


216
217
218
# File 'lib/rex/proto/http/web_socket.rb', line 216

def to_wschannel(**kwargs)
  Channel.new(self, **kwargs)
end

#wsclose(opts = {}) ⇒ Object

Close the WebSocket. If the underlying TCP socket is still active a WebSocket CONNECTION_CLOSE request will be sent and then it will wait for a CONNECTION_CLOSE response. Once completed the underlying TCP socket will be closed.


224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
# File 'lib/rex/proto/http/web_socket.rb', line 224

def wsclose(opts = {})
  return if closed? # there's nothing to do if the underlying TCP socket has already been closed

  # this implementation doesn't handle the optional close reasons at all
  frame = Frame.new(header: { opcode: Opcode::CONNECTION_CLOSE })
  # close frames must be masked
  # see: https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.1
  frame.mask!
  put_wsframe(frame, opts = opts)
  while (frame = get_wsframe(opts))
    break if frame.opcode == Opcode::CONNECTION_CLOSE
    # all other frames are dropped after our connection close request is sent
  end

  close # close the underlying TCP socket
end

#wsloop(opts = {}, &block) ⇒ Object

Run a loop to handle data from the remote end of the websocket. The loop will automatically handle fragmentation unmasking payload data and ping requests. When the remote connection is closed, the loop will exit. If specified the block will be passed data chunks and their data types.


246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
# File 'lib/rex/proto/http/web_socket.rb', line 246

def wsloop(opts = {}, &block)
  buffer = ''
  buffer_type = nil

  # since web sockets have their own tear down exchange, use a synchronization lock to ensure we aren't closed until
  # either the remote socket is closed or the teardown takes place
  @wsstream_lock = Rex::ReadWriteLock.new
  @wsstream_lock.synchronize_read do
    while (frame = get_wsframe(opts))
      frame.unmask! if frame.header.masked == 1

      case frame.header.opcode
      when Opcode::CONNECTION_CLOSE
        put_wsframe(Frame.new(header: { opcode: Opcode::CONNECTION_CLOSE }).tap { |f| f.mask! }, opts = opts)
        break
      when Opcode::CONTINUATION
        # a continuation frame can only be sent for a data frames
        # see: https://datatracker.ietf.org/doc/html/rfc6455#section-5.4
        raise WebSocketError, 'Received an unexpected continuation frame (uninitialized buffer)' if buffer_type.nil?

        buffer << frame.payload_data
      when Opcode::BINARY
        raise WebSocketError, 'Received an unexpected binary frame (incomplete buffer)' unless buffer_type.nil?

        buffer = frame.payload_data
        buffer_type = :binary
      when Opcode::TEXT
        raise WebSocketError, 'Received an unexpected text frame (incomplete buffer)' unless buffer_type.nil?

        buffer = frame.payload_data
        buffer_type = :text
      when Opcode::PING
        # see: https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.2
        put_wsframe(frame.dup.tap { |f| f.header.opcode = Opcode::PONG }, opts = opts)
      end

      next unless frame.header.fin == 1

      if block_given?
        # text data is UTF-8 encoded
        # see: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
        buffer.force_encoding('UTF-8') if buffer_type == :text
        # release the stream lock before entering the callback, allowing it to close the socket if desired
        @wsstream_lock.unlock_read
        begin
          block.call(buffer, buffer_type)
        ensure
          @wsstream_lock.lock_read
        end
      end

      buffer = ''
      buffer_type = nil
    end
  end

  close
end