feat(websocket): add support for reason in WebSocket.close()

falcon/asgi/app.py

@@ -976,7 +976,13 @@ async def _handle_websocket(self, ver, scope, receive, send):
  # we don't support, so bail out. This also fulfills the ASGI
  # spec requirement to only process the request after
  # receiving and verifying the first event.
- await send({'type': EventType.WS_CLOSE, 'code': WSCloseCode.SERVER_ERROR})
+ await send(
+ {
+ 'type': EventType.WS_CLOSE,
+ 'code': WSCloseCode.SERVER_ERROR,
+ 'reason': 'Internal Server Error',
+ }
+ )
  return
 
  req = self._request_type(scope, receive, options=self.req_options)
@@ -988,6 +994,7 @@ async def _handle_websocket(self, ver, scope, receive, send):
  send,
  self.ws_options.media_handlers,
  self.ws_options.max_receive_queue,
+ self.ws_options.default_close_reasons,
  )
 
  on_websocket = None

falcon/asgi/ws.py

@@ -49,6 +49,7 @@ class WebSocket:
  '_asgi_send',
  '_buffered_receiver',
  '_close_code',
+ '_close_reasons',
  '_supports_accept_headers',
  '_mh_bin_deserialize',
  '_mh_bin_serialize',
@@ -69,6 +70,7 @@ def __init__(
  Union[media.BinaryBaseHandlerWS, media.TextBaseHandlerWS],
  ],
  max_receive_queue: int,
+ default_close_reasons: Dict[Optional[int], str],
  ):
  self._supports_accept_headers = ver != '2.0'
 
@@ -94,6 +96,7 @@ def __init__(
  self._mh_bin_serialize = mh_bin.serialize
  self._mh_bin_deserialize = mh_bin.deserialize
 
+ self._close_reasons = default_close_reasons
  self._state = _WebSocketState.HANDSHAKE
  self._close_code = None # type: Optional[int]
 
@@ -257,10 +260,18 @@ async def close(self, code: Optional[int] = None) -> None:
  if self.closed:
  return
 
+ if code in self._close_reasons:
+ reason = self._close_reasons[code]
+ elif 3100 <= code <= 3999:
+ reason = falcon.util.code_to_http_status(code - 3000)
+ else:
+ reason = ''
+
  await self._asgi_send(
  {
  'type': EventType.WS_CLOSE,
  'code': code,
+ 'reason': reason,
  }
  )
 
@@ -512,6 +523,9 @@ class WebSocketOptions:
  unhandled error is raised while handling a WebSocket connection
  (default ``1011``). For a list of valid close codes and ranges,
  see also: https://tools.ietf.org/html/rfc6455#section-7.4
+ default_close_reasons (dict): A default mapping between the Websocket
+ close code and the reason why the connection is close. HTTPerrors
+ are not included and will be rendered automatically with HTTP status.
  media_handlers (dict): A dict-like object for configuring media handlers
  according to the WebSocket payload type (TEXT vs. BINARY) of a
  given message. See also: :ref:`ws_media_handlers`.
@@ -528,7 +542,12 @@ class WebSocketOptions:
 
  """
 
- __slots__ = ['error_close_code', 'max_receive_queue', 'media_handlers']
+ __slots__ = [
+ 'error_close_code',
+ 'default_close_reasons',
+ 'max_receive_queue',
+ 'media_handlers',
+ ]
 
  def __init__(self):
  try:
@@ -557,6 +576,12 @@ def __init__(self):
  #
  self.error_close_code: int = WSCloseCode.SERVER_ERROR
 
+ self.default_close_reasons: Dict[int, str] = {
+ 1000: 'Normal Closure',
+ 1011: 'Internal Server Error',
+ 3011: 'Internal Server Error',
+ }
+
  # NOTE(kgriffs): The websockets library itself will buffer, so we keep
  # this value fairly small by default to mitigate buffer bloat. But in
  # the case that we have a large spillover from the websocket server's

falcon/testing/helpers.py

@@ -408,6 +408,7 @@ def __init__(self):
  self._state = _WebSocketState.CONNECT
  self._disconnect_emitted = False
  self._close_code = None
+ self._close_reason = None
  self._accepted_subprotocol = None
  self._accepted_headers = None
  self._collected_server_events = deque()
@@ -427,6 +428,10 @@ def closed(self) -> bool:
  def close_code(self) -> int:
  return self._close_code
 
+ @property
+ def close_reason(self) -> str:
+ return self._close_reason
+
  @property
  def subprotocol(self) -> str:
  return self._accepted_subprotocol
@@ -464,12 +469,14 @@ async def wait_ready(self, timeout: Optional[int] = 5):
  # NOTE(kgriffs): This is a coroutine just in case we need it to be
  # in a future code revision. It also makes it more consistent
  # with the other methods.
- async def close(self, code: Optional[int] = None):
+ async def close(self, code: Optional[int] = None, reason: Optional[str] = None):
  """Close the simulated connection.
 
  Keyword Args:
  code (int): The WebSocket close code to send to the application
  per the WebSocket spec (default: ``1000``).
+ reason (str): The WebSocket close reason to send to the application
+ per the WebSocket spec (default: empty string).
  """
 
  # NOTE(kgriffs): Give our collector a chance in case the
@@ -488,8 +495,12 @@ async def close(self, code: Optional[int] = None):
  if code is None:
  code = WSCloseCode.NORMAL
 
+ if reason is None:
+ reason = ''
+
  self._state = _WebSocketState.CLOSED
  self._close_code = code
+ self._close_reason = reason
 
  async def send_text(self, payload: str):
  """Send a message to the app with a Unicode string payload.
@@ -727,6 +738,7 @@ async def _collect(self, event: Dict[str, Any]):
  self._state = _WebSocketState.DENIED
 
  desired_code = event.get('code', WSCloseCode.NORMAL)
+ reason = event.get('reason', '')
  if desired_code == WSCloseCode.SERVER_ERROR or (
  3000 <= desired_code < 4000
  ):
@@ -735,12 +747,16 @@ async def _collect(self, event: Dict[str, Any]):
  # different raised error types or to pass through a
  # raised HTTPError status code.
  self._close_code = desired_code
+ self._close_reason = reason
  else:
  # NOTE(kgriffs): Force the close code to this since it is
  # similar to what happens with a real web server (the HTTP
  # connection is closed with a 403 and there is no websocket
  # close code).
  self._close_code = WSCloseCode.FORBIDDEN
+ self._close_reason = falcon.util.code_to_http_status(
+ WSCloseCode.FORBIDDEN - 3000
+ )
 
  self._event_handshake_complete.set()
 
@@ -755,6 +771,7 @@ async def _collect(self, event: Dict[str, Any]):
  if event_type == EventType.WS_CLOSE:
  self._state = _WebSocketState.CLOSED
  self._close_code = event.get('code', WSCloseCode.NORMAL)
+ self._close_reason = event.get('reason', '')
  else:
  assert event_type == EventType.WS_SEND
  self._collected_server_events.append(event)
@@ -780,7 +797,11 @@ def _create_checked_disconnect(self) -> Dict[str, Any]:
  )
 
  self._disconnect_emitted = True
