Fix edge cases with UTF-8 strings in ChecksumResultSet

Author: Noremac201

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ChecksumResultSet.java

@@ -329,7 +329,7 @@ private void putString(String stringValue) {
       if (buffer == null || (buffer.capacity() < MAX_BUFFER_SIZE && buffer.capacity() < length)) {
         // Create a ByteBuffer with a maximum buffer size.
         // This buffer is re-used for all string values in the result set.
-        buffer = ByteBuffer.allocate(Math.min(MAX_BUFFER_SIZE, length));
+        buffer = ByteBuffer.allocate(Math.max(4, Math.min(MAX_BUFFER_SIZE, length)));
       } else {
         buffer.clear();
       }
@@ -349,8 +349,8 @@ private void putString(String stringValue) {
         buffer.flip();
         // Put the bytes from the buffer into the digest.
         digest.update(buffer);
-        // Flip the buffer again, so we can repeat and write to the start of the buffer again.
-        buffer.flip();
+        // Clear the buffer, so we can repeat and write to the start of the buffer again.
+        buffer.clear();
       }
     }
   }

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ChecksumResultSetTest.java

@@ -448,4 +448,74 @@ public void testRetry() {
           () -> resultSet.retry(abortedException));
     }
   }
+
+  @Test
+  public void testEmptyString() {
+    Type type = Type.struct(StructField.of("stringVal", Type.string()));
+    Struct row = Struct.newBuilder().set("stringVal").to("").build();
+
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    Statement statement = Statement.of("select * from foo");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+    ReadWriteTransaction transaction = mock(ReadWriteTransaction.class);
+    when(transaction.runWithRetry(any(Callable.class)))
+        .thenAnswer(invocationOnMock -> ((Callable<?>) invocationOnMock.getArgument(0)).call());
+    when(transaction.getStatementExecutor()).thenReturn(mock(StatementExecutor.class));
+
+    ResultSet queryResult = ResultSets.forRows(type, ImmutableList.of(row));
+    ChecksumResultSet resultSet =
+        new ChecksumResultSet(
+            transaction,
+            DirectExecuteResultSet.ofResultSet(queryResult),
+            parsedStatement,
+            AnalyzeMode.NONE);
+    assertTrue(resultSet.next());
+  }
+
+  @Test
+  public void testSingleCharMultiByteString() {
+    Type type = Type.struct(StructField.of("stringVal", Type.string()));
+    // "ä" is 1 char, but 2 bytes in UTF-8.
+    Struct row = Struct.newBuilder().set("stringVal").to("ä").build();
+
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    Statement statement = Statement.of("select * from foo");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+    ReadWriteTransaction transaction = mock(ReadWriteTransaction.class);
+    when(transaction.runWithRetry(any(Callable.class)))
+        .thenAnswer(invocationOnMock -> ((Callable<?>) invocationOnMock.getArgument(0)).call());
+    when(transaction.getStatementExecutor()).thenReturn(mock(StatementExecutor.class));
+
+    ResultSet queryResult = ResultSets.forRows(type, ImmutableList.of(row));
+    ChecksumResultSet resultSet =
+        new ChecksumResultSet(
+            transaction,
+            DirectExecuteResultSet.ofResultSet(queryResult),
+            parsedStatement,
+            AnalyzeMode.NONE);
+    assertTrue(resultSet.next());
+  }
+
+  @Test
+  public void testLongMixedUtf8String() {
+    Type type = Type.struct(StructField.of("stringVal", Type.string()));
+    Struct row = Struct.newBuilder().set("stringVal").to("aaa\uD841\uDF0E").build();
+
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    Statement statement = Statement.of("select * from foo");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+    ReadWriteTransaction transaction = mock(ReadWriteTransaction.class);
+    when(transaction.runWithRetry(any(Callable.class)))
+        .thenAnswer(invocationOnMock -> ((Callable<?>) invocationOnMock.getArgument(0)).call());
+    when(transaction.getStatementExecutor()).thenReturn(mock(StatementExecutor.class));
+
+    ResultSet queryResult = ResultSets.forRows(type, ImmutableList.of(row));
+    ChecksumResultSet resultSet =
+        new ChecksumResultSet(
+            transaction,
+            DirectExecuteResultSet.ofResultSet(queryResult),
+            parsedStatement,
+            AnalyzeMode.NONE);
+    assertTrue(resultSet.next());
+  }
 }