WebSocket#close accepts code and reason arguments

spec/std/http/web_socket_spec.cr

@@ -275,14 +275,16 @@ describe HTTP::WebSocket do
   end
 
   describe "close" do
-    it "closes with message" do
-      message = "bye"
+    it "closes with code and response" do
+      response = "bye"
+      code = 4000_i16
       io = IO::Memory.new
       ws = HTTP::WebSocket::Protocol.new(io)
-      ws.close(message)
+      ws.close(code, response)
       bytes = io.to_slice
       (bytes[0] & 0x0f).should eq(0x8) # CLOSE frame
-      String.new(bytes[2, bytes[1]]).should eq(message)
+      IO::ByteFormat::NetworkEndian.decode(Int16, bytes[2, 3]).should eq(code)
+      String.new(bytes[4, bytes[1] - 2]).should eq(response)
     end
 
     it "closes without message" do

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

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

src/http/web_socket.cr

@@ -53,7 +53,7 @@ class HTTP::WebSocket
   def on_binary(&@on_binary : Bytes ->)
   end
 
-  def on_close(&@on_close : String ->)
+  def on_close(&@on_close : Int16?, String? ->)
   end
 
   protected def check_open
@@ -90,18 +90,26 @@ class HTTP::WebSocket
     end
   end
 
-  def close(message = nil)
+  # Closes the websocket.
+  def close
     return if closed?
     @closed = true
-    @ws.close(message)
+    @ws.close
+  end
+
+  # Closes the websocket with *code* and *reason*.
+  def close(code : Int16, reason : String? = nil)
+    return if closed?
+    @closed = true
+    @ws.close(code, reason)
   end
 
   def run
     loop do
       begin
         info = @ws.receive(@buffer)
       rescue IO::EOFError
-        @on_close.try &.call("")
+        @on_close.try &.call(nil, nil)
         break
       end
 
@@ -135,9 +143,16 @@ class HTTP::WebSocket
       when Protocol::Opcode::CLOSE
         @current_message.write @buffer[0, info.size]
         if info.final
-          message = @current_message.to_s
-          @on_close.try &.call(message)
-          close(message) unless closed?
+          begin
+            code = @current_message.read_bytes(Int16, IO::ByteFormat::NetworkEndian)
+            reason = @current_message.gets_to_end
+            @on_close.try &.call(code, reason)
+          rescue IO::EOFError
+            @on_close.try &.call(nil, nil)
+          ensure
+            close unless closed?
+          end
+
           @current_message.clear
           break
         end

src/http/web_socket/protocol.cr

@@ -229,12 +229,15 @@ class HTTP::WebSocket::Protocol
     end
   end
 
-  def close(message = nil)
-    if message
-      send(message.to_slice, Opcode::CLOSE)
-    else
-      send(Bytes.empty, Opcode::CLOSE)
-    end
+  def close
+    send(Bytes.empty, Opcode::CLOSE)
+  end
+
+  def close(code : Int16, reason : String? = nil)
+    slice = Bytes.new(sizeof(Int16) + (reason ? reason.bytesize : 0))
+    IO::ByteFormat::NetworkEndian.encode(code, slice)
+    reason.to_slice.copy_to(slice + sizeof(Int16)) if reason
+    send(slice, Opcode::CLOSE)
   end
 
   def self.new(host : String, path : String, port = nil, tls = false, headers = HTTP::Headers.new)