prism/test/test_websocket_integration.rb

require 'minitest/autorun'
require 'prism'
require 'socket'
require 'timeout'

# Full stack: a real Prism::Server accepting a real TCP connection, doing
# the actual HTTP Upgrade handshake by hand (not through any test helper),
# exchanging masked/unmasked WebSocket frames, and broadcasting through a
# Hub shared by two independently connected clients - the same style of
# proof the project already leans on for server.rb/rate limiting (see
# test_server.rb), since framing bugs only ever show up on the real wire.
class WebSocketIntegrationTest < Minitest::Test
  Frame = Prism::WebSocket::Frame
  Handshake = Prism::WebSocket::Handshake

  class EchoChannel
    def self.hub
      @hub ||= Prism::WebSocket::Hub.new
    end

    def on_open(connection)
      self.class.hub.subscribe('echo_room', connection)
      connection.write_text('welcome')
    end

    def on_message(connection, message)
      if message == 'broadcast'
        self.class.hub.broadcast('echo_room', 'broadcasted!')
      else
        connection.write_text("echo: #{message}")
      end
    end

    def on_close(connection)
      self.class.hub.unsubscribe_all(connection)
    end
  end

  # Exercises Hub.new(max_subscribers_per_stream:) (see test_websocket_hub.rb
  # for the unit-level proof) over a real wire: the channel itself decides
  # what a rejected subscribe means, here mirroring the real
  # examples/blog/app/channels/post_activity_channel.rb pattern of closing
  # with 1008 (policy violation) when Hub#subscribe returns false.
  class CappedRoomChannel
    def self.hub
      @hub ||= Prism::WebSocket::Hub.new(max_subscribers_per_stream: 1)
    end

    def on_open(connection)
      return connection.close(1008, 'room is full') unless self.class.hub.subscribe('only_room', connection)
      connection.write_text('welcome')
    end

    def on_message(connection, message)
      connection.write_text("echo: #{message}")
    end

    def on_close(connection)
      self.class.hub.unsubscribe_all(connection)
    end
  end

  class WebSocketApp < Prism::Application
    websocket '/cable', EchoChannel
    websocket '/cable-capped', EchoChannel, max_message_size: 16
    websocket '/capped-room', CappedRoomChannel

    get '/ping' do
      html 'pong'
    end
  end

  def setup
    @server = Prism::Server.new(WebSocketApp, port: 0).start
    @port = @server.port
  end

  def teardown
    @server.stop
  end

  def test_ordinary_http_routes_still_work_alongside_a_registered_websocket_path
    socket = TCPSocket.new('127.0.0.1', @port)
    socket.write("GET /ping HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n")
    response = socket.read
    socket.close
    assert_includes response, '200'
    assert_includes response, 'pong'
  end

  def test_full_handshake_and_echo_round_trip
    socket = connect_and_upgrade('/cable')

    assert_equal 'welcome', Frame.decode(socket).payload

    socket.write(Frame.encode(Frame::TEXT, 'hello', mask: true))
    assert_equal 'echo: hello', Frame.decode(socket).payload
  ensure
    socket&.close
  end

  def test_unregistered_path_falls_through_to_the_ordinary_app_and_404s
    socket = TCPSocket.new('127.0.0.1', @port)
    socket.write(upgrade_request('/no-such-socket'))
    response = read_headers(socket)
    socket.close
    assert_includes response, '404'
  end

  def test_broadcast_reaches_every_subscribed_connection
    a = connect_and_upgrade('/cable')
    b = connect_and_upgrade('/cable')
    assert_equal 'welcome', Frame.decode(a).payload
    assert_equal 'welcome', Frame.decode(b).payload

    a.write(Frame.encode(Frame::TEXT, 'broadcast', mask: true))

    assert_equal 'broadcasted!', Frame.decode(a).payload
    assert_equal 'broadcasted!', Frame.decode(b).payload
  ensure
    a&.close
    b&.close
  end

  def test_close_handshake_is_echoed_back
    socket = connect_and_upgrade('/cable')
    Frame.decode(socket) # welcome

    socket.write(Frame.encode(Frame::CLOSE, '', mask: true))
    reply = Frame.decode(socket)
    assert_equal Frame::CLOSE, reply.opcode
  ensure
    socket&.close
  end

  # Application.websocket(path, handler, max_message_size:) end-to-end: a
  # per-endpoint cap lower than Connection::DEFAULT_MAX_MESSAGE_SIZE is
  # actually enforced by the real Server, and the server sends a 1009
  # close frame (RFC 6455 7.4.1) back before dropping the connection,
  # rather than merely dropping it silently.
  def test_a_per_endpoint_max_message_size_is_enforced_over_the_real_wire
    socket = connect_and_upgrade('/cable-capped')
    Frame.decode(socket) # welcome

    socket.write(Frame.encode(Frame::TEXT, 'this payload is well over sixteen bytes', mask: true))
    reply = Frame.decode(socket)
    assert_equal Frame::CLOSE, reply.opcode
    assert_equal 1009, reply.payload.byteslice(0, 2).unpack1('n')
  ensure
    socket&.close
  end

  # The default (unset) endpoint is unaffected by another endpoint's
  # custom max_message_size - proves the option is stored per path, not
  # as global mutable state on Connection::DEFAULT_MAX_MESSAGE_SIZE.
  def test_other_endpoints_keep_the_default_max_message_size
    socket = connect_and_upgrade('/cable')
    Frame.decode(socket) # welcome