- return {'type': EventType.WS_DISCONNECT, 'code': self._close_code}
+ return {
+ 'type': EventType.WS_DISCONNECT,
+ 'code': self._close_code,
+ 'reason': self._close_reason,
+ }
 
 
 # get_encoding_from_headers() is Copyright 2016 Kenneth Reitz, and is

tests/asgi/test_ws.py

@@ -1093,3 +1093,74 @@ def test_msgpack_missing():
 
  with pytest.raises(RuntimeError):
  handler.deserialize(b'{}')
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('reason', ['Client closing connection', '', None])
+async def test_client_close_with_reason(reason, conductor):
+ class Resource:
+ def __init__(self):
+ pass
+
+ async def on_websocket(self, req, ws):
+ await ws.accept()
+ while True:
+ try:
+ await ws.receive_data()
+
+ except falcon.WebSocketDisconnected:
+ break
+
+ resource = Resource()
+ conductor.app.add_route('/', resource)
+
+ async with conductor as c:
+ async with c.simulate_ws('/') as ws:
+ await ws.close(4099, reason)
+
+ assert ws.close_code == 4099
+ if reason:
+ assert ws.close_reason == reason
+ else:
+ assert ws.close_reason == ''
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('no_default', [True, False])
+@pytest.mark.parametrize('code', [None, 1011, 4099, 4042, 3405])
+async def test_no_reason_mapping(no_default, code, conductor):
+ class Resource:
+ def __init__(self):
+ pass
+
+ async def on_websocket(self, req, ws):
+ await ws.accept()
+ await ws.close(code)
+
+ resource = Resource()
+ conductor.app.add_route('/', resource)
+ if no_default:
+ conductor.app.ws_options.default_close_reasons = {}
+ else:
+ conductor.app.ws_options.default_close_reasons[4099] = '4099 reason'
+
+ async with conductor as c:
+ with pytest.raises(falcon.WebSocketDisconnected):
+ async with c.simulate_ws('/') as ws:
+ await ws.receive_data()
+
+ if code:
+ assert ws.close_code == code
+ else:
+ assert ws.close_code == CloseCode.NORMAL
+
+ if 3100 <= ws.close_code <= 3999:
+ assert ws.close_reason == falcon.util.code_to_http_status(ws.close_code - 3000)
+ elif (
+ no_default
+ or ws.close_code not in conductor.app.ws_options.default_close_reasons
+ ):
+ assert ws.close_reason == ''
+ else:
+ reason = conductor.app.ws_options.default_close_reasons[ws.close_code]
+ assert ws.close_reason == reason