From 18782e67fd6a288439f88226ecc12df261a4ff36 Mon Sep 17 00:00:00 2001
From: Matej Ferencevic <matej.ferencevic@memgraph.io>
Date: Wed, 3 May 2017 14:16:56 +0200
Subject: [PATCH] Fixed bolt encoder test.

Summary: Fixed encoder off-by-one limit.

Reviewers: teon.banek

Reviewed By: teon.banek

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D335
---
 .../bolt/v1/encoder/base_encoder.hpp          |  2 +-
 tests/unit/bolt_encoder.cpp                   | 45 ++++++++++---------
 2 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/src/communication/bolt/v1/encoder/base_encoder.hpp b/src/communication/bolt/v1/encoder/base_encoder.hpp
index 312cbcc04..f72559786 100644
--- a/src/communication/bolt/v1/encoder/base_encoder.hpp
+++ b/src/communication/bolt/v1/encoder/base_encoder.hpp
@@ -97,7 +97,7 @@ class BaseEncoder : public Loggable {
       uint8_t len = size;
       WriteRAW(underlying_cast(Marker8[typ]));
       WriteRAW(len);
-    } else if (size <= 65536) {
+    } else if (size <= 65535) {
       uint16_t len = size;
       WriteRAW(underlying_cast(Marker16[typ]));
       WriteValue(len);
diff --git a/tests/unit/bolt_encoder.cpp b/tests/unit/bolt_encoder.cpp
index f1c311244..fa43ae139 100644
--- a/tests/unit/bolt_encoder.cpp
+++ b/tests/unit/bolt_encoder.cpp
@@ -15,27 +15,32 @@ using query::TypedValue;
 constexpr const int SIZE = 131072;
 uint8_t data[SIZE];
 
-void CheckTypeSize(std::vector<uint8_t> &v, int typ, uint64_t size) {
-  if (size <= 15) {
-    uint8_t len = size;
-    len &= 0x0F;
-    len += type_tiny_magic[typ];
-    CheckOutput(v, &len, 1, false);
-  } else if (size <= 255) {
-    uint8_t len = size;
-    CheckOutput(v, &type_8_magic[typ], 1, false);
-    CheckOutput(v, &len, 1, false);
-  } else if (size <= 65536) {
-    uint16_t len = size;
-    len = bswap(len);
-    CheckOutput(v, &type_16_magic[typ], 1, false);
-    CheckOutput(v, reinterpret_cast<const uint8_t *>(&len), 2, false);
-  } else {
-    uint32_t len = size;
-    len = bswap(len);
-    CheckOutput(v, &type_32_magic[typ], 1, false);
-    CheckOutput(v, reinterpret_cast<const uint8_t *>(&len), 4, false);
+uint64_t GetBigEndianInt(std::vector<uint8_t> &v, uint8_t len, uint8_t offset = 1) {
+  uint64_t ret = 0;
+  v.erase(v.begin(), v.begin() + offset);
+  for (int i = 0; i < len; ++i) {
+    ret <<= 8;
+    ret += v[i];
   }
+  v.erase(v.begin(), v.begin() + len);
+  return ret;
+}
+
+void CheckTypeSize(std::vector<uint8_t> &v, int typ, uint64_t size) {
+  uint64_t len;
+  if ((v[0] & 0xF0) == type_tiny_magic[typ]) {
+    len = v[0] & 0x0F;
+    v.erase(v.begin(), v.begin() + 1);
+  } else if (v[0] == type_8_magic[typ]) {
+    len = GetBigEndianInt(v, 1);
+  } else if (v[0] == type_16_magic[typ]) {
+    len = GetBigEndianInt(v, 2);
+  } else if (v[0] == type_32_magic[typ]) {
+    len = GetBigEndianInt(v, 4);
+  } else {
+    FAIL() << "Got wrong marker!";
+  }
+  ASSERT_EQ(len, size);
 }
 
 void CheckRecordHeader(std::vector<uint8_t> &v, uint64_t size) {