From 8cf69738bf6bbbc28c8eaf5b06c5cf8a0a44ff22 Mon Sep 17 00:00:00 2001
From: Matej Ferencevic <matej.ferencevic@memgraph.io>
Date: Wed, 6 Sep 2017 12:51:56 +0200
Subject: [PATCH] Fixed bolt decoder chunked buffer implementation.

Summary:
Bolt client now uses new chunked decoder.

Fixed bolt session to use new chunked decoder.

Fixed chunked decoder buffer test.

Fixed bolt session test.

Removed debug message from client.

Fixed bolt encoder to comply with specification.

Reviewers: buda, mislav.bradac

Reviewed By: buda

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D751
---
 CHANGELOG.md                                  |  1 +
 src/communication/bolt/client.hpp             | 30 ++++------
 src/communication/bolt/v1/constants.hpp       |  3 +-
 .../v1/decoder/chunked_decoder_buffer.hpp     | 59 ++++++++++---------
 .../v1/encoder/chunked_encoder_buffer.hpp     | 45 +++++++-------
 src/communication/bolt/v1/session.hpp         | 50 ++++++++--------
 .../bolt/v1/states/handshake.hpp              |  1 -
 tests/unit/bolt_chunked_decoder_buffer.cpp    | 30 ++--------
 tests/unit/bolt_chunked_encoder_buffer.cpp    | 12 ++--
 tests/unit/bolt_session.cpp                   | 16 -----
 10 files changed, 106 insertions(+), 141 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 906fa7afe..a7636c9b1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@
 ### Bug Fixes and Other Changes
 
 * Keywords appearing in header (named expressions) keep original case.
+* Our Bolt protocol implementation is now completely compatible with the protocol version 1 specification. (https://boltprotocol.org/v1/)
 
 ## v0.7.0
 
diff --git a/src/communication/bolt/client.hpp b/src/communication/bolt/client.hpp
index b5e464c24..d105adef2 100644
--- a/src/communication/bolt/client.hpp
+++ b/src/communication/bolt/client.hpp
@@ -75,14 +75,11 @@ class Client {
     DLOG(INFO) << "Sending init message";
     if (!encoder_.MessageInit(client_name, {{"scheme", "basic"},
                                             {"principal", username},
-                                            {"secret", password}})) {
+                                            {"credentials", password}})) {
       throw ClientSocketException();
     }
 
     DLOG(INFO) << "Reading init message response";
