Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

1
# -*- coding: binary -*-
2

3
require 'bindata'
4
require 'rex/post/channel'
5

6
module Rex::Proto::Http::WebSocket
7
  class WebSocketError < StandardError
8
  end
9

10
  class ConnectionError < WebSocketError
11
    def initialize(msg: 'The WebSocket connection failed', http_response: nil)
12
      @message = msg
13
      @http_response = http_response
14
    end
15

16
    attr_accessor :message, :http_response
17
    alias to_s message
18
  end
19

20
  # This defines the interface that the standard socket is extended with to provide WebSocket functionality. It should be
21
  # used on a socket when the server has already successfully handled a WebSocket upgrade request.
22
  module Interface
23
    #
24
    # A channel object that allows reading and writing either text or binary data directly to the remote peer.
25
    #
26
    class Channel
27
      include Rex::Post::Channel::StreamAbstraction
28

29
      module SocketInterface
30
        include Rex::Post::Channel::SocketAbstraction::SocketInterface
31

32
        def type?
33
          'tcp'
34
        end
35
      end
36

37
      # The socket parameters describing the underlying connection.
38
      # @!attribute [r] params
39
      #   @return [Rex::Socket::Parameters]
40
      attr_reader :params
41

42
      # @param [WebSocket::Interface] websocket the WebSocket that this channel is being opened on
43
      # @param [nil, Symbol] read_type the data type(s) to read from the WebSocket, one of :binary, :text or nil (for both
44
      #   binary and text)
45
      # @param [Symbol] write_type the data type to write to the WebSocket
46
      def initialize(websocket, read_type: nil, write_type: :binary)
47
        initialize_abstraction
48

49
        # a read type of nil will handle both binary and text frames that are received
50
        raise ArgumentError, 'read_type must be nil, :binary or :text' unless [nil, :binary, :text].include?(read_type)
51
        raise ArgumentError, 'write_type must be :binary or :text' unless %i[binary text].include?(write_type)
52

53
        @websocket = websocket
54
        @read_type = read_type
55
        @write_type = write_type
56
        @mutex = Mutex.new
57

58
        # beware of: https://github.com/rapid7/rex-socket/issues/32
59
        _, localhost, localport = websocket.getlocalname
60
        _, peerhost, peerport = Rex::Socket.from_sockaddr(websocket.getpeername)
61
        @params = Rex::Socket::Parameters.from_hash({
62
          'LocalHost' => localhost,
63
          'LocalPort' => localport,
64
          'PeerHost' => peerhost,
65
          'PeerPort' => peerport,
66
          'SSL' => websocket.respond_to?(:sslctx) && !websocket.sslctx.nil?
67
        })
68

69
        @thread = Rex::ThreadFactory.spawn("WebSocketChannel(#{localhost}->#{peerhost})", false) do
70
          websocket.wsloop do |data, data_type|
71
            next unless @read_type.nil? || data_type == @read_type
72

73
            data = on_data_read(data, data_type)
74
            next if data.nil?
75

76
            rsock.syswrite(data)
77
          end
78

79
          close
80
        end
81

82
        lsock.extend(SocketInterface)
83
        lsock.channel = self
84

85
        rsock.extend(SocketInterface)
86
        rsock.channel = self
87
      end
88

89
      def closed?
90
        @websocket.nil?
91
      end
92

93
      def close
94
        @mutex.synchronize do
95
          return if closed?
96

97
          @websocket.wsclose
98
          @websocket = nil
99
        end
100

101
        cleanup_abstraction
102
      end
103

104
      #
105
      # Close the channel for write operations. This sends a CONNECTION_CLOSE request, after which (per RFC 6455 section
106
      # 5.5.1) this side must not send any more data frames.
107
      #
108
      def close_write
109
        if closed?
110
          raise IOError, 'Channel has been closed.', caller
111
        end
112

113
        @websocket.put_wsframe(Frame.new(header: { opcode: Opcode::CONNECTION_CLOSE }))
114
      end
115

116
      #
117
      # Write *buf* to the channel, optionally truncating it to *length* bytes.
118
      #
119
      # @param [String] buf The data to write to the channel.
120
      # @param [Integer] length An optional length to truncate *data* to before
121
      #   sending it.
122
      def write(buf, length = nil)
123
        if closed?
124
          raise IOError, 'Channel has been closed.', caller
125
        end
126

127
        if !length.nil? && buf.length >= length
