diff --git a/src/WebSocket.h b/src/WebSocket.h index 2aedf39..22bbe97 100644 --- a/src/WebSocket.h +++ b/src/WebSocket.h @@ -88,7 +88,7 @@ public: return true; } - /* Emit close event, stat passive timeout */ + /* Emit close event, start passive timeout */ void close(int code, std::string_view message = {}) { static const int MAX_CLOSE_PAYLOAD = 123; int length = std::min(MAX_CLOSE_PAYLOAD, message.length()); @@ -96,11 +96,17 @@ public: // todo: here we start a timeout and handle it accordingly in the timeout handler WebSocketData *webSocketData = (WebSocketData *) static_dispatch(us_ssl_socket_ext, us_socket_ext)((SOCKET_TYPE *) this); + + /* We postpone any FIN sending to either drainage or uncorking */ webSocketData->isShuttingDown = true; /* Format and send the close frame */ char closePayload[MAX_CLOSE_PAYLOAD + 2]; int closePayloadLength = (int) WebSocketProtocol>::formatClosePayload(closePayload, code, message.data(), length); + + // but what if we are NOT corked, THEN we can FIN here if we succeeded + + // if we are corked and send returns true we cannot know for sure if we can fin send(std::string_view(closePayload, closePayloadLength), OpCode::CLOSE); // why should we fin here? @@ -110,7 +116,9 @@ public: WebSocketContextData *webSocketContextData = (WebSocketContextData *) static_dispatch(us_ssl_socket_context_ext, us_socket_context_ext)( (SOCKET_CONTEXT_TYPE *) static_dispatch(us_ssl_socket_get_context, us_socket_get_context)((SOCKET_TYPE *) this) ); - webSocketContextData->closeHandler(this, code, message); + if (webSocketContextData->closeHandler) { + webSocketContextData->closeHandler(this, code, message); + } } }; diff --git a/src/WebSocketContext.h b/src/WebSocketContext.h index 6944497..1d1b9d4 100644 --- a/src/WebSocketContext.h +++ b/src/WebSocketContext.h @@ -21,7 +21,7 @@ #include "WebSocketContextData.h" #include "WebSocketProtocol.h" #include "WebSocketData.h" -#include "AsyncSocket.h" +#include "WebSocket.h" namespace uWS { @@ -99,9 +99,11 @@ private: } /* Emit message event & break if we are closed or shut down when returning */ - webSocketContextData->messageHandler((WebSocket *) s, std::string_view(data, length), (uWS::OpCode) opCode); - if (us_socket_is_closed((us_socket *) s) || webSocketData->isShuttingDown) { - return true; + if (webSocketContextData->messageHandler) { + webSocketContextData->messageHandler((WebSocket *) s, std::string_view(data, length), (uWS::OpCode) opCode); + if (us_socket_is_closed((us_socket *)s) || webSocketData->isShuttingDown) { + return true; + } } } else { /* Allocate fragment buffer up front first time */ @@ -150,9 +152,11 @@ private: } /* Emit message and check for shutdown or close */ - webSocketContextData->messageHandler((WebSocket *) s, std::string_view(data, length), (uWS::OpCode) opCode); - if (us_socket_is_closed((us_socket *) s) || webSocketData->isShuttingDown) { - return true; + if (webSocketContextData->messageHandler) { + webSocketContextData->messageHandler((WebSocket *) s, std::string_view(data, length), (uWS::OpCode) opCode); + if (us_socket_is_closed((us_socket *)s) || webSocketData->isShuttingDown) { + return true; + } } /* If we shutdown or closed, this will be taken care of elsewhere */ @@ -247,18 +251,20 @@ private: webSocket->cork(); /* We need the websocket data */ - WebSocketData *wsState = (WebSocketData *) (static_dispatch(us_ssl_socket_ext, us_socket_ext)(s)); + WebSocketData *webSocketData = (WebSocketData *) (static_dispatch(us_ssl_socket_ext, us_socket_ext)(s)); /* This parser has virtually no overhead */ - uWS::WebSocketProtocol>::consume(data, length, wsState, s); + uWS::WebSocketProtocol>::consume(data, length, (WebSocketState *) webSocketData, s); + + // todo: we need to check for close and shutdown here? as we just emitted a bunch of message/close events! // todo: check for failures here just like for HTTP webSocket->uncork(); - // I guess we need to check drain here + // I guess we need to check drain here - emit drain if we had to poll for writable - // are we shutdown? - if (wsState->isShuttingDown) { + // are we shutdown? can onnly call this if we did succeed uncork! + if (webSocketData->isShuttingDown) { webSocket->shutdown(); } @@ -268,18 +274,34 @@ private: /* Handle HTTP write out (note: SSL_read may trigger this spuriously, the app need to handle spurious calls) */ static_dispatch(us_ssl_socket_context_on_writable, us_socket_context_on_writable)(getSocketContext(), [](auto *s) { - std::cout << "websocket writable" << std::endl; - - // we need to drain here! + // check if we already shut down the us socket and return? AsyncSocket *webSocket = (AsyncSocket *) s; + WebSocketData *webSocketData = (WebSocketData *)(static_dispatch(us_ssl_socket_ext, us_socket_ext)(s)); - // check for failures and shutdown just like in data event - webSocket->write(nullptr, 0); // drainage - also check for shutdown! - - // call drain here + /* Drain as much as possible */ + webSocket->write(nullptr, 0); + /* Are we in (WebSocket) shutdown mode? As in, have we called WebSocket::close? */ + if (webSocketData->isShuttingDown) { + /* Check if we just now drained completely */ + if (webSocket->getBufferedAmount() == 0) { + /* Now perform the actual TCP/TLS shutdown which was postponed due to backpressure */ + webSocket->shutdown(); + /* Set us to not shutting down so to avoid any spurious extra calls */ + webSocketData->isShuttingDown = false; + } + } else { + /* Call drain event even though nothing might actually changed */ + WebSocketContextData *webSocketContextData = (WebSocketContextData *) static_dispatch(us_ssl_socket_context_ext, us_socket_context_ext)( + static_dispatch(us_ssl_socket_get_context, us_socket_get_context)((SOCKET_TYPE *)s) + ); + if (webSocketContextData->drainHandler) { + webSocketContextData->drainHandler((WebSocket *) s); + } + /* No need to check for closed here as we leave the handler immediately*/ + } return s; }); @@ -315,8 +337,7 @@ private: } public: - - // we do not need SSL options as we come from adoptions + /* WebSocket contexts are always child contexts to a HTTP context so no SSL options are needed as they are inherited */ static WebSocketContext *create(Loop *loop, SOCKET_CONTEXT_TYPE *parentSocketContext) { WebSocketContext *webSocketContext = (WebSocketContext *)static_dispatch(us_create_child_ssl_socket_context, us_create_child_socket_context)(parentSocketContext, sizeof(WebSocketContextData)); if (!webSocketContext) { @@ -327,7 +348,6 @@ public: new ((WebSocketContextData *) static_dispatch(us_ssl_socket_context_ext, us_socket_context_ext)((SOCKET_CONTEXT_TYPE *)webSocketContext)) WebSocketContextData; return webSocketContext->init(); } - }; } diff --git a/src/WebSocketContextData.h b/src/WebSocketContextData.h index abeda42..a390bc1 100644 --- a/src/WebSocketContextData.h +++ b/src/WebSocketContextData.h @@ -29,9 +29,9 @@ template struct WebSocket; template struct WebSocketContextData { /* The callbacks for this context */ - std::function *, std::string_view, uWS::OpCode)> messageHandler; - std::function *)> drainHandler; - std::function *, int, std::string_view)> closeHandler; + std::function *, std::string_view, uWS::OpCode)> messageHandler = nullptr; + std::function *)> drainHandler = nullptr; + std::function *, int, std::string_view)> closeHandler = nullptr; }; }