Reject truncated Thrift containers before allocation

iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java

@@ -35,6 +35,9 @@ public final class RpcMessages {
       "Frame size (%d) larger than protect max size (%d)%s!";
   public static final String FRAME_ERROR_STRING_LENGTH_EXCEEDED =
       "String length (%d) larger than protect max size (%d)%s!";
+  public static final String
+      EXCEPTION_REQUIRED_READ_SIZE_ARG_EXCEEDS_REMAINING_FRAME_SIZE_ARG_ARG_9C0541EE =
+          "Required read size (%d) exceeds remaining frame size (%d)%s!";
 
   // TElasticFramedTransport - SSL
   public static final String NON_SSL_TO_SSL_PORT =

iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java

@@ -32,6 +32,9 @@ public final class RpcMessages {
       "帧大小 (%d) 超过保护最大值 (%d)%s！";
   public static final String FRAME_ERROR_STRING_LENGTH_EXCEEDED =
       "字符串长度 (%d) 超过保护最大值 (%d)%s！";
+  public static final String
+      EXCEPTION_REQUIRED_READ_SIZE_ARG_EXCEEDS_REMAINING_FRAME_SIZE_ARG_ARG_9C0541EE =
+          "请求读取的大小 (%d) 超过当前帧剩余大小 (%d)%s！";
 
   // TElasticFramedTransport - SSL
   public static final String NON_SSL_TO_SSL_PORT =

iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java

@@ -245,7 +245,9 @@ private enum FrameError {
     TLS_REQUEST(RpcMessages.FRAME_ERROR_TLS_REQUEST),
     NEGATIVE_FRAME_SIZE(RpcMessages.FRAME_ERROR_NEGATIVE_FRAME_SIZE),
     FRAME_SIZE_EXCEEDED(RpcMessages.FRAME_ERROR_FRAME_SIZE_EXCEEDED),
-    STRING_LENGTH_EXCEEDED(RpcMessages.FRAME_ERROR_STRING_LENGTH_EXCEEDED);
+    STRING_LENGTH_EXCEEDED(RpcMessages.FRAME_ERROR_STRING_LENGTH_EXCEEDED),
+    INSUFFICIENT_FRAME_DATA(
+        RpcMessages.EXCEPTION_REQUIRED_READ_SIZE_ARG_EXCEEDS_REMAINING_FRAME_SIZE_ARG_ARG_9C0541EE);
 
     private final String messageFormat;
 
@@ -255,7 +257,9 @@ private enum FrameError {
 
     void throwException(long size, String remoteInfo, int maxSize) throws TTransportException {
       String message =
-          (this == FRAME_SIZE_EXCEEDED || this == STRING_LENGTH_EXCEEDED)
+          (this == FRAME_SIZE_EXCEEDED
+                  || this == STRING_LENGTH_EXCEEDED
+                  || this == INSUFFICIENT_FRAME_DATA)
               ? String.format(messageFormat, size, maxSize, remoteInfo)
               : String.format(messageFormat, size, remoteInfo);
       throw new TTransportException(TTransportException.CORRUPTED_DATA, message);
@@ -308,18 +312,32 @@ public void updateKnownMessageSize(long size) throws TTransportException {
 
   @Override
   public void checkReadBytesAvailable(long numBytes) throws TTransportException {
+    // RPC messages are flushed as complete frames. Container checks pass their minimum encoded
+    // size here, before generated code allocates the container. Compare it with actual buffered
+    // data, not just the configured frame cap. Use readBuffer directly because copyBinary makes
+    // this transport's getBytesRemainingInBuffer() return -1.
+    int remaining = readBuffer.getBytesRemainingInBuffer();
+    FrameError error;
+    int limit;
     if (numBytes >= thriftMaxFrameSize) {
-      SocketAddress remoteAddress = null;
-      if (underlying instanceof TSocket) {
-        remoteAddress = ((TSocket) underlying).getSocket().getRemoteSocketAddress();
-      }
-      String remoteInfo =
-          (remoteAddress == null)
-              ? RpcMessages.EMPTY_MESSAGE
-              : RpcMessages.REMOTE_ADDRESS_PREFIX + remoteAddress;
-      close();
-      FrameError.STRING_LENGTH_EXCEEDED.throwException(numBytes, remoteInfo, thriftMaxFrameSize);
+      error = FrameError.STRING_LENGTH_EXCEEDED;
+      limit = thriftMaxFrameSize;
+    } else if (numBytes > remaining) {
+      error = FrameError.INSUFFICIENT_FRAME_DATA;
+      limit = remaining;
+    } else {
+      return;
+    }
+    SocketAddress remoteAddress = null;
+    if (underlying instanceof TSocket) {
+      remoteAddress = ((TSocket) underlying).getSocket().getRemoteSocketAddress();
     }
+    String remoteInfo =
+        (remoteAddress == null)
+            ? RpcMessages.EMPTY_MESSAGE
+            : RpcMessages.REMOTE_ADDRESS_PREFIX + remoteAddress;
+    close();
+    error.throwException(numBytes, remoteInfo, limit);
   }
 
   @Override

iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TElasticFramedTransportReadTest.java

@@ -0,0 +1,266 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.rpc;
+
+import org.apache.iotdb.rpc.i18n.RpcMessages;
+import org.apache.iotdb.service.rpc.thrift.IClientRPCService;
+import org.apache.iotdb.service.rpc.thrift.TSInsertRecordReq;
+
+import org.apache.thrift.TException;
+import org.apache.thrift.protocol.TBinaryProtocol;
+import org.apache.thrift.protocol.TCompactProtocol;
+import org.apache.thrift.protocol.TField;
+import org.apache.thrift.protocol.TList;
+import org.apache.thrift.protocol.TMap;
+import org.apache.thrift.protocol.TMessage;
+import org.apache.thrift.protocol.TMessageType;
+import org.apache.thrift.protocol.TProtocol;
+import org.apache.thrift.protocol.TSet;
+import org.apache.thrift.protocol.TStruct;
+import org.apache.thrift.protocol.TType;
+import org.apache.thrift.transport.TMemoryBuffer;
+import org.apache.thrift.transport.TMemoryInputTransport;
+import org.apache.thrift.transport.TTransport;
+import org.apache.thrift.transport.TTransportException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+
+@RunWith(Parameterized.class)
+public class TElasticFramedTransportReadTest {
+  private static final int MAX_FRAME_SIZE = 1024 * 1024;
+  private static final int ELEMENT_COUNT = 200_001;
+
+  private final boolean compact;
+  private final boolean copyBinary;
+  private final boolean snappy;
+
+  @Parameterized.Parameters(name = "compact={0}, copyBinary={1}, snappy={2}")
+  public static Collection<Object[]> parameters() {
+    Collection<Object[]> parameters = new ArrayList<>();
+    for (boolean compact : new boolean[] {false, true}) {
+      for (boolean copyBinary : new boolean[] {false, true}) {
+        for (boolean snappy : new boolean[] {false, true}) {
+          parameters.add(new Object[] {compact, copyBinary, snappy});
+        }
+      }
+    }
+    return parameters;
+  }
+
+  public TElasticFramedTransportReadTest(boolean compact, boolean copyBinary, boolean snappy) {
+    this.compact = compact;
+    this.copyBinary = copyBinary;
+    this.snappy = snappy;
+  }
+
+  @Test
+  public void testTruncatedRequestRejectedBeforeContainerAllocation() throws Exception {
+    byte[] frame =
+        serialize(
+            p -> {
+              p.writeMessageBegin(new TMessage("insertRecord", TMessageType.CALL, 1));
+              p.writeStructBegin(new TStruct("insertRecord_args"));
+              p.writeFieldBegin(new TField("req", TType.STRUCT, (short) 1));
+              p.writeStructBegin(new TStruct("TSInsertRecordReq"));
+              p.writeFieldBegin(new TField("columnCategoryies", TType.LIST, (short) 8));
+              p.writeListBegin(new TList(TType.BYTE, ELEMENT_COUNT));
+              // Flush a complete frame that ends at the container header, without any elements.
+            });
+    try (TElasticFramedTransport transport = transport(new TMemoryInputTransport(frame))) {
+      TProtocol protocol = protocol(transport);
+      assertEquals("insertRecord", protocol.readMessageBegin().name);
+      IClientRPCService.insertRecord_args args = new IClientRPCService.insertRecord_args();
+      TTransportException exception =
+          assertThrows(TTransportException.class, () -> args.read(protocol));
+      assertNotNull(args.req);
+      assertNull(args.req.columnCategoryies);
+      assertInsufficientFrame(exception, ELEMENT_COUNT, 0);
+    }
+  }
+
+  @Test
+  public void testLargeCompleteRequestsInSuccessiveFrames() throws Exception {
+    TSInsertRecordReq large =
+        new TSInsertRecordReq(1, "root.sg.d", Collections.emptyList(), ByteBuffer.allocate(0), 1)
+            .setColumnCategoryies(Collections.nCopies(ELEMENT_COUNT, (byte) 0));
+    TSInsertRecordReq empty =
+        new TSInsertRecordReq(1, "root.sg.d", Collections.emptyList(), ByteBuffer.allocate(0), 2)
+            .setColumnCategoryies(Collections.emptyList());
+    TMemoryBuffer wire = new TMemoryBuffer(128);
+    try (TElasticFramedTransport output = transport(wire)) {
+      IClientRPCService.Client client = new IClientRPCService.Client(protocol(output));
+      client.send_insertRecord(large);
+      client.send_insertRecord(empty);
+    }
+    try (TElasticFramedTransport input =
+        transport(new TMemoryInputTransport(Arrays.copyOf(wire.getArray(), wire.length())))) {
+      TProtocol protocol = protocol(input);
+      for (TSInsertRecordReq expected : new TSInsertRecordReq[] {large, empty}) {
+        assertEquals("insertRecord", protocol.readMessageBegin().name);
+        IClientRPCService.insertRecord_args args = new IClientRPCService.insertRecord_args();
+        args.read(protocol);
+        protocol.readMessageEnd();
+        assertEquals(expected, args.req);
+      }
+    }
+  }
+
+  @Test
+  public void testTruncatedContainerHeaders() throws Exception {
+    for (byte elementType : new byte[] {TType.BYTE, TType.STRUCT, TType.LIST, TType.MAP}) {
+      for (byte containerType : new byte[] {TType.LIST, TType.SET, TType.MAP}) {
+        byte[] frame =
+            serialize(
+                p -> {
+                  if (containerType == TType.LIST) {
+                    p.writeListBegin(new TList(elementType, ELEMENT_COUNT));
+                  } else if (containerType == TType.SET) {
+                    p.writeSetBegin(new TSet(elementType, ELEMENT_COUNT));
+                  } else {
+                    p.writeMapBegin(new TMap(TType.BYTE, elementType, ELEMENT_COUNT));
+                  }
+                });
+        try (TElasticFramedTransport input = transport(new TMemoryInputTransport(frame))) {
+          TProtocol p = protocol(input);
+          long minimumBytes =
+              (long) ELEMENT_COUNT
+                  * (p.getMinSerializedSize(elementType) + (containerType == TType.MAP ? 1 : 0));
+          TTransportException exception =
+              assertThrows(
+                  TTransportException.class,
+                  () -> {
+                    if (containerType == TType.LIST) {
+                      p.readListBegin();
+                    } else if (containerType == TType.SET) {
+                      p.readSetBegin();
+                    } else {
+                      p.readMapBegin();
+                    }
+                  });
+          assertInsufficientFrame(exception, minimumBytes, 0);
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testCompleteBinaryAtFrameBoundary() throws Exception {
+    ByteBuffer binary = ByteBuffer.wrap(new byte[ELEMENT_COUNT]);
+    byte[] frame =
+        serialize(
+            p -> {
+              p.writeString("");
+              p.writeBinary(binary);
+            });
+    try (TElasticFramedTransport input = transport(new TMemoryInputTransport(frame))) {
+      TProtocol protocol = protocol(input);
+      assertEquals("", protocol.readString());
+      assertEquals(binary, protocol.readBinary());
+      input.checkReadBytesAvailable(0);
+    }
+  }
+
+  @Test
+  public void testTruncatedStringAndBinary() throws Exception {
+    TMemoryBuffer buffer = new TMemoryBuffer(128);
+    protocol(buffer).writeBinary(ByteBuffer.wrap(new byte[ELEMENT_COUNT]));
+    int headerLength = buffer.length() - ELEMENT_COUNT;
+    byte[] frame = serialize(p -> p.getTransport().write(buffer.getArray(), 0, headerLength));
+    for (boolean binary : new boolean[] {false, true}) {
+      try (TElasticFramedTransport input = transport(new TMemoryInputTransport(frame))) {
+        TProtocol p = protocol(input);
+        TTransportException exception =
+            assertThrows(
+                TTransportException.class,
+                () -> {
+                  if (binary) {
+                    p.readBinary();
+                  } else {
+                    p.readString();
+                  }
+                });
+        assertInsufficientFrame(exception, ELEMENT_COUNT, 0);
+      }
+    }
+  }
+
+  @Test
+  public void testExistingMaximumReadSizeProtection() throws Exception {
+    try (TElasticFramedTransport input = transport(new TMemoryInputTransport(new byte[0]))) {
+      TTransportException exception =
+          assertThrows(
+              TTransportException.class, () -> input.checkReadBytesAvailable(MAX_FRAME_SIZE));
+      assertEquals(TTransportException.CORRUPTED_DATA, exception.getType());
+      assertEquals(
+          String.format(
+              RpcMessages.FRAME_ERROR_STRING_LENGTH_EXCEEDED, MAX_FRAME_SIZE, MAX_FRAME_SIZE, ""),
+          exception.getMessage());
+    }
+  }
+
+  private void assertInsufficientFrame(
+      TTransportException exception, long required, int remaining) {
+    assertEquals(TTransportException.CORRUPTED_DATA, exception.getType());
+    assertEquals(
+        String.format(
+            RpcMessages
+                .EXCEPTION_REQUIRED_READ_SIZE_ARG_EXCEEDS_REMAINING_FRAME_SIZE_ARG_ARG_9C0541EE,
+            required,
+            remaining,
+            ""),
+        exception.getMessage());
+  }
+
+  private TElasticFramedTransport transport(TTransport underlying) throws TTransportException {
+    return snappy
+        ? new TSnappyElasticFramedTransport(underlying, 128, MAX_FRAME_SIZE, copyBinary)
+        : new TElasticFramedTransport(underlying, 128, MAX_FRAME_SIZE, copyBinary);
+  }
+
+  private TProtocol protocol(TTransport transport) {
+    return compact ? new TCompactProtocol(transport) : new TBinaryProtocol(transport);
+  }
+
+  private byte[] serialize(ProtocolWriter writer) throws TException {
+    TMemoryBuffer wire = new TMemoryBuffer(128);
+    try (TElasticFramedTransport output = transport(wire)) {
+      writer.write(protocol(output));
+      output.flush();
+    }
+    return Arrays.copyOf(wire.getArray(), wire.length());
+  }
+
+  @FunctionalInterface
+  private interface ProtocolWriter {
+    void write(TProtocol protocol) throws TException;
+  }
+}