Fix STR8 encoding and underflow check in packString

xerial · xerial · commit 9cc49c07f40b · 2016-04-08T10:30:59.000-07:00
diff --git a/msgpack-core/src/main/java/org/msgpack/core/MessagePack.java b/msgpack-core/src/main/java/org/msgpack/core/MessagePack.java
@@ -16,6 +16,7 @@
 package org.msgpack.core;
 
 import org.msgpack.core.buffer.ArrayBufferInput;
+import org.msgpack.core.buffer.ArrayBufferOutput;
 import org.msgpack.core.buffer.ChannelBufferInput;
 import org.msgpack.core.buffer.ChannelBufferOutput;
 import org.msgpack.core.buffer.InputStreamBufferInput;
diff --git a/msgpack-core/src/main/java/org/msgpack/core/MessagePacker.java b/msgpack-core/src/main/java/org/msgpack/core/MessagePacker.java
@@ -457,6 +457,7 @@ private int encodeStringToBufferAt(int pos, String s)
         ByteBuffer bb = buffer.sliceAsByteBuffer(pos, buffer.size() - pos);
         int startPosition = bb.position();
         CharBuffer in = CharBuffer.wrap(s);
+        encoder.reset();
         CoderResult cr = encoder.encode(in, bb, true);
         if (cr.isError()) {
             try {
@@ -466,7 +467,8 @@ private int encodeStringToBufferAt(int pos, String s)
                 throw new MessageStringCodingException(e);
             }
         }
-        if (cr.isUnderflow() || cr.isOverflow()) {
+        if (!cr.isUnderflow() || cr.isOverflow()) {
+            // Underflow should be on to ensure all of the input string is encoded
             return -1;
         }
         return bb.position() - startPosition;
@@ -499,7 +501,7 @@ else if (s.length() < (1 << 8)) {
             // keep 2-byte header region and write raw string
             int written = encodeStringToBufferAt(position + 2, s);
             if (written >= 0) {
-                if (written < (1 << 8)) {
+                if (str8FormatSupport && written < (1 << 8)) {
                     buffer.putByte(position++, STR8);
                     buffer.putByte(position++, (byte) written);
                     position += written;
diff --git a/msgpack-core/src/test/scala/org/msgpack/core/MessagePackerTest.scala b/msgpack-core/src/test/scala/org/msgpack/core/MessagePackerTest.scala
@@ -282,4 +282,39 @@ class MessagePackerTest extends MessagePackSpec {
       .writePayload(payload)
       .close()
   }
+
+  "pack small string with STR8" in {
+    val packer = new PackerConfig().newBufferPacker()
+    packer.packString("Hello. This is a string longer than 32 characters!")
+    val b = packer.toByteArray
+
+    val unpacker = MessagePack.newDefaultUnpacker(b)
+    val f = unpacker.getNextFormat
+    f shouldBe MessageFormat.STR8
+  }
+
+  "be able to disable STR8 for backward compatibility" in {
+    val config = new PackerConfig()
+            .withSmallStringOptimizationThreshold(0) // Disable small string optimiztion
+            .withStr8FormatSupport(false)
+
+    {
+      val packer = config.newBufferPacker()
+      packer.packString("Hello. This is a string longer than 32 characters!")
+      val unpacker = MessagePack.newDefaultUnpacker(packer.toByteArray)
+      val f = unpacker.getNextFormat
+      f shouldBe MessageFormat.STR16
+    }
+
+    {
+      val packer2 = config.newBufferPacker()
+      packer2.packString("small string")
+      packer2.flush()
+      val unpacker2 = MessagePack.newDefaultUnpacker(packer2.toByteArray)
+      val f2 = unpacker2.getNextFormat
+      f2 shouldNot be (MessageFormat.STR8)
+      val s = unpacker2.unpackString()
+      s shouldBe "small string"
+    }
+  }
 }
