Merge 'Websocket server fixes' from Povilas Kanapickas

This PR fixes several issues in websocket frame parser. The most important one is that websocket parser didn't properly support partial packets. That is, the parser would fail if it received part of the websocket frame in one chunk of data and then received the rest of the frame separately.

The rest of fixes are general cleanups and fixes for bugs that can only be reproduced when using websocket parser in client mode (there will be more PRs in this area).

Unit tests have been added to better cover this area of the websocket parser.

The PR is best reviewed commit by commit.

Closes #2535

* https://github.com/scylladb/seastar:
  websocket: Remove unnecessary condition in frame parsing
  websocket: Fix logic when parsing header
  websocket: Avoid memory copy when full websocket frames are received
  websocket: Fix websocket frame parsing on partial packets

Author: xemul

include/seastar/websocket/server.hh

@@ -135,6 +135,7 @@ class websocket_parser {
     sstring _buffer;
     std::unique_ptr<frame_header> _header;
     uint64_t _payload_length;
+    uint64_t _consumed_payload_length = 0;
     uint32_t _masking_key;
     buff_t _result;
 
@@ -144,6 +145,9 @@ class websocket_parser {
     static future<consumption_result_t> stop(buff_t data) {
         return make_ready_future<consumption_result_t>(stop_consuming(std::move(data)));
     }
+    uint64_t remaining_payload_length() const {
+        return _payload_length - _consumed_payload_length;
+    }
 
     // Removes mask from payload given in p.
     void remove_mask(buff_t& p, size_t n) {

src/websocket/server.cc

@@ -207,23 +207,27 @@ future<websocket_parser::consumption_result_t> websocket_parser::operator()(
     }
     if (_state == parsing_state::flags_and_payload_data) {
         if (_buffer.length() + data.size() >= 2) {
-            if (_buffer.length() < 2) {
-                size_t hlen = _buffer.length();
-                _buffer.append(data.get(), 2 - hlen);
-                data.trim_front(2 - hlen);
-                _header = std::make_unique<frame_header>(_buffer.data());
-                _buffer = {};
-
-                // https://datatracker.ietf.org/doc/html/rfc6455#section-5.1
-                // We must close the connection if data isn't masked.
-                if ((!_header->masked) ||
-                        // RSVX must be 0
-                        (_header->rsv1 | _header->rsv2 | _header->rsv3) ||
-                        // Opcode must be known.
-                        (!_header->is_opcode_known())) {
-                    _cstate = connection_state::error;
-                    return websocket_parser::stop(std::move(data));
-                }
+            // _buffer.length() is less than 2 when entering this if body due to how
+            // the rest of code is structured. The else branch will never increase
+            // _buffer.length() to >=2 and other paths to this condition will always
+            // have buffer cleared.
+            assert(_buffer.length() < 2);
+
+            size_t hlen = _buffer.length();
+            _buffer.append(data.get(), 2 - hlen);
+            data.trim_front(2 - hlen);
+            _header = std::make_unique<frame_header>(_buffer.data());
+            _buffer = {};
+
+            // https://datatracker.ietf.org/doc/html/rfc6455#section-5.1
+            // We must close the connection if data isn't masked.
+            if ((!_header->masked) ||
+                    // RSVX must be 0
+                    (_header->rsv1 | _header->rsv2 | _header->rsv3) ||
+                    // Opcode must be known.
+                    (!_header->is_opcode_known())) {
+                _cstate = connection_state::error;
+                return websocket_parser::stop(std::move(data));
             }
             _state = parsing_state::payload_length_and_mask;
         } else {
@@ -238,35 +242,54 @@ future<websocket_parser::consumption_result_t> websocket_parser::operator()(
                 size_t hlen = _buffer.length();
                 _buffer.append(data.get(), required_bytes - hlen);
                 data.trim_front(required_bytes - hlen);
-
-                _payload_length = _header->length;
-                char const *input = _buffer.data();
-                if (_header->length == 126) {
-		    _payload_length = consume_be<uint16_t>(input);
-                } else if (_header->length == 127) {
-		    _payload_length = consume_be<uint64_t>(input);
-                }
-
-		_masking_key = consume_be<uint32_t>(input);
-                _buffer = {};
             }
+            _payload_length = _header->length;
+            char const *input = _buffer.data();
+            if (_header->length == 126) {
+                _payload_length = consume_be<uint16_t>(input);
+            } else if (_header->length == 127) {
+                _payload_length = consume_be<uint64_t>(input);
+            }
+
+            _masking_key = consume_be<uint32_t>(input);
+            _buffer = {};
             _state = parsing_state::payload;
         } else {
             _buffer.append(data.get(), data.size());
             return websocket_parser::dont_stop();
         }
     }
     if (_state == parsing_state::payload) {
-        if (_payload_length > data.size()) {
-            _payload_length -= data.size();
-            remove_mask(data, data.size());
-            _result = std::move(data);
-            return websocket_parser::stop(buff_t(0));
+        if (data.size() < remaining_payload_length()) {
+            // data has insufficient data to complete the frame - consume data.size() bytes
+            if (_result.empty()) {
+                _result = temporary_buffer<char>(remaining_payload_length());
+                _consumed_payload_length = 0;
+            }
+            std::copy(data.begin(), data.end(), _result.get_write() + _consumed_payload_length);
+            _consumed_payload_length += data.size();
+            return websocket_parser::dont_stop();
         } else {
-            _result = data.clone();
+            // data has sufficient data to complete the frame - consume remaining_payload_length()
+            auto consumed_bytes = remaining_payload_length();
+            if (_result.empty()) {
+                // Try to avoid memory copies in case when network packets contain one or more full
+                // websocket frames.
+                if (consumed_bytes == data.size()) {
+                    _result = std::move(data);
+                    data = temporary_buffer<char>(0);
+                } else {
+                    _result = data.share();
+                    _result.trim(consumed_bytes);
+                    data.trim_front(consumed_bytes);
+                }
+            } else {
+                std::copy(data.begin(), data.begin() + consumed_bytes,
+                          _result.get_write() + _consumed_payload_length);
+                data.trim_front(consumed_bytes);
+            }
             remove_mask(_result, _payload_length);
-            data.trim_front(_payload_length);
-            _payload_length = 0;
+            _consumed_payload_length = 0;
             _state = parsing_state::flags_and_payload_data;
             return websocket_parser::stop(std::move(data));
         }

tests/unit/websocket_test.cc

@@ -136,19 +136,36 @@ SEASTAR_TEST_CASE(test_websocket_handler_registration) {
         output.flush().get();
         input.read_exactly(186).get();
 
-        // Sending and receiving a websocket frame
-        const auto ws_frame = std::string(
-            "\202\204"  // 1000 0002 1000 0100
-            "TEST"      // Masking Key
-            "\0\0\0\0", 10); // Masked Message - TEST
-        const auto rs_frame = std::string(
-            "\202\004" // 1000 0002 0000 0100
-            "TEST", 6);    // Message - TEST
-        output.write(ws_frame).get();
-        output.flush().get();
+        unsigned ws_frame_len = 10;
+        for (unsigned split_i = 0; split_i < ws_frame_len - 1; ++split_i) {
+            // The loop tests various combinations of partial websocket frame coming in
+
+            // Sending and receiving a websocket frame
+            const std::string ws_frame = std::string(
+                "\202\204"  // 1000 0002 1000 0100
+                "TEST"      // Masking Key
+                "\0\0\0\0", ws_frame_len); // Masked Message - TEST
+            const auto rs_frame = std::string(
+                "\202\004" // 1000 0002 0000 0100
+                "TEST", 6);    // Message - TEST
+
+            if (split_i == 0) {
+                output.write(ws_frame).get();
+                output.flush().get();
+            } else {
+                output.write(ws_frame.substr(0, split_i)).get();
+                output.flush().get();
 
-        auto response = input.read_exactly(6).get();
-        auto response_str = std::string(response.begin(), response.end());
-        BOOST_REQUIRE_EQUAL(rs_frame, response_str);
+                // ensure that server attempts to read before the second part of the frame lands
+                sleep(std::chrono::milliseconds(100)).get();
+
+                output.write(ws_frame.substr(split_i)).get();
+                output.flush().get();
+            }
+
+            auto response = input.read_exactly(6).get();
+            auto response_str = std::string(response.begin(), response.end());
+            BOOST_REQUIRE_EQUAL(rs_frame, response_str);
+        }
     });
 }