From f17ed977c33663723735c64af7518e64ab3bcde8 Mon Sep 17 00:00:00 2001
From: Dimitris Mandalidis <dimitris.mandalidis@gmail.com>
Date: Wed, 29 Jul 2020 12:44:01 +0300
Subject: [PATCH] Fix cast in StringCodec buffer size allocation #1367

Original pull request: #1368.
---
 .../io/lettuce/core/codec/StringCodec.java    | 27 ++++++++++++++-----
 .../core/codec/StringCodecUnitTests.java      | 12 +++++++++
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/main/java/io/lettuce/core/codec/StringCodec.java b/src/main/java/io/lettuce/core/codec/StringCodec.java
index 4ea4f4844c..fff9a9d241 100644
--- a/src/main/java/io/lettuce/core/codec/StringCodec.java
+++ b/src/main/java/io/lettuce/core/codec/StringCodec.java
@@ -98,9 +98,9 @@ public void encode(String str, ByteBuf target) {
             return;
         }
 
-        CharsetEncoder encoder = CharsetUtil.encoder(charset);
-        int length = (int) ((double) str.length() * encoder.maxBytesPerChar());
+        int length = calculateStringBytes(str, false);
         target.ensureWritable(length);
+        CharsetEncoder encoder = CharsetUtil.encoder(charset);
         try {
             final ByteBuffer dstBuf = target.nioBuffer(0, length);
             final int pos = dstBuf.position();
@@ -122,8 +122,7 @@ public void encode(String str, ByteBuf target) {
     public int estimateSize(Object keyOrValue) {
 
         if (keyOrValue instanceof String) {
-            CharsetEncoder encoder = CharsetUtil.encoder(charset);
-            return (int) (encoder.averageBytesPerChar() * ((String) keyOrValue).length());
+            return calculateStringBytes((String) keyOrValue, true);
         }
         return 0;
     }
@@ -164,8 +163,7 @@ private ByteBuffer encodeAndAllocateBuffer(String key) {
             return ByteBuffer.wrap(EMPTY);
         }
 
-        CharsetEncoder encoder = CharsetUtil.encoder(charset);
-        ByteBuffer buffer = ByteBuffer.allocate((int) (encoder.maxBytesPerChar() * key.length()));
+        ByteBuffer buffer = ByteBuffer.allocate(calculateStringBytes(key, false));
 
         ByteBuf byteBuf = Unpooled.wrappedBuffer(buffer);
         byteBuf.clear();
@@ -174,5 +172,20 @@ private ByteBuffer encodeAndAllocateBuffer(String key) {
 
         return buffer;
     }
-
+    
+    /**
+     * Calculate either the maximum number of bytes a string may occupy in a given character set or
+     * the average number of bytes it may hold.
+     * @param encoder the character set encoder (from which char to byte count association is inferred)
+     * @param value the actual value (must be not null)
+     * @param estimate whether the caller needs for an estimation or an actual value needed by buffer allocation
+     * @return the calculated string byte count
+     */
+    int calculateStringBytes(String value, boolean estimate) {
+        CharsetEncoder encoder = CharsetUtil.encoder(charset);
+        if (estimate) {
+            return (int) (encoder.averageBytesPerChar() * value.length());
+        }
+        return (int) encoder.maxBytesPerChar() * value.length();
+    }
 }
diff --git a/src/test/java/io/lettuce/core/codec/StringCodecUnitTests.java b/src/test/java/io/lettuce/core/codec/StringCodecUnitTests.java
index 6099e64a9d..530c52931f 100644
--- a/src/test/java/io/lettuce/core/codec/StringCodecUnitTests.java
+++ b/src/test/java/io/lettuce/core/codec/StringCodecUnitTests.java
@@ -18,12 +18,14 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import java.nio.ByteBuffer;
+import java.nio.charset.CharsetEncoder;
 import java.nio.charset.StandardCharsets;
 
 import org.junit.jupiter.api.Test;
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
+import io.netty.util.CharsetUtil;
 
 /**
  * @author Mark Paluch
@@ -116,4 +118,14 @@ void estimateSize() {
         assertThat(new StringCodec(StandardCharsets.US_ASCII).estimateSize(teststring)).isEqualTo(teststring.length());
         assertThat(new StringCodec(StandardCharsets.ISO_8859_1).estimateSize(teststring)).isEqualTo(teststring.length());
     }
+    
+    @Test
+    public void calculateStringSize() {
+        assertThat(new StringCodec(StandardCharsets.UTF_8).calculateStringBytes(teststring, false))
+            .isEqualTo(teststring.length() * 3);
+        assertThat(new StringCodec(StandardCharsets.US_ASCII).calculateStringBytes(teststring, false))
+            .isEqualTo(teststring.length());
+        assertThat(new StringCodec(StandardCharsets.ISO_8859_1).calculateStringBytes(teststring, false))
+            .isEqualTo(teststring.length());
+    }
 }