From c118cd38c7a18b136521090d61924b2303f2b994 Mon Sep 17 00:00:00 2001
From: cigui <cgcg96@sjtu.edu.cn>
Date: Fri, 15 Jul 2022 16:20:03 +0800
Subject: [PATCH] fix: compatibility mode for bin format family (#110)

---
 README.md                  |  2 +-
 lib/encoder.js             | 23 ++++++++++++++++++++++-
 test/compatibility-mode.js | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 448ac53..e330218 100644
--- a/README.md
+++ b/README.md
@@ -108,7 +108,7 @@ options:
 
 - `forceFloat64`, a boolean to that forces all floats to be encoded as 64-bits floats. Defaults to false.
 - `sortKeys`, a boolean to force a determinate keys order
-- `compatibilityMode`, a boolean that enables "compatibility mode" which doesn't use str 8 format. Defaults to false.
+- `compatibilityMode`, a boolean that enables "compatibility mode" which doesn't use bin format family and str 8 format. Defaults to false.
 - `disableTimestampEncoding`, a boolean that when set disables the encoding of Dates into the [timestamp extension type](https://github.com/msgpack/msgpack/blob/master/spec.md#timestamp-extension-type). Defaults to false.
 - `preferMap`, a boolean that forces all maps to be decoded to `Map`s rather than plain objects. This ensures that `decode(encode(new Map())) instanceof Map` and that iteration order is preserved. Defaults to false.
 - `protoAction`, a string which can be `error|ignore|remove` that determines what happens when decoding a plain object with a `__proto__` property which would cause prototype poisoning. `error` (default) throws an error, `remove` removes the property, `ignore` (not recommended) allows the property, thereby causing prototype poisoning on the decoded object.
diff --git a/lib/encoder.js b/lib/encoder.js
index 098b954..0ed7360 100644
--- a/lib/encoder.js
+++ b/lib/encoder.js
@@ -22,7 +22,8 @@ module.exports = function buildEncode (encodingTypes, options) {
       }
       // weird hack to support Buffer
       // and Buffer-like objects
-      return bl([getBufferHeader(obj.length), obj])
+      const _getBufferHeader = options.compatibilityMode ? getCompatibleBufferHeader : getBufferHeader
+      return bl([_getBufferHeader(obj.length), obj])
     }
     if (Array.isArray(obj)) return encodeArray(obj, encode)
     if (typeof obj === 'object') return encodeExt(obj, encodingTypes) || encodeObject(obj, options, encode)
@@ -225,6 +226,26 @@ function getBufferHeader (length) {
   return header
 }
 
+function getCompatibleBufferHeader (length) {
+  let header
+  if (length <= 0x1f) {
+    // fix raw header: 101XXXXX
+    header = Buffer.allocUnsafe(1)
+    header[0] = 0xa0 | length
+  } else if (length <= 0xffff) {
+    // raw 16 header: 0xda, XXXXXXXX, XXXXXXXX
+    header = Buffer.allocUnsafe(3)
+    header[0] = 0xda
+    header.writeUInt16BE(length, 1)
+  } else {
+    // raw 32 header: 0xdb, XXXXXXXX, XXXXXXXX, XXXXXXXX, XXXXXXXX
+    header = Buffer.allocUnsafe(5)
+    header[0] = 0xdb
+    header.writeUInt32BE(length, 1)
+  }
+  return header
+}
+
 function encodeNumber (obj, options) {
   let buf
   if (isFloat(obj)) return encodeFloat(obj, options.forceFloat64)
diff --git a/test/compatibility-mode.js b/test/compatibility-mode.js
index 743fee7..185a6ad 100644
--- a/test/compatibility-mode.js
+++ b/test/compatibility-mode.js
@@ -3,6 +3,13 @@
 const test = require('tape').test
 const msgpack = require('../')
 
+function buildBuffer (size) {
+  const buf = Buffer.allocUnsafe(size)
+  buf.fill('a')
+
+  return buf
+}
+
 test('encode/compatibility mode', function (t) {
   const compatEncoder = msgpack({
     compatibilityMode: true
@@ -37,4 +44,30 @@ test('encode/compatibility mode', function (t) {
     t.deepEqual(buf1, buf2, 'must be equal for two byte strings')
     t.end()
   })
+
+  const fixRawBuffer = buildBuffer(1)
+  const raw16Buffer = buildBuffer(Math.pow(2, 16) - 1)
+  const raw32Buffer = buildBuffer(Math.pow(2, 16) + 1)
+
+  t.test('compat. encoding a Buffer of length ' + fixRawBuffer.length, function (t) {
+    // fix raw header: 0xa0 | 1 = 0xa1
+    const buf = compatEncoder.encode(fixRawBuffer)
+    t.equal(buf[0], 0xa1, 'must have the proper header (fix raw)')
+    t.equal(buf.toString('utf8', 1, Buffer.byteLength(fixRawBuffer) + 1), fixRawBuffer.toString('utf8'), 'must decode correctly')
+    t.end()
+  })
+
+  t.test('compat. encoding a Buffer of length ' + raw16Buffer.length, function (t) {
+    const buf = compatEncoder.encode(raw16Buffer)
+    t.equal(buf[0], 0xda, 'must have the proper header (raw 16)')
+    t.equal(buf.toString('utf8', 3, Buffer.byteLength(raw16Buffer) + 3), raw16Buffer.toString('utf8'), 'must decode correctly')
+    t.end()
+  })
+
+  t.test('compat. encoding a Buffer of length ' + raw32Buffer.length, function (t) {
+    const buf = compatEncoder.encode(raw32Buffer)
+    t.equal(buf[0], 0xdb, 'must have the proper header (raw 32)')
+    t.equal(buf.toString('utf8', 5, Buffer.byteLength(raw32Buffer) + 5), raw32Buffer.toString('utf8'), 'must decode correctly')
+    t.end()
+  })
 })