128
          buf = buf[0..length]
129
        end
130

131
        length = buf.length
132
        buf = on_data_write(buf)
133
        if @write_type == :binary
134
          @websocket.put_wsbinary(buf)
135
        elsif @write_type == :text
136
          @websocket.put_wstext(buf)
137
        end
138

139
        length
140
      end
141

142
      #
143
      # This provides a hook point that is called when data is read from the WebSocket peer. Subclasses can intercept and
144
      # process the data. The default functionality does nothing.
145
      #
146
      # @param [String] data the data that was read
147
      # @param [Symbol] data_type the type of data that was received, either :binary or :text
148
      # @return [String, nil] if a string is returned, it's passed through the channel
149
      def on_data_read(data, _data_type)
150
        data
151
      end
152

153
      #
154
      # This provides a hook point that is called when data is written to the WebSocket peer. Subclasses can intercept and
155
      # process the data. The default functionality does nothing.
156
      #
157
      # @param [String] data the data that is being written
158
      # @return [String, nil] if a string is returned, it's passed through the channel
159
      def on_data_write(data)
160
        data
161
      end
162
    end
163

164
    #
165
    # Send a WebSocket::Frame to the peer.
166
    #
167
    # @param [WebSocket::Frame] frame the frame to send to the peer.
168
    def put_wsframe(frame, opts = {})
169
      put(frame.to_binary_s, opts = opts)
170
    end
171

172
    #
173
    # Build a WebSocket::Frame representing the binary data and send it to the peer.
174
    #
175
    # @param [String] value the binary value to use as the frame payload.
176
    def put_wsbinary(value, opts = {})
177
      put_wsframe(Frame.from_binary(value), opts = opts)
178
    end
179

180
    #
181
    # Build a WebSocket::Frame representing the text data and send it to the peer.
182
    #
183
    # @param [String] value the binary value to use as the frame payload.
184
    def put_wstext(value, opts = {})
185
      put_wsframe(Frame.from_text(value), opts = opts)
186
    end
187

188
    #
189
    # Read a WebSocket::Frame from the peer.
190
    #
191
    # @return [Nil, WebSocket::Frame] the frame that was received from the peer.
192
    def get_wsframe(_opts = {})
193
      frame = Frame.new
194
      frame.header.read(self)
195
      payload_data = ''
196
      while payload_data.length < frame.payload_len
197
        chunk = read(frame.payload_len - payload_data.length)
198
        if chunk.empty? # no partial reads!
199
          elog('WebSocket::Interface#get_wsframe: received an empty websocket payload data chunk')
200
          return nil
201
        end
202

203
        payload_data << chunk
204
      end
205
      frame.payload_data.assign(payload_data)
206
      frame
207
    rescue ::IOError
208
      wlog('WebSocket::Interface#get_wsframe: encountered an IOError while reading a websocket frame')
209
      nil
210
    end
211

212
    #
213
    # Build a channel to allow reading and writing from the WebSocket. This provides high level functionality so the
214
    # caller needn't worry about individual frames.
215
    #
216
    # @return [WebSocket::Interface::Channel]
217
    def to_wschannel(**kwargs)
218
      Channel.new(self, **kwargs)
219
    end
220

221
    #
222
    # Close the WebSocket. If the underlying TCP socket is still active a WebSocket CONNECTION_CLOSE request will be sent
223
    # and then it will wait for a CONNECTION_CLOSE response. Once completed the underlying TCP socket will be closed.
224
    #
225
    def wsclose(opts = {})
226
      return if closed? # there's nothing to do if the underlying TCP socket has already been closed
227

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

240
      close # close the underlying TCP socket
241
    end
242

243
    #
244
    # Run a loop to handle data from the remote end of the websocket. The loop will automatically handle fragmentation
245
    # unmasking payload data and ping requests. When the remote connection is closed, the loop will exit. If specified the
246
    # block will be passed data chunks and their data types.
247
    #
248
    def wsloop(opts = {}, &block)
249
      buffer = ''
250
      buffer_type = nil
251

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

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

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

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

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

284
          next unless frame.header.fin == 1
285

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

299
          buffer = ''
300
          buffer_type = nil
301
        end
302
      end
303

304
      close
305
    end
306

307
    def close
308
      # if #wsloop was ever called, a synchronization lock will have been initialized
309
      @wsstream_lock.lock_write unless @wsstream_lock.nil?