-    if (!GetDataByChunk()) {
-      throw ClientSocketException();
-    }
     Signature signature;
     DecodedValue metadata;
     if (!ReadMessage(&signature, &metadata)) {
@@ -91,6 +88,7 @@ class Client {
     if (signature != Signature::Success) {
       throw ClientInvalidDataException();
     }
+    DLOG(INFO) << "Metadata of init message response: " << metadata;
   }
 
   Client(const Client &) = delete;
@@ -109,9 +107,6 @@ class Client {
     encoder_.MessagePullAll();
 
     DLOG(INFO) << "Reading run message response";
-    if (!GetDataByChunk()) {
-      throw ClientSocketException();
-    }
     Signature signature;
     DecodedValue fields;
     if (!ReadMessage(&signature, &fields)) {
@@ -213,25 +208,25 @@ class Client {
   }
 
   bool GetDataByChunk() {
-    // If there is more data in the buffer then don't read data.
-    if (decoder_buffer_.Size() > 0) return true;
-
     ChunkState state;
-    while ((state = decoder_buffer_.GetChunk()) == ChunkState::Partial) {
+    while ((state = decoder_buffer_.GetChunk()) != ChunkState::Done) {
+      if (state == ChunkState::Whole) {
+        // The chunk is whole, no need to read more data.
+        continue;
+      }
       auto buff = buffer_.Allocate();
       int ret = socket_.Read(buff.data, buff.len);
       if (ret == -1) return false;
       buffer_.Written(ret);
     }
-
-    if (state == ChunkState::Whole) {
-      return true;
-    }
-    return false;
+    return true;
   }
 
   bool ReadMessage(Signature *signature, DecodedValue *ret) {
     Marker marker;
+    if (!GetDataByChunk()) {
+      return false;
+    }
     if (!decoder_.ReadMessageHeader(signature, &marker)) {
       return false;
     }
@@ -255,9 +250,6 @@ class Client {
     while (true) {
       Signature signature;
       DecodedValue data;
-      if (!GetDataByChunk()) {
-        throw ClientInvalidDataException();
-      }
       if (!ReadMessage(&signature, &data)) {
         throw ClientInvalidDataException();
       }
diff --git a/src/communication/bolt/v1/constants.hpp b/src/communication/bolt/v1/constants.hpp
index 3fd5a8ecf..1cf9a3c57 100644
--- a/src/communication/bolt/v1/constants.hpp
+++ b/src/communication/bolt/v1/constants.hpp
@@ -10,8 +10,7 @@ namespace communication::bolt {
 static constexpr size_t CHUNK_HEADER_SIZE = 2;
 static constexpr size_t MAX_CHUNK_SIZE = 65535;
 static constexpr size_t CHUNK_END_MARKER_SIZE = 2;
-static constexpr size_t WHOLE_CHUNK_SIZE =
-    CHUNK_HEADER_SIZE + MAX_CHUNK_SIZE + CHUNK_END_MARKER_SIZE;
+static constexpr size_t WHOLE_CHUNK_SIZE = CHUNK_HEADER_SIZE + MAX_CHUNK_SIZE;
 
 /**
  * Handshake size defined in the Bolt protocol.
diff --git a/src/communication/bolt/v1/decoder/chunked_decoder_buffer.hpp b/src/communication/bolt/v1/decoder/chunked_decoder_buffer.hpp
index 971087d69..5bde14377 100644
--- a/src/communication/bolt/v1/decoder/chunked_decoder_buffer.hpp
+++ b/src/communication/bolt/v1/decoder/chunked_decoder_buffer.hpp
@@ -22,11 +22,11 @@ enum class ChunkState : uint8_t {
   // The chunk isn't complete, we have to read more data
   Partial,
 
-  // The chunk is invalid, it's tail isn't 0x00 0x00
-  Invalid,
-
   // The chunk is whole and correct and has been loaded into the buffer
-  Whole
+  Whole,
+
+  // The chunk size is 0 meaning that the message is done
+  Done
 };
 
 /**
@@ -44,7 +44,9 @@ class ChunkedDecoderBuffer {
   using StreamBufferT = io::network::StreamBuffer;
 
  public:
-  ChunkedDecoderBuffer(Buffer<> &buffer) : buffer_(buffer) {}
+  ChunkedDecoderBuffer(Buffer<> &buffer) : buffer_(buffer) {
+    data_.reserve(MAX_CHUNK_SIZE);
+  }
 
   /**
    * Reads data from the internal buffer.
@@ -55,9 +57,13 @@ class ChunkedDecoderBuffer {
    *          false otherwise
    */
   bool Read(uint8_t *data, size_t len) {
-    if (len > size_ - pos_) return false;
+    if (len > Size()) return false;
     memcpy(data, &data_[pos_], len);
     pos_ += len;
+    if (Size() == 0) {
+      pos_ = 0;
+      data_.clear();
+    }
     return true;
   }
 
@@ -72,25 +78,23 @@ class ChunkedDecoderBuffer {
    *          false otherwise
    */
   bool Peek(uint8_t *data, size_t len, size_t offset = 0) {
-    if (len + offset > size_ - pos_) return false;
+    if (len + offset > Size()) return false;
     memcpy(data, &data_[pos_ + offset], len);
     return true;
   }
 
   /**
    * Gets a chunk from the underlying raw data buffer.
-   * When getting a chunk this function validates the chunk.
-   * If the chunk isn't yet finished the function just returns false.
-   * If the chunk is finished (all data has been read) and the chunk isn't
-   * valid, then the function automatically deletes the invalid chunk
-   * from the underlying buffer and returns false.
    *
-   * @returns true if a chunk was successfully copied into the internal
-   *          buffer, false otherwise
+   * @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) {
       DLOG(WARNING) << "Size < 2";
       return ChunkState::Partial;
@@ -99,22 +103,22 @@ class ChunkedDecoderBuffer {
     size_t chunk_size = data[0];
     chunk_size <<= 8;
     chunk_size += data[1];
-    if (size < chunk_size + 4) {
+
+    if (chunk_size == 0) {
+      // The message is done.
+      buffer_.Shift(2);
+      return ChunkState::Done;
+    }
+
+    if (size < chunk_size + 2) {
       DLOG(WARNING) << fmt::format(
           "Chunk size is {} but only have {} data bytes.", chunk_size, size);
       return ChunkState::Partial;
     }
 
-    if (data[chunk_size + 2] != 0 || data[chunk_size + 3] != 0) {
-      DLOG(WARNING) << "Invalid chunk!";
-      buffer_.Shift(chunk_size + 4);
-      return ChunkState::Invalid;
-    }
-
-    pos_ = 0;
-    size_ = chunk_size;
-    memcpy(data_, data + 2, size - 4);
-    buffer_.Shift(chunk_size + 4);
+    data_.reserve(data_.size() + chunk_size);
+    std::copy(data + 2, data + chunk_size + 2, std::back_inserter(data_));
+    buffer_.Shift(chunk_size + 2);
 
     return ChunkState::Whole;
   }
@@ -124,12 +128,11 @@ class ChunkedDecoderBuffer {
    *
    * @returns size of available data
    */
-  size_t Size() { return size_ - pos_; }
+  size_t Size() { return data_.size() - pos_; }
 
  private:
   Buffer<> &buffer_;
-  uint8_t data_[MAX_CHUNK_SIZE];
-  size_t size_{0};
+  std::vector<uint8_t> data_;
   size_t pos_{0};
 };
 }
diff --git a/src/communication/bolt/v1/encoder/chunked_encoder_buffer.hpp b/src/communication/bolt/v1/encoder/chunked_encoder_buffer.hpp
index 7c3665ed8..6e6b72fbc 100644
--- a/src/communication/bolt/v1/encoder/chunked_encoder_buffer.hpp
+++ b/src/communication/bolt/v1/encoder/chunked_encoder_buffer.hpp
@@ -26,11 +26,11 @@ namespace communication::bolt {
  * the currently stored data to the Socket. Chunking prepends data length and
  * appends chunk end marker (0x00 0x00).
  *
- * | chunk header | --- chunk --- | end marker | ---------- another chunk ... |
- * | ------------- whole chunk ----------------| ---------- another chunk ... |
+ * | chunk header | --- chunk --- | another chunk | -- end marker -- |
+ * | ------- whole chunk -------- |  whole chunk  | chunk of size 0  |
  *
- * | ------------------------ message --------------------------------------- |
- * | ------------------------ buffer  --------------------------------------- |
+ * | --------------------------- message --------------------------- |
+ * | --------------------------- buffer  --------------------------- |
  *
  * The current implementation stores the whole message into a single buffer
  * which is std::vector.
@@ -53,9 +53,7 @@ class ChunkedEncoderBuffer {
     while (n > 0) {
       // Define number of bytes  which will be copied into chunk because
       // chunk is a fixed length array.
-      auto size = n < MAX_CHUNK_SIZE + CHUNK_HEADER_SIZE - pos_
-                      ? n
-                      : MAX_CHUNK_SIZE + CHUNK_HEADER_SIZE - pos_;
+      auto size = n < WHOLE_CHUNK_SIZE - pos_ ? n : WHOLE_CHUNK_SIZE - pos_;
 
       // Copy size values to chunk array.
       std::memcpy(chunk_.data() + pos_, values, size);
@@ -67,15 +65,19 @@ class ChunkedEncoderBuffer {
 
       // If chunk is full copy it into the message buffer and make space for
       // other incomming values that are left in the values array.
-      if (pos_ == CHUNK_HEADER_SIZE + MAX_CHUNK_SIZE) Chunk();
+      if (pos_ == WHOLE_CHUNK_SIZE) Chunk(false);
     }
   }
 
   /**
    * Wrap the data from chunk array (append header and end marker) and put
    * the whole chunk into the buffer.
+   *
+   * @param message_done if set to true then chunk appends an end message
+   *                     marker to the chunk, should always be set to true
+   *                     (the default value), false is used only internally
    */
-  void Chunk() {
+  void Chunk(bool message_done = true) {
     // 1. Write the size of the chunk (CHUNK HEADER).
     uint16_t size = pos_ - CHUNK_HEADER_SIZE;
     // Write the higher byte.
@@ -83,24 +85,25 @@ class ChunkedEncoderBuffer {
     // Write the lower byte.
     chunk_[1] = size & 0xFF;
 
-    // 2. Write last two bytes in the whole chunk (CHUNK END MARKER).
-    // The last two bytes are always 0x00 0x00.
-    chunk_[pos_++] = 0x00;
-    chunk_[pos_++] = 0x00;
+    // 2. Determine the final size for the end marker.
+    if (message_done) size_ += 2;
 
-    debug_assert(pos_ <= WHOLE_CHUNK_SIZE,
-                 "Internal variable pos_ is bigger than the whole chunk size.");
-
-    // 3. Remember first chunk size.
-    if (first_chunk_size_ == -1) first_chunk_size_ = pos_;
-
-    // 4. Copy whole chunk into the buffer.
+    // 3. Copy whole chunk into the buffer.
     size_ += pos_;
     buffer_.reserve(size_);
     std::copy(chunk_.begin(), chunk_.begin() + pos_,
               std::back_inserter(buffer_));
 
-    // 5. Cleanup.
+    // 4. Insert message end marker.
+    if (message_done) {
+      buffer_.push_back(0);
+      buffer_.push_back(0);
+    }
+
+    // 5. Remember first chunk size.
+    if (first_chunk_size_ == -1) first_chunk_size_ = size_;
+
+    // 6. Cleanup.
     //     * pos_ has to be reset to the size of chunk header (reserved
     //       space for the chunk size)
     pos_ = CHUNK_HEADER_SIZE;
diff --git a/src/communication/bolt/v1/session.hpp b/src/communication/bolt/v1/session.hpp
index 8926a05ec..680aa972e 100644
--- a/src/communication/bolt/v1/session.hpp
+++ b/src/communication/bolt/v1/session.hpp
@@ -79,34 +79,31 @@ class Session {
    * Goes through the bolt states in order to execute commands from the client.
    */
   void Execute() {
-    // while there is data in the buffers
-    while (buffer_.size() > 0 || decoder_buffer_.Size() > 0) {
-      if (LIKELY(connected_)) {
-        DLOG(INFO) << fmt::format("Decoding chunk of size {}", buffer_.size());
-        auto chunk_state = decoder_buffer_.GetChunk();
-        if (chunk_state == ChunkState::Partial) {
-          DLOG(WARNING) << "Chunk isn't complete!";
-          return;
-        } else if (chunk_state == ChunkState::Invalid) {
-          DLOG(WARNING) << "Chunk is invalid!";
-          ClientFailureInvalidData();
-          return;
-        }
-        // if chunk_state == ChunkState::Whole then we continue with
-        // execution of the select below
-      } else if (buffer_.size() < HANDSHAKE_SIZE) {
+    if (UNLIKELY(!handshake_done_)) {
+      if (buffer_.size() < HANDSHAKE_SIZE) {
         DLOG(WARNING) << fmt::format("Received partial handshake of size {}",
                                      buffer_.size());
         return;
-      } else {
-        DLOG(WARNING) << fmt::format("Decoding handshake of size {}",
-                                     buffer_.size());
+      }
+      DLOG(WARNING) << fmt::format("Decoding handshake of size {}",
+                                   buffer_.size());
+      state_ = StateHandshakeRun(*this);
+      if (UNLIKELY(state_ == State::Close)) {
+        ClientFailureInvalidData();
+        return;
+      }
+      handshake_done_ = true;
+    }
+
+    ChunkState chunk_state;
+    while ((chunk_state = decoder_buffer_.GetChunk()) != ChunkState::Partial) {
+      if (chunk_state == ChunkState::Whole) {
+        // The chunk is whole, we need to read one more chunk
+        // (the 0x00 0x00 end marker).
+        continue;
       }
 
       switch (state_) {
-        case State::Handshake:
-          state_ = StateHandshakeRun(*this);
-          break;
         case State::Init:
           state_ = StateInitRun(*this);
           break;
@@ -119,15 +116,16 @@ class Session {
         case State::ErrorWaitForRollback:
           state_ = StateErrorRun(*this, state_);
           break;
-        case State::Close:
-          // This state is handled below
+        default:
+          // State::Handshake is handled above
+          // State::Close is handled below
           break;
       }
 
       // State::Close is handled here because we always want to check for
       // it after the above select. If any of the states above return a
       // State::Close then the connection should be terminated immediately.
-      if (state_ == State::Close) {
+      if (UNLIKELY(state_ == State::Close)) {
         ClientFailureInvalidData();
         return;
       }
@@ -200,7 +198,7 @@ class Session {
   Decoder<ChunkedDecoderBuffer> decoder_{decoder_buffer_};
 
   io::network::Epoll::Event event_;
-  bool connected_{false};
+  bool handshake_done_{false};
   State state_{State::Handshake};
   // GraphDbAccessor of active transaction in the session, can be null if there
   // is no associated transaction.
diff --git a/src/communication/bolt/v1/states/handshake.hpp b/src/communication/bolt/v1/states/handshake.hpp
index 4f307da93..1394e5ea2 100644
--- a/src/communication/bolt/v1/states/handshake.hpp
+++ b/src/communication/bolt/v1/states/handshake.hpp
@@ -28,7 +28,6 @@ State StateHandshakeRun(Session &session) {
     DLOG(WARNING) << "Couldn't write handshake response!";
     return State::Close;
   }
-  session.connected_ = true;
 
   // Delete data from buffer. It is guaranteed that there will more than, or
   // equal to 20 bytes (HANDSHAKE_SIZE) in the buffer.
diff --git a/tests/unit/bolt_chunked_decoder_buffer.cpp b/tests/unit/bolt_chunked_decoder_buffer.cpp
index 6d23b01cf..3aec55aff 100644
--- a/tests/unit/bolt_chunked_decoder_buffer.cpp
+++ b/tests/unit/bolt_chunked_decoder_buffer.cpp
@@ -24,6 +24,7 @@ TEST(BoltBuffer, CorrectChunk) {
   buffer.Written(1004);
 
   ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Whole);
+  ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Done);
 
   ASSERT_EQ(decoder_buffer.Read(tmp, 1000), true);
   for (int i = 0; i < 1000; ++i) EXPECT_EQ(data[i], tmp[i]);
@@ -45,6 +46,7 @@ TEST(BoltBuffer, CorrectChunkTrailingData) {
   buffer.Written(2004);
 
   ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Whole);
+  ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Done);
 
   ASSERT_EQ(decoder_buffer.Read(tmp, 1000), true);
   for (int i = 0; i < 1000; ++i) EXPECT_EQ(data[i], tmp[i]);
@@ -54,26 +56,6 @@ TEST(BoltBuffer, CorrectChunkTrailingData) {
   for (int i = 0; i < 1000; ++i) EXPECT_EQ(data[i + 1002], leftover[i]);
 }
 
-TEST(BoltBuffer, InvalidChunk) {
-  BufferT buffer;
-  DecoderBufferT decoder_buffer(buffer);
-  StreamBufferT sb = buffer.Allocate();
-
-  sb.data[0] = 0x03;
-  sb.data[1] = 0xe8;
-  memcpy(sb.data + 2, data, 2002);
-  sb.data[1002] = 1;
-  sb.data[1003] = 1;
-  buffer.Written(2004);
-
-  ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Invalid);
-
-  ASSERT_EQ(buffer.size(), 1000);
-
-  uint8_t *tmp = buffer.data();
-  for (int i = 0; i < 1000; ++i) EXPECT_EQ(data[i + 1002], tmp[i]);
-}
-
 TEST(BoltBuffer, GraduallyPopulatedChunk) {
   uint8_t tmp[2000];
   BufferT buffer;
@@ -83,13 +65,12 @@ TEST(BoltBuffer, GraduallyPopulatedChunk) {
   sb.data[0] = 0x03;
   sb.data[1] = 0xe8;
   buffer.Written(2);
-  ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Partial);
 
   for (int i = 0; i < 5; ++i) {
+    ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Partial);
     sb = buffer.Allocate();
     memcpy(sb.data, data + 200 * i, 200);
     buffer.Written(200);
-    ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Partial);
   }
 
   sb = buffer.Allocate();
@@ -97,6 +78,7 @@ TEST(BoltBuffer, GraduallyPopulatedChunk) {
   sb.data[1] = 0;
   buffer.Written(2);
   ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Whole);
+  ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Done);
 
   ASSERT_EQ(decoder_buffer.Read(tmp, 1000), true);
   for (int i = 0; i < 1000; ++i) EXPECT_EQ(data[i], tmp[i]);
@@ -113,13 +95,12 @@ TEST(BoltBuffer, GraduallyPopulatedChunkTrailingData) {
   sb.data[0] = 0x03;
   sb.data[1] = 0xe8;
   buffer.Written(2);
-  ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Partial);
 
   for (int i = 0; i < 5; ++i) {
+    ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Partial);
     sb = buffer.Allocate();
     memcpy(sb.data, data + 200 * i, 200);
     buffer.Written(200);
-    ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Partial);
   }
 
   sb = buffer.Allocate();
@@ -132,6 +113,7 @@ TEST(BoltBuffer, GraduallyPopulatedChunkTrailingData) {
   buffer.Written(1000);
 
   ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Whole);
+  ASSERT_EQ(decoder_buffer.GetChunk(), ChunkStateT::Done);
 
   ASSERT_EQ(decoder_buffer.Read(tmp, 1000), true);
   for (int i = 0; i < 1000; ++i) EXPECT_EQ(data[i], tmp[i]);
diff --git a/tests/unit/bolt_chunked_encoder_buffer.cpp b/tests/unit/bolt_chunked_encoder_buffer.cpp
index 90982bfb8..72db20f52 100644
--- a/tests/unit/bolt_chunked_encoder_buffer.cpp
+++ b/tests/unit/bolt_chunked_encoder_buffer.cpp
@@ -21,8 +21,10 @@ static constexpr auto WCS = communication::bolt::WHOLE_CHUNK_SIZE;
  * @param size of data array
  * @param element expected element in all positions of chunk data array
  *        (all data bytes in tested chunk should be equal to element)
+ * @param final_chunk if set to true then check for 0x00 0x00 after the chunk
  */
-void VerifyChunkOfOnes(uint8_t *data, int size, uint8_t element) {
+void VerifyChunkOfOnes(uint8_t *data, int size, uint8_t element,
+                       bool final_chunk = true) {
   // first two bytes are size (big endian)
   uint8_t lower_byte = size & 0xFF;
   uint8_t higher_byte = (size & 0xFF00) >> 8;
@@ -37,8 +39,10 @@ void VerifyChunkOfOnes(uint8_t *data, int size, uint8_t element) {
 
   // last two bytes should be zeros
   // next to header and data
-  ASSERT_EQ(*(data + CHS + size), 0x00);
-  ASSERT_EQ(*(data + CHS + size + 1), 0x00);
+  if (final_chunk) {
+    ASSERT_EQ(*(data + CHS + size), 0x00);
+    ASSERT_EQ(*(data + CHS + size + 1), 0x00);
+  }
 }
 
 TEST(BoltChunkedEncoderBuffer, OneSmallChunk) {
@@ -105,7 +109,7 @@ TEST(BoltChunkedEncoderBuffer, OneAndAHalfOfMaxChunk) {
   // the output array should look like this:
   // [0xFF, 0xFF, 1, 1, ... , 1, 0, 0, 0x86, 0xA1, 1, 1, ... , 1, 0, 0]
   auto output = socket.output.data();
-  VerifyChunkOfOnes(output, MCS, element);
+  VerifyChunkOfOnes(output, MCS, element, false);
   VerifyChunkOfOnes(output + WCS, size - MCS, element);
 }
 
diff --git a/tests/unit/bolt_session.cpp b/tests/unit/bolt_session.cpp
index 88a5c39ed..1d1861341 100644
--- a/tests/unit/bolt_session.cpp
+++ b/tests/unit/bolt_session.cpp
@@ -614,22 +614,6 @@ TEST(BoltSession, PartialChunk) {
   PrintOutput(output);
 }
 
-TEST(BoltSession, InvalidChunk) {
-  INIT_VARS;
-  ExecuteHandshake(session, output);
-  ExecuteInit(session, output);
-
-  // this will write 0x00 0x02 0x00 0x02 0x00 0x02
-  // that is a chunk of good size, but it's invalid because the last
-  // two bytes are 0x00 0x02 and they should be 0x00 0x00
-  for (int i = 0; i < 3; ++i) WriteChunkHeader(session, 2);
-  session.Execute();
-
-  ASSERT_EQ(session.state_, StateT::Close);
-  ASSERT_FALSE(session.socket_.IsOpen());
-  CheckFailureMessage(output);
-}
-
 TEST(BoltSession, ExplicitTransactionValidQueries) {
   // It is not really easy to check if we commited or aborted transaction except
   // by faking GraphDb/TxEngine...