@@ -71,7 +71,7 @@ bool WebSocketConnection::performHandshake() const {
7171 << " Origin: http://" " \r\n\r\n "
7272
7373 const std::string request = handshake.str ();
74- send (sock, request.c_str (), request.size (), 0 );
74+ send (sock, request.c_str (), static_cast < int >( request.size () ), 0 );
7575
7676 char buffer[BUFFER_SIZE];
7777 if (const int bytes_received = recv (sock, buffer, BUFFER_SIZE, 0 ); bytes_received > 0 ) {
@@ -140,49 +140,72 @@ void WebSocketConnection::sendMessage(const std::string &message) const {
140140 frame.push_back (message[i] ^ masking_key[i % 4 ]);
141141 }
142142
143- send (sock, reinterpret_cast <const char *>(frame.data ()), frame.size (), 0 );
143+ send (sock, reinterpret_cast <const char *>(frame.data ()), static_cast < int >( frame.size () ), 0 );
144144}
145145
146- // TODO: check/rework that because it throws an error if the message is too long
147146std::string WebSocketConnection::receiveMessage () const {
148- char buffer[BUFFER_SIZE];
149- int bytes_received = recv (sock, buffer, BUFFER_SIZE, 0 );
150- if (bytes_received <= 0 ) {
151- return " "
152- }
147+ std::vector<unsigned char > buffer (BUFFER_SIZE);
148+ std::vector<unsigned char > message_data;
153149
154- if ((buffer[0 ] & 0x0F ) != 0x1 ) {
155- return " "
156- }
150+ while (true ) {
151+ int bytes_received = recv (sock, reinterpret_cast <char *>(buffer.data ()), static_cast <int >(buffer.size ()), 0 );
152+ if (bytes_received <= 0 ) {
153+ return " "
154+ }
157155
158- int payload_len = buffer[1 ] & 0x7F ;
159- int offset = 2 ;
156+ if ((buffer[0 ] & 0x0F ) != 0x1 ) {
157+ return " "
158+ }
160159
161- if (payload_len == 126 ) {
162- payload_len = buffer[2 ] << 8 | buffer[3 ];
163- offset += 2 ;
164- } else if (payload_len == 127 ) {
165- payload_len = 0 ;
166- for (int i = 0 ; i < 8 ; ++i) {
167- payload_len = payload_len << 8 | buffer[offset++];
160+ int payload_len = buffer[1 ] & 0x7F ;
161+ int offset = 2 ;
162+
163+ if (payload_len == 126 ) {
164+ payload_len = buffer[2 ] << 8 | buffer[3 ];
165+ offset += 2 ;
166+ } else if (payload_len == 127 ) {
167+ payload_len = 0 ;
168+ for (int i = 0 ; i < 8 ; ++i) {
169+ payload_len = payload_len << 8 | buffer[offset++];
170+ }
168171 }
169- }
170172
171- bool mask = buffer[1 ] & 0x80 ;
172- unsigned char masking_key[4 ];
173- if (mask) {
174- std::memcpy (masking_key, buffer + offset, 4 );
175- offset += 4 ;
176- }
173+ // check if the message is masked
174+ const bool mask = buffer[1 ] & 0x80 ;
175+ unsigned char masking_key[4 ] = {};
176+ if (mask) {
177+ std::memcpy (masking_key, buffer.data () + offset, 4 );
178+ offset += 4 ;
179+ }
177180
178- std::string message (buffer + offset, payload_len);
179- if (mask) {
180- for (int i = 0 ; i < payload_len; ++i) {
181- message[i] ^= masking_key[i % 4 ];
181+ // read the actual payload
182+ int message_offset = offset;
183+ while (message_data.size () < payload_len) {
184+ const int chunk_size = std::min (payload_len - static_cast <int >(message_data.size ()),
185+ bytes_received - message_offset);
186+ message_data.insert (message_data.end (), buffer.begin () + message_offset,
187+ buffer.begin () + message_offset + chunk_size);
188+
189+ // if we've read the entire payload, break out of the loop
190+ if (message_data.size () >= payload_len) {
191+ break ;
192+ }
193+
194+ bytes_received = recv (sock, reinterpret_cast <char *>(buffer.data ()), static_cast <int >(buffer.size ()), 0 );
195+ if (bytes_received <= 0 ) {
196+ return " "
197+ }
198+ message_offset = 0 ; // reset message offset for the next chunk
182199 }
183- }
184200
185- return message;
201+ if (mask) {
202+ for (int i = 0 ; i < payload_len; ++i) {
203+ message_data[i] ^= masking_key[i % 4 ];
204+ }
205+ }
206+
207+ return {message_data.begin (), message_data.end ()};
208+ }
186209}
187210
188211void WebSocketConnection::closeConnection () {
0 commit comments