140 lines
3.9 KiB
C++
140 lines
3.9 KiB
C++
// Copyright 2022 Memgraph Ltd.
|
|
//
|
|
// Use of this software is governed by the Business Source License
|
|
// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
|
|
// License, and you may not use this file except in compliance with the Business Source License.
|
|
//
|
|
// As of the Change Date specified in that file, in accordance with
|
|
// the Business Source License, use of this software will be governed
|
|
// by the Apache License, Version 2.0, included in the file
|
|
// licenses/APL.txt.
|
|
|
|
#pragma once
|
|
|
|
#include <algorithm>
|
|
#include <cstring>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
#include <fmt/format.h>
|
|
|
|
#include "communication/bolt/v1/constants.hpp"
|
|
|
|
namespace memgraph::communication::bolt {
|
|
|
|
/**
|
|
* This class is used as the return value of the GetChunk function of the
|
|
* ChunkedDecoderBuffer. It represents the 3 situations that can happen when
|
|
* reading a chunk.
|
|
*/
|
|
enum class ChunkState : uint8_t {
|
|
// The chunk isn't complete, we have to read more data
|
|
Partial,
|
|
|
|
// The chunk is whole and correct and has been loaded into the buffer
|
|
Whole,
|
|
|
|
// The chunk size is 0 meaning that the message is done
|
|
Done
|
|
};
|
|
|
|
/**
|
|
* @brief ChunkedDecoderBuffer
|
|
*
|
|
* Has methods for getting chunks and reading their data.
|
|
*
|
|
* Getting a chunk copies the chunk into the internal buffer from which
|
|
* the data can then be read. While getting a chunk the buffer checks the
|
|
* chunk for validity and then copies only data from the chunk. The headers
|
|
* aren't copied so that the decoder can read only the raw encoded data.
|
|
*/
|
|
template <typename TBuffer>
|
|
class ChunkedDecoderBuffer {
|
|
public:
|
|
ChunkedDecoderBuffer(TBuffer &buffer) : buffer_(buffer) { data_.reserve(kChunkMaxDataSize); }
|
|
|
|
/**
|
|
* Reads data from the internal buffer.
|
|
*
|
|
* @param data a pointer to where the data should be read
|
|
* @param len the length of data that should be read
|
|
* @returns true if exactly len of data was copied into data,
|
|
* false otherwise
|
|
*/
|
|
bool Read(uint8_t *data, size_t len) {
|
|
if (len > Size()) return false;
|
|
memcpy(data, &data_[pos_], len);
|
|
pos_ += len;
|
|
if (Size() == 0) {
|
|
pos_ = 0;
|
|
data_.clear();
|
|
}
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Peeks data from the internal buffer.
|
|
* Reads data, but doesn't remove it from the buffer.
|
|
*
|
|
* @param data a pointer to where the data should be read
|
|
* @param len the length of data that should be read
|
|
* @param offset offset from the beginning of the data
|
|
* @returns true if exactly len of data was copied into data,
|
|
* false otherwise
|
|
*/
|
|
bool Peek(uint8_t *data, size_t len, size_t offset = 0) {
|
|
if (len + offset > Size()) return false;
|
|
memcpy(data, &data_[pos_ + offset], len);
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Gets a chunk from the underlying raw data buffer.
|
|
*
|
|
* @returns ChunkState::Partial if the chunk isn't whole
|
|
* ChunkState::Whole if the chunk is whole
|
|
* ChunkState::Done if the chunk size is 0 (that signals that the
|
|
* message is whole)
|
|
*/
|
|
ChunkState GetChunk() {
|
|
uint8_t *data = buffer_.data();
|
|
size_t size = buffer_.size();
|
|
|
|
if (size < 2) {
|
|
return ChunkState::Partial;
|
|
}
|
|
|
|
size_t chunk_size = data[0];
|
|
chunk_size <<= 8;
|
|
chunk_size += data[1];
|
|
|
|
if (chunk_size == 0) {
|
|
// The message is done.
|
|
buffer_.Shift(2);
|
|
return ChunkState::Done;
|
|
}
|
|
|
|
if (size < chunk_size + 2) {
|
|
return ChunkState::Partial;
|
|
}
|
|
|
|
std::copy(data + 2, data + chunk_size + 2, std::back_inserter(data_));
|
|
buffer_.Shift(chunk_size + 2);
|
|
|
|
return ChunkState::Whole;
|
|
}
|
|
|
|
/**
|
|
* Gets the size of currently available data in the loaded chunk.
|
|
*
|
|
* @returns size of available data
|
|
*/
|
|
size_t Size() { return data_.size() - pos_; }
|
|
|
|
private:
|
|
TBuffer &buffer_;
|
|
std::vector<uint8_t> data_;
|
|
size_t pos_{0};
|
|
};
|
|
} // namespace memgraph::communication::bolt
|