Correctly support websocket close codes

spec/std/http/web_socket_spec.cr

@@ -278,14 +278,41 @@ describe HTTP::WebSocket do
   end
 
   describe "close" do
+    it "closes with code" do
+      io = IO::Memory.new
+      ws = HTTP::WebSocket::Protocol.new(io)
+      ws.close(4020)
+      bytes = io.to_slice
+      (bytes[0] & 0x0f).should eq(0x8) # CLOSE frame
+      bytes[1].should eq(2)            # 2 bytes code
+      bytes[2].should eq(0x0f)
+      bytes[3].should eq(0xb4)
+    end
+
     it "closes with message" do
       message = "bye"
       io = IO::Memory.new
       ws = HTTP::WebSocket::Protocol.new(io)
-      ws.close(message)
+      ws.close(nil, message)
       bytes = io.to_slice
       (bytes[0] & 0x0f).should eq(0x8) # CLOSE frame
-      String.new(bytes[2, bytes[1]]).should eq(message)
+      bytes[1].should eq(5)            # 2 + message.bytesize
+      bytes[2].should eq(0x03)
+      bytes[3].should eq(0xe8)
+      String.new(bytes[4..6]).should eq(message)
+    end
+
+    it "closes with message and code" do
+      message = "4020"
+      io = IO::Memory.new
+      ws = HTTP::WebSocket::Protocol.new(io)
+      ws.close(4020, message)
+      bytes = io.to_slice
+      (bytes[0] & 0x0f).should eq(0x8) # CLOSE frame
+      bytes[1].should eq(6)            # 2 + message.bytesize
+      bytes[2].should eq(0x0f)
+      bytes[3].should eq(0xb4)
+      String.new(bytes[4..7]).should eq(message)
     end
 
     it "closes without message" do
@@ -301,6 +328,7 @@ describe HTTP::WebSocket do
   each_ip_family do |family, _, any_address|
     it "negotiates over HTTP correctly" do
       address_chan = Channel(Socket::IPAddress).new
+      close_chan = Channel({Int32, String}).new
 
       f = spawn do
         http_ref = nil
@@ -313,8 +341,9 @@ describe HTTP::WebSocket do
             ws.send("pong #{str}")
           end
 
-          ws.on_close do
+          ws.on_close do |code, message|
             http_ref.not_nil!.close
+            close_chan.send({code, message})
           end
         end
 
@@ -332,11 +361,15 @@ describe HTTP::WebSocket do
       random = Random::Secure.hex
       ws2.on_message do |str|
         str.should eq("pong #{random}")
-        ws2.close
+        ws2.close(4020, "close message")
       end
       ws2.send(random)
 
       ws2.run
+
+      code, message = close_chan.receive
+      code.should eq(4020)
+      message.should eq("close message")
     end
 
     it "negotiates over HTTPS correctly" do

src/compiler/crystal/tools/playground/server.cr

@@ -490,7 +490,7 @@ module Crystal::Playground
         origin = context.request.headers["Origin"]
         if !accept_request?(origin)
           Log.warn { "Invalid Request Origin: #{origin}" }
-          ws.close "Invalid Request Origin"
+          ws.close HTTP::WebSocket::CloseCodes::PolicyViolation, "Invalid Request Origin"
         else
           @sessions_key += 1
           @sessions[@sessions_key] = session = Session.new(ws, @sessions_key, @port)

src/http/web_socket.cr

@@ -60,7 +60,7 @@ class HTTP::WebSocket
   def on_binary(&@on_binary : Bytes ->)
   end
 
-  def on_close(&@on_close : String ->)
+  def on_close(&@on_close : Int32, String ->)
   end
 
   protected def check_open
@@ -97,18 +97,24 @@ class HTTP::WebSocket
     end
   end
 
-  def close(message = nil)
+  @[Deprecated("Use WebSocket#close(code, message) instead")]
+  def close(message)
+    close(nil, message)
+  end
+
+  def close(code : Int? = nil, message = nil)
     return if closed?
     @closed = true
-    @ws.close(message)
+    @ws.close(code, message)
   end
 
   def run
     loop do
       begin
         info = @ws.receive(@buffer)
       rescue
-        @on_close.try &.call("")
+        @on_close.try &.call(CloseCodes::AbnormalClosure, "")
+        @closed = true
         break
       end
 
@@ -142,9 +148,18 @@ class HTTP::WebSocket
       when .close?
         @current_message.write @buffer[0, info.size]
         if info.final
-          message = @current_message.to_s
-          @on_close.try &.call(message)
-          close(message)
+          @current_message.rewind
+
+          if @current_message.size >= 2
+            code = @current_message.read_bytes(UInt16, IO::ByteFormat::NetworkEndian).to_i
+          else
+            code = CloseCodes::NoStatusReceived
+          end
+          message = @current_message.gets_to_end
+
+          @on_close.try &.call(code, message)
+          close(code)
+
           @current_message.clear
           break
         end

src/http/web_socket/close_codes.cr

@@ -0,0 +1,17 @@
+module HTTP::WebSocket::CloseCodes
+  NormalClosure           = 1000
+  GoingAway               = 1001
+  ProtocolError           = 1002
+  UnsupportedData         = 1003
+  NoStatusReceived        = 1005
+  AbnormalClosure         = 1006
+  InvalidFramePayloadData = 1007
+  PolicyViolation         = 1008
+  MessageTooBig           = 1009
+  MandatoryExtension      = 1010
+  InternalServerError     = 1011
+  ServiceRestart          = 1012
+  TryAgainLater           = 1013
+  BadGateway              = 1014
+  TLSHandshake            = 1015
+end

src/http/web_socket/protocol.cr

@@ -233,13 +233,27 @@ class HTTP::WebSocket::Protocol
     end
   end
 
-  def close(message = nil)
+  def close(code : Int? = nil, message = nil)
     return if @io.closed?
+
     if message
-      send(message.to_slice, Opcode::CLOSE)
+      message = message.to_slice
+      code ||= CloseCodes::NormalClosure
+
+      payload = Bytes.new(2 + message.size)
+      IO::ByteFormat::NetworkEndian.encode(code.to_u16, payload)
+      message.copy_to(payload + 2)
     else
-      send(Bytes.empty, Opcode::CLOSE)
+      if code
+        payload = Bytes.new(2)
+        IO::ByteFormat::NetworkEndian.encode(code.to_u16, payload)
+      else
+        payload = Bytes.empty
+      end
     end
+
+    send(payload, Opcode::CLOSE)
+
     @io.close if @sync_close
   end