310
      begin
311
        super
312
      ensure
313
        @wsstream_lock.unlock_write unless @wsstream_lock.nil?
314
      end
315
    end
316
  end
317

318
  class Opcode < BinData::Bit4
319
    CONTINUATION = 0
320
    TEXT = 1
321
    BINARY = 2
322
    CONNECTION_CLOSE = 8
323
    PING = 9
324
    PONG = 10
325

326
    default_parameter assert: -> { !Opcode.name(value).nil? }
327

328
    def self.name(value)
329
      constants.select { |c| c.upcase == c }.find { |c| const_get(c) == value }
330
    end
331

332
    def to_sym
333
      self.class.name(value)
334
    end
335
  end
336

337
  class Frame < BinData::Record
338
    endian :big
339

340
    struct :header do
341
      endian :big
342
      hide   :rsv1, :rsv2, :rsv3
343

344
      bit1   :fin, initial_value: 1
345
      bit1   :rsv1
346
      bit1   :rsv2
347
      bit1   :rsv3
348
      opcode :opcode
349
      bit1   :masked
350
      bit7   :payload_len_sm
351
      uint16 :payload_len_md, onlyif: -> { payload_len_sm == 126 }
352
      uint64 :payload_len_lg, onlyif: -> { payload_len_sm == 127 }
353
      uint32 :masking_key, onlyif: -> { masked == 1 }
354
    end
355
    string :payload_data, read_length: -> { payload_len }
356

357
    class << self
358
      private
359

360
      def from_opcode(opcode, payload, last: true, mask: true)
361
        frame = Frame.new(header: { fin: (last ? 1 : 0), opcode: opcode })
362
        frame.payload_len = payload.length
363
        frame.payload_data = payload
364

365
        case mask
366
        when TrueClass
367
          frame.mask!
368
        when Integer
369
          frame.mask!(mask)
370
        when FalseClass
371
        else
372
          raise ArgumentError, 'mask must be true, false or an integer (literal masking key)'
373
        end
374

375
        frame
376
      end
377
    end
378

379
    def self.apply_masking_key(data, mask)
380
      mask = [mask].pack('N').each_byte.to_a
381
      xored = ''
382
      data.each_byte.each_with_index do |byte, index|
383
        xored << (byte ^ mask[index % 4]).chr
384
      end
385

386
      xored
387
    end
388

389
    def self.from_binary(value, last: true, mask: true)
390
      from_opcode(Opcode::BINARY, value, last: last, mask: mask)
391
    end
392

393
    def self.from_text(value, last: true, mask: true)
394
      from_opcode(Opcode::TEXT, value, last: last, mask: mask)
395
    end
396

397
    #
398
    # Update the frame instance in place to apply a masking key to the payload data as defined in RFC 6455 section 5.3.
399
    #
400
    # @param [nil, Integer] key either an explicit 32-bit masking key or nil to generate a random one
401
    # @return [String] the masked payload data is returned
402
    def mask!(key = nil)
403
      header.masked.assign(1)
404
      key = rand(1..0xffffffff) if key.nil?
405
      header.masking_key.assign(key)
406
      payload_data.assign(self.class.apply_masking_key(payload_data, header.masking_key))
407
      payload_data.value
408
    end
409

410
    #
411
    # Update the frame instance in place to apply a masking key to the payload data as defined in RFC 6455 section 5.3.
412
    #
413
    # @return [String] the unmasked payload data is returned
414
    def unmask!
415
      payload_data.assign(self.class.apply_masking_key(payload_data, header.masking_key))
416
      header.masked.assign(0)
417
      payload_data.value
418
    end
419

420
    def payload_len
421
      case header.payload_len_sm
422
      when 127
423
        header.payload_len_lg
424
      when 126
425
        header.payload_len_md
426
      else
427
        header.payload_len_sm
428
      end
429
    end
430

431
    def payload_len=(value)
432
      if value < 126
433
        header.payload_len_sm.assign(value)
434
      elsif value < 0xffff
435
        header.payload_len_sm.assign(126)
436
        header.payload_len_md.assign(value)
437
      elsif value < 0x7fffffffffffffff
438
        header.payload_len_sm.assign(127)
439
        header.payload_len_lg.assign(value)
440
      else
441
        raise ArgumentError, 'payload length is outside the acceptable range'
442
      end
443
    end
444
  end
445
end
446

447