Ensure WebsocketConnection::send releases source stream on failure

SmingHub · Mar 27, 2024 · cb508a7 · cb508a7
1 parent 99b2ccd
commit cb508a7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/Sming/Components/Network/src/Network/Http/Websocket/WebsocketConnection.cpp b/Sming/Components/Network/src/Network/Http/Websocket/WebsocketConnection.cpp
@@ -194,6 +194,9 @@ bool WebsocketConnection::send(const char* message, size_t length, ws_frame_type
 
 bool WebsocketConnection::send(IDataSourceStream* source, ws_frame_type_t type, bool useMask, bool isFin)
 {
+	// Ensure source gets destroyed if we return prematurely
+	std::unique_ptr<IDataSourceStream> sourceRef(source);
+
 	if(source == nullptr) {
 		return false;
 	}
@@ -279,10 +282,11 @@ bool WebsocketConnection::send(IDataSourceStream* source, ws_frame_type_t type,
 
 	// send the header
 	if(!connection->send(reinterpret_cast<const char*>(packet), packetLength)) {
-		delete source;
 		return false;
 	}
 
+	// Pass stream to connection
+	sourceRef.release();
 	return connection->send(source);
 }
 

diff --git a/Sming/Components/Network/src/Network/Http/Websocket/WebsocketConnection.h b/Sming/Components/Network/src/Network/Http/Websocket/WebsocketConnection.h
@@ -109,14 +109,14 @@ class WebsocketConnection
 
 	/**
 	 * @brief Sends websocket message from a stream
-	 * @param stream
+	 * @param source The stream to send - we get ownership of the stream
 	 * @param type
 	 * @param useMask MUST be true for client connections
 	 * @param isFin true if this is the final frame
 	 *
 	 * @retval bool true on success
 	 */
-	bool send(IDataSourceStream* stream, ws_frame_type_t type = WS_FRAME_TEXT, bool useMask = false, bool isFin = true);
+	bool send(IDataSourceStream* source, ws_frame_type_t type = WS_FRAME_TEXT, bool useMask = false, bool isFin = true);
 
 	/**
 	 * @brief Broadcasts a message to all active websocket connections