Skip to content

Commit

Permalink
Introduce #close_write and #shutdown.
Browse files Browse the repository at this point in the history
  • Loading branch information
ioquatix committed Sep 11, 2024
1 parent 2f5dfa1 commit c08d760
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 3 deletions.
8 changes: 8 additions & 0 deletions lib/protocol/websocket/close_frame.rb
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,14 @@ def pack(code = nil, reason = nil)
end
end

# Generate a suitable reply.
# @returns [CloseFrame]
def reply(code = Error::NO_ERROR, reason = "")
frame = CloseFrame.new
frame.pack(code, reason)
return frame
end

# Apply this frame to the specified connection.
def apply(connection)
connection.receive_close(self)
Expand Down
23 changes: 22 additions & 1 deletion lib/protocol/websocket/connection.rb
Original file line number Diff line number Diff line change
Expand Up @@ -89,13 +89,34 @@ def closed?
@state == :closed
end

# Immediately transition the connection to the closed state *and* close the underlying connection.
# Immediately transition the connection to the closed state *and* close the underlying connection. Any data not yet read will be lost.
def close(...)
close!(...)

@framer.close
end

# Close the connection gracefully, sending a close frame with the specified error code and reason. If an error occurs while sending the close frame, the connection will be closed immediately. You may continue to read data from the connection after calling this method, but you should not write any more data.
#
# @parameter error [Error | Nil] The error that occurred, if any.
def close_write(error = nil)
if error
send_close(Error::INTERNAL_ERROR, error.message)
else
send_close
end
end

# Close the connection gracefully. This will send a close frame and wait for the remote end to respond with a close frame. Any data received after the close frame is sent will be ignored. If you want to process this data, use {#close_write} instead, and read the data before calling {#close}.
def shutdown
send_close unless @state == :closed

# `read_frame` will return nil after receiving a close frame:
while read_frame
# Drain the connection.
end
end

# Read a frame from the framer, and apply it to the connection.
def read_frame
return nil if closed?
Expand Down
18 changes: 16 additions & 2 deletions lib/protocol/websocket/error.rb
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,24 @@ class Error < HTTP::Error
# Indicates that an endpoint is terminating the connection due to a protocol error.
PROTOCOL_ERROR = 1002

# Indicates that an endpoint is terminating the connection because it has received a type of data it cannot accept.
# Indicates that an endpoint is terminating the connection because it has received a type of data it cannot accept. (e.g., an endpoint that understands only text data MAY send this if it receives a binary message).
INVALID_DATA = 1003

# There are other status codes but most of them are "implementation specific".

# Indicates that an endpoint is terminating the connection because it has received data within a message that was not consistent with the type of the message (e.g., non-UTF-8 data within a text message).
INVALID_PAYLOAD = 1007

# Indicates that an endpoint is terminating the connection because it has received a message that violates its policy. This is a generic status code that can be returned when there is no other more suitable status code (e.g., 1003 or 1009) or if there is a need to hide specific details about the policy.
POLICY_VIOLATION = 1008

# Indicates that an endpoint is terminating the connection because it has received a message that is too big for it to process.
MESSAGE_TOO_LARGE = 1009

# Indicates that an endpoint (client) is terminating the connection because it has expected the server to negotiate one or more extension, but the server didn't return them in the response message of the WebSocket handshake. The list of extensions that are needed should appear in the /reason/ part of the Close frame. Note that this status code is not used by the server, because it can fail the WebSocket handshake instead.
MISSING_EXTENSION = 1010

# Indicates that a server is terminating the connection because it encountered an unexpected condition that prevented it from fulfilling the request.
INTERNAL_ERROR = 1011
end

# Raised by stream or connection handlers, results in GOAWAY frame which signals termination of the current connection. You *cannot* recover from this exception, or any exceptions subclassed from it.
Expand Down
32 changes: 32 additions & 0 deletions test/protocol/websocket/connection.rb
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,38 @@
end
end

with "#close_write" do
it "can close the write side of the connection" do
connection.close_write

frame = client.read_frame
expect(frame).to be_a(Protocol::WebSocket::CloseFrame)

client.write_frame(frame.reply)
client.close

frame = connection.read_frame
expect(frame).to be_a(Protocol::WebSocket::CloseFrame)

expect(connection).to be(:closed?)
end
end

with "#shutdown" do
it "can shutdown the connection" do
data_frame = Protocol::WebSocket::TextFrame.new(true).tap do |frame|
frame.pack("Hello World!")
end
client.write_frame(data_frame)

close_frame = Protocol::WebSocket::CloseFrame.new.tap(&:pack)
client.write_frame(close_frame)

connection.shutdown
expect(connection).to be(:closed?)
end
end

it "doesn't generate mask by default" do
expect(connection.mask).to be == nil
end
Expand Down

0 comments on commit c08d760

Please sign in to comment.