Branch data     Line data    Source code

       1                 :             : /*
       2                 :             :  * Copyright 2025 Jinwoo Sung
       3                 :             :  *
       4                 :             :  * Licensed under the Apache License, Version 2.0 (the "License");
       5                 :             :  * you may not use this file except in compliance with the License.
       6                 :             :  * You may obtain a copy of the License at
       7                 :             :  *
       8                 :             :  *     http://www.apache.org/licenses/LICENSE-2.0
       9                 :             :  *
      10                 :             :  * Unless required by applicable law or agreed to in writing, software
      11                 :             :  * distributed under the License is distributed on an "AS IS" BASIS,
      12                 :             :  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      13                 :             :  * See the License for the specific language governing permissions and
      14                 :             :  * limitations under the License.
      15                 :             :  */
      16                 :             : 
      17                 :             : #include "wirestead/framer/length_prefix_framer.hpp"
      18                 :             : 
      19                 :             : #include <stdexcept>
      20                 :             : #include <utility>
      21                 :             : 
      22                 :             : #include "wirestead/diagnostics/logger.hpp"
      23                 :             : 
      24                 :             : namespace wirestead {
      25                 :             : namespace framer {
      26                 :             : 
      27                 :          16 : LengthPrefixFramer::LengthPrefixFramer(size_t prefix_bytes, Endian endian, size_t max_length,
      28                 :          16 :                                        bool length_includes_prefix)
      29                 :          16 :     : prefix_bytes_(prefix_bytes),
      30                 :          16 :       endian_(endian),
      31                 :          16 :       max_length_(max_length),
      32                 :          16 :       length_includes_prefix_(length_includes_prefix) {
      33   [ +  +  +  +  :          16 :   if (prefix_bytes != 1 && prefix_bytes != 2 && prefix_bytes != 4) {
                   +  + ]
      34                 :           2 :     throw std::invalid_argument("LengthPrefixFramer: prefix_bytes must be 1, 2 or 4");
      35                 :             :   }
      36         [ +  + ]:          14 :   if (max_length == 0) {
      37                 :           1 :     throw std::invalid_argument("LengthPrefixFramer: max_length must be greater than 0");
      38                 :             :   }
      39                 :          22 : }
      40                 :             : 
      41                 :          26 : bool LengthPrefixFramer::read_length(size_t offset, size_t& out) const {
      42         [ +  + ]:          26 :   if (buffer_.size() - offset < prefix_bytes_) return false;
      43                 :             : 
      44                 :          16 :   uint64_t value = 0;
      45         [ +  + ]:          54 :   for (size_t i = 0; i < prefix_bytes_; ++i) {
      46         [ +  + ]:          38 :     const size_t idx = endian_ == Endian::Big ? offset + i : offset + (prefix_bytes_ - 1 - i);
      47                 :          38 :     value = (value << 8) | buffer_[idx];
      48                 :             :   }
      49                 :          16 :   out = static_cast<size_t>(value);
      50                 :          16 :   return true;
      51                 :             : }
      52                 :             : 
      53                 :          15 : void LengthPrefixFramer::push_bytes(memory::ConstByteSpan data) {
      54         [ +  - ]:          15 :   if (!data.empty()) {
      55                 :          15 :     buffer_.insert(buffer_.end(), data.data(), data.data() + data.size());
      56                 :             :   }
      57                 :             : 
      58                 :          15 :   size_t consumed = 0;
      59                 :             :   while (true) {
      60                 :          26 :     size_t declared = 0;
      61         [ +  + ]:          26 :     if (!read_length(consumed, declared)) break;
      62                 :             : 
      63                 :             :     // Both conventions exist for what the length counts. A frame claiming to
      64                 :             :     // be shorter than its own header cannot be either, so it is corruption.
      65                 :          16 :     size_t payload_len = declared;
      66         [ +  + ]:          16 :     if (length_includes_prefix_) {
      67         [ +  + ]:           2 :       if (declared < prefix_bytes_) {
      68                 :           1 :         WIRESTEAD_LOG_WARNING("framer", "length_prefix",
      69                 :             :                               "Declared length is shorter than the prefix itself; resynchronising");
      70                 :           1 :         reset();
      71                 :           2 :         return;
      72                 :             :       }
      73                 :           1 :       payload_len = declared - prefix_bytes_;
      74                 :             :     }
      75                 :             : 
      76         [ +  + ]:          15 :     if (payload_len > max_length_) {
      77                 :             :       // Never allocate what a bad or hostile header asked for. Without a sync
      78                 :             :       // word there is nothing to hunt for, so the only honest recovery is to
      79                 :             :       // drop what we have and start over.
      80                 :           1 :       WIRESTEAD_LOG_WARNING("framer", "length_prefix",
      81                 :             :                             "Declared payload length exceeds max_length; dropping the buffer");
      82                 :           1 :       reset();
      83                 :           1 :       return;
      84                 :             :     }
      85                 :             : 
      86                 :          14 :     const size_t frame_end = consumed + prefix_bytes_ + payload_len;
      87         [ +  + ]:          14 :     if (buffer_.size() < frame_end) break;  // wait for the rest of the payload
      88                 :             : 
      89         [ +  - ]:          11 :     if (on_message_) {
      90                 :          11 :       on_message_(memory::ConstByteSpan(buffer_.data() + consumed + prefix_bytes_, payload_len));
      91                 :             :     }
      92                 :          11 :     consumed = frame_end;
      93                 :          11 :   }
      94                 :             : 
      95         [ +  + ]:          13 :   if (consumed > 0) {
      96                 :           9 :     buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<std::ptrdiff_t>(consumed));
      97                 :             :   }
      98                 :             : }
      99                 :             : 
     100                 :          11 : void LengthPrefixFramer::on_message(MessageCallback cb) { on_message_ = std::move(cb); }
     101                 :             : 
     102                 :           3 : void LengthPrefixFramer::reset() { buffer_.clear(); }
     103                 :             : 
     104                 :             : }  // namespace framer
     105                 :             : }  // namespace wirestead