crystal-lang · Dec 20, 2016 · Dec 20, 2016 · Dec 20, 2016
Showing with 67 additions and 1 deletion.

+51 −1 src/http/web_socket.cr

+16 −0 src/http/web_socket/protocol.cr
diff --git a/src/http/web_socket.cr b/src/http/web_socket.cr
@@ -1,4 +1,6 @@
 class HTTP::WebSocket
+  getter? closed = false
+
   # :nodoc:
   def initialize(io : IO)
     initialize(Protocol.new(io))
@@ -37,23 +39,55 @@ class HTTP::WebSocket
     new(Protocol.new(host, path, port, tls))
   end
 
+  def on_ping(&@on_ping : String ->)
+  end
+
+  def on_pong(&@on_pong : String ->)
+  end
+
   def on_message(&@on_message : String ->)
   end
 
   def on_close(&@on_close : String ->)
   end
 
+  protected def check_open
+    raise IO::Error.new "closed socket" if closed?
+  end
+
   def send(message)
+    check_open
     @ws.send(message)
   end
 
+  # It's possible to send a PING frame, which the client must respond to
+  # with a PONG, or the server can send an unsolicited PONG frame
+  # which the client should not respond to.
+  #
+  # See `#pong`.
+  def ping(message = nil)
+    check_open
+    @ws.ping(message)
+  end
+
+  # Server can send an unsolicited PONG frame which the client should not respond to.
+  #
+  # See `#ping`.
+  def pong(message = nil)
+    check_open
+    @ws.pong(message)
+  end
+
   def stream(binary = true, frame_size = 1024)
+    check_open
     @ws.stream(binary: binary, frame_size: frame_size) do |io|
       yield io
     end
   end
 
   def close(message = nil)
+    return if closed?
+    @closed = true
     @ws.close(message)
   end
 
@@ -67,6 +101,20 @@ class HTTP::WebSocket
       end
 
       case info.opcode
+      when Protocol::Opcode::PING
+        @current_message.write @buffer[0, info.size]
+        if info.final
+          message = @current_message.to_s
+          @on_ping.try &.call(message)
+          pong(message) unless closed?
+          @current_message.clear
+        end
+      when Protocol::Opcode::PONG
+        @current_message.write @buffer[0, info.size]
+        if info.final
+          @on_pong.try &.call(@current_message.to_s)
+          @current_message.clear
+        end
       when Protocol::Opcode::TEXT
         @current_message.write @buffer[0, info.size]
         if info.final
@@ -76,7 +124,9 @@ class HTTP::WebSocket
       when Protocol::Opcode::CLOSE
         @current_message.write @buffer[0, info.size]
         if info.final
-          @on_close.try &.call(@current_message.to_s)
+          message = @current_message.to_s
+          @on_close.try &.call(message)
+          close(message) unless closed?
           @current_message.clear
           break
         end

diff --git a/src/http/web_socket/protocol.cr b/src/http/web_socket/protocol.cr
@@ -213,6 +213,22 @@ class HTTP::WebSocket::Protocol
     (@header[1] & 0x80_u8) != 0_u8
   end
 
+  def ping(message = nil)
+    if message
+      send(message.to_slice, Opcode::PING)
+    else
+      send(Bytes.empty, Opcode::PING)
+    end
+  end
+
+  def pong(message = nil)
+    if message
+      send(message.to_slice, Opcode::PONG)
+    else
+      send(Bytes.empty, Opcode::PONG)
+    end
+  end
+
   def close(message = nil)
     if message
       send(message.to_slice, Opcode::CLOSE)