RubyでつくるWebSocketサーバーにprotocol-websocket gemを取り入れる その6

RubyでつくるWebSocketサーバーにprotocol-websocket gemを取り入れる その5 - @ledsun blog で機能するようになりました。 リファクタリングします。 TCPServerっぽいインターフェースにします。

まずはWebSocketクラス。 確立されたWebSocketコネクションを表すクラスです。

require 'protocol/websocket/headers'
require 'protocol/websocket/framer'
require 'protocol/websocket/text_frame'

class WebSocket
  include Protocol::WebSocket::Headers

  def initialize(socket)
    @socket = socket
  end

  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

次にWebSocketServerクラス WebSocketの接続を待ち受けるクラスです。

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

class WebSocketServer
  include Protocol::WebSocket::Headers

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

  def accept
    socket = @server.accept

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

    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.send_response socket

    WebSocket.new socket
  end

  private

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

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

最後にアプリケーション

require_relative 'web_socket.rb'
require_relative 'web_socket_server.rb'

server = WebSocketServer.new 'localhost', 2345

loop do
  socket = server.accept
  puts "Received: #{socket.gets}"

  socket.write "Loud and clear!"
  puts "Sent message"
  socket.close
end

どうです？かっこいいでしょう。