RubyでWebSocketライブラリ

RubyでWebSocketクライアントを書く その3 - @ledsun blog までの諸諸の調査をまとめます。 このライブラリの目的はdRubyのTCPトランスポートを置き換えるために、TCPSocket、TCPServerのインターフェースに近い形で、WebSocketを扱う事です。

require 'socket'
require 'protocol/websocket/headers'
require 'protocol/websocket/framer'
require 'protocol/websocket/text_frame'
require_relative 'upgrade_request'
require_relative 'http_response'

module WANDS
  # This is a class that represents WebSocket, which has the same interface as TCPSocket.
  #
  # The WebSocket class is responsible for reading and writing messages to the server.
  #
  # Example usage:
  #
  # web_socket = WebSocket.open('localhost', 2345)
  # web_socket.write("Hello World!")
  #
  # puts web_socket.gets
  #
  # web_socket.close
  #
  class WebSocket
    include Protocol::WebSocket::Headers
    attr_reader :remote_address

    def self.open(host, port)
      socket = TCPSocket.new('localhost', 2345)
      request = UpgradeRequest.new
      socket.write(request.to_s)
      socket.flush

      response = HTTPResponse.new
      response.parse(socket)

      request.verify response

      self.new(socket)
    end

    def initialize(socket)
      @socket = socket
      @remote_address = socket.remote_address
    end

    # @return [String]
    def gets
      framer = Protocol::WebSocket::Framer.new(@socket)
      frame = framer.read_frame
      raise 'frame is not a text' unless frame.is_a? Protocol::WebSocket::TextFrame
      frame.unpack
    end

    def write(message)
      frame = Protocol::WebSocket::TextFrame.new(true, message)
      frame.write(@socket)
    end

    def close = @socket&.close
  end
end

require 'socket'
require 'protocol/websocket/headers'
require 'webrick/httprequest'
require 'webrick/httpresponse'
require 'webrick/config'
require_relative 'web_socket'

module WANDS
  # The WebSocketServer class is responsible for accepting WebSocket connections.
  # This class has the same interface as TCPServer.
  #
  # Example usage:
  #
  # server = WebSocketServer.new('localhost', 2345)
  # loop do
  #  begin
  #   socket = server.accept
  #   next unless socket
  #   puts "Accepted connection from #{socket.remote_address.ip_address} #{socket.remote_address.ip_port}"
  #
  #   received_message = socket.gets
  #   puts "Received: #{received_message}"
  #
  #   socket.write received_message
  #   socket.close
  #  rescue WEBrick::HTTPStatus::EOFError => e
  #   STDERR.puts e.message
  #  rescue Errno::ECONNRESET => e
  #   STDERR.puts "#{e.message} #{socket.remote_address.ip_address} #{socket.remote_address.ip_port}"
  #  rescue EOFError => e
  #   STDERR.puts "#{e.message} #{socket.remote_address.ip_address} #{socket.remote_address.ip_port}"
  #  end
  # end
  #
  class WebSocketServer
    include Protocol::WebSocket::Headers

    def initialize(hostname, port)
      @tcp_server = TCPServer.new hostname, port
    end

    def accept
      socket = @tcp_server.accept

      headers = read_headers_from socket
      unless headers["upgrade"].include? PROTOCOL
        socket.close
        raise "Not a websocket request"
      end

      response = response_to headers
      response.send_response socket

      WebSocket.new socket
    rescue WEBrick::HTTPStatus::BadRequest => e
      STDERR.puts e.message
      socket.write "HTTP/1.1 400 Bad Request\r\n\r\n"
      socket.close
    end

    private

    def read_headers_from(socket)
      request = WEBrick::HTTPRequest.new(WEBrick::Config::HTTP)
      request.parse(socket)
      request.header
    end

    def response_to(headers)
      response_key = calculate_accept_nonce_from headers
      response = WEBrick::HTTPResponse.new(WEBrick::Config::HTTP)
      response.status = 101
      response.upgrade! PROTOCOL
      response[SEC_WEBSOCKET_ACCEPT] = response_key

      response
    end

    def calculate_accept_nonce_from(headers)
      key = headers[SEC_WEBSOCKET_KEY].first
      Nounce.accept_digest(key)
    end
  end
end

require 'erb'
require 'protocol/websocket/headers'

module WANDS
  # The request is used to upgrade the HTTP connection to a WebSocket connection.
  class UpgradeRequest
    include ::Protocol::WebSocket::Headers

    TEMPLATE = <<~REQUEST
      GET / HTTP/1.1
      Connection: Upgrade
      Upgrade: websocket
      Sec-WebSocket-Version: 13
      Sec-WebSocket-Key: <%= @key %>

    REQUEST

    ERB = ERB.new(TEMPLATE).freeze

    def initialize
      @key = Nounce.generate_key
    end

    def to_s
      ERB.result(binding).gsub(/\r?\n/, "\r\n")
    end

    def verify(response)
      accept_digest = response.header[SEC_WEBSOCKET_ACCEPT].first
      accept_digest == Nounce.accept_digest(@key) || raise("Invalid accept digest")
    end
  end
end

module WANDS
  # This is a class that parses the response from the server and stores the headers in a hash.
  # The parse and header methods in this class are modeled on WEBrick::HTTPRequest.
  #
  # The expected HTTP response string is:
  #
  # HTTP/1.1 101 Switching Protocols
  # Upgrade: websocket
  # Connection: Upgrade
  # Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=
  # Sec-WebSocket-Protocol: chat
  # Sec-WebSocket-Version: 13
  #

  # Example usage:
  #
  # response = HTTPResponse.new
  # response.parse(socket)
  # response.header["upgrade"] # => ["websocket"]
  # response.header["connection"] # => ["Upgrade"]
  # response.header["sec-websocket-accept"] # => ["s3pPLMBiTxaQ9kYGzzhZRbK+xOo="]
  #
  class HTTPResponse
    attr_reader :header

    def parse(stream)
      @response = read_from stream
      @header = headers_of @response
    end

    def to_s
      @response
    end

    private

    def read_from(stream)
      response_string = ""
      while (line = stream.gets) != "\r\n"
        response_string += line
      end

      response_string
    end

    # Parse the headers from the HTTP response string.
    def headers_of(response_string)
      # Split the response string into headers and body.
      headers, _body = response_string.split("\r\n\r\n", 2)

      # Split the headers into lines.
      headers_lines = headers.split("\r\n")

      # The first line is the status line.
      # We don't need it, so we remove it from the headers.
      _status_line = headers_lines.shift

      # Parse the headers into a hash.
      headers_lines.map do |line|
        # Split the line into header name and value.
        header_name, value = line.split(': ', 2)
        [header_name.downcase, [value.strip]]
      end.to_h
    end
  end
end