Handle short reads from InputStream

java/tsfile/src/main/java/org/apache/tsfile/utils/BytesUtils.java

@@ -900,7 +900,19 @@ public static byte[] safeReadInputStreamToBytes(int count, InputStream in) throw
     byte[] bytes = new byte[count];
     int readCount = 0;
     while (readCount < count) {
-      readCount += in.read(bytes, readCount, count - readCount);
+      int currentRead = in.read(bytes, readCount, count - readCount);
+      if (currentRead < 0) {
+        throw new IOException(Messages.format("error.utils.rwio_short_read", count, readCount));
+      }
+      if (currentRead == 0) {
+        int value = in.read();
+        if (value < 0) {
+          throw new IOException(Messages.format("error.utils.rwio_short_read", count, readCount));
+        }
+        bytes[readCount++] = (byte) value;
+      } else {
+        readCount += currentRead;
+      }
     }
     return bytes;
   }

java/tsfile/src/main/java/org/apache/tsfile/utils/ReadWriteIOUtils.java

@@ -82,14 +82,14 @@ private ReadWriteIOUtils() {}
 
   /** read a bool from inputStream. */
   public static boolean readBool(InputStream inputStream) throws IOException {
-    int flag = inputStream.read();
+    int flag = readRequiredByte(inputStream);
     return flag == 1;
   }
 
   // used by generated code
   @SuppressWarnings("unused")
   public static boolean readBoolean(InputStream inputStream) throws IOException {
-    int flag = inputStream.read();
+    int flag = readRequiredByte(inputStream);
     return flag == 1;
   }
 
@@ -521,13 +521,13 @@ public static int sizeToWrite(String s) {
 
   /** read a byte var from inputStream. */
   public static byte readByte(InputStream inputStream) throws IOException {
-    return (byte) inputStream.read();
+    return (byte) readRequiredByte(inputStream);
   }
 
   /** read a short var from inputStream. */
   public static short readShort(InputStream inputStream) throws IOException {
     byte[] bytes = new byte[SHORT_LEN];
-    int readLen = inputStream.read(bytes);
+    int readLen = readAtMost(inputStream, bytes);
     if (readLen != SHORT_LEN) {
       throw new IOException(Messages.format(RETURN_ERROR_KEY, SHORT_LEN, readLen));
     }
@@ -542,7 +542,7 @@ public static short readShort(ByteBuffer buffer) {
   /** Read a float var from inputStream. */
   public static float readFloat(InputStream inputStream) throws IOException {
     byte[] bytes = new byte[FLOAT_LEN];
-    int readLen = inputStream.read(bytes);
+    int readLen = readAtMost(inputStream, bytes);
     if (readLen != FLOAT_LEN) {
       throw new IOException(Messages.format(RETURN_ERROR_KEY, FLOAT_LEN, readLen));
     }
@@ -559,7 +559,7 @@ public static float readFloat(ByteBuffer byteBuffer) {
   /** Read a double var from inputStream. */
   public static double readDouble(InputStream inputStream) throws IOException {
     byte[] bytes = new byte[DOUBLE_LEN];
-    int readLen = inputStream.read(bytes);
+    int readLen = readAtMost(inputStream, bytes);
     if (readLen != DOUBLE_LEN) {
       throw new IOException(Messages.format(RETURN_ERROR_KEY, DOUBLE_LEN, readLen));
     }
@@ -576,7 +576,7 @@ public static double readDouble(ByteBuffer byteBuffer) {
   /** Read an int var from inputStream. */
   public static int readInt(InputStream inputStream) throws IOException {
     byte[] bytes = new byte[INT_LEN];
-    int readLen = inputStream.read(bytes);
+    int readLen = readAtMost(inputStream, bytes);
     if (readLen != INT_LEN) {
       throw new IOException(Messages.format(RETURN_ERROR_KEY, INT_LEN, readLen));
     }
@@ -603,21 +603,9 @@ public static int read(ByteBuffer buffer) {
   /** Read a long var from inputStream. */
   public static long readLong(InputStream inputStream) throws IOException {
     byte[] bytes = new byte[LONG_LEN];
-    int readLen = 0;
-    while (readLen < LONG_LEN) {
-      int count = inputStream.read(bytes, readLen, LONG_LEN - readLen);
-      if (count < 0) {
-        throw new EOFException(Messages.format(RETURN_ERROR_KEY, LONG_LEN, readLen));
-      }
-      if (count == 0) {
-        int value = inputStream.read();
-        if (value < 0) {
-          throw new EOFException(Messages.format(RETURN_ERROR_KEY, LONG_LEN, readLen));
-        }
-        bytes[readLen++] = (byte) value;
-      } else {
-        readLen += count;
-      }
+    int readLen = readAtMost(inputStream, bytes);
+    if (readLen != LONG_LEN) {
+      throw new EOFException(Messages.format(RETURN_ERROR_KEY, LONG_LEN, readLen));
     }
     return BytesUtils.bytesToLong(bytes);
   }
@@ -636,7 +624,7 @@ public static String readString(InputStream inputStream) throws IOException {
       return "";
     }
     byte[] bytes = new byte[strLength];
-    int readLen = inputStream.read(bytes, 0, strLength);
+    int readLen = readAtMost(inputStream, bytes);
     if (readLen != strLength) {
       throw new IOException(Messages.format(RETURN_ERROR_KEY, strLength, readLen));
     }
@@ -652,13 +640,33 @@ public static String readVarIntString(InputStream inputStream) throws IOExceptio
       return "";
     }
     byte[] bytes = new byte[strLength];
-    int readLen = inputStream.read(bytes, 0, strLength);
+    int readLen = readAtMost(inputStream, bytes);
     if (readLen != strLength) {
       throw new IOException(Messages.format(RETURN_ERROR_KEY, strLength, readLen));
     }
     return new String(bytes, 0, strLength, TSFileConfig.STRING_CHARSET);
   }
 
+  private static int readAtMost(InputStream inputStream, byte[] bytes) throws IOException {
+    int readLen = 0;
+    while (readLen < bytes.length) {
+      int count = inputStream.read(bytes, readLen, bytes.length - readLen);
+      if (count < 0) {
+        break;
+      }
+      if (count == 0) {
+        int value = inputStream.read();
+        if (value < 0) {
+          break;
+        }
+        bytes[readLen++] = (byte) value;
+      } else {
+        readLen += count;
+      }
+    }
+    return readLen;
+  }
+
   /** Read string from byteBuffer. */
   public static String readString(ByteBuffer buffer) {
     return readStringWithLength(buffer, readInt(buffer));
@@ -718,15 +726,18 @@ public static String readStringFromDirectByteBuffer(ByteBuffer buffer)
    */
   public static byte[] readBytes(InputStream inputStream, int length) throws IOException {
     byte[] bytes = new byte[length];
-    int offset = 0;
-    int len;
-    while (bytes.length - offset > 0
-        && (len = inputStream.read(bytes, offset, bytes.length - offset)) != NO_BYTE_TO_READ) {
-      offset += len;
-    }
+    readAtMost(inputStream, bytes);
     return bytes;
   }
 
+  private static int readRequiredByte(InputStream inputStream) throws IOException {
+    int value = inputStream.read();
+    if (value < 0) {
+      throw new EOFException(Messages.format(RETURN_ERROR_KEY, 1, 0));
+    }
+    return value;
+  }
+
   public static Map<String, String> readMap(ByteBuffer buffer) {
     int length = readInt(buffer);
     if (length == NO_BYTE_TO_READ) {
@@ -829,20 +840,31 @@ public static byte[] readByteBufferWithSelfDescriptionLength(ByteBuffer buffer)
   public static int readAsPossible(TsFileInput input, long position, ByteBuffer buffer)
       throws IOException {
     int length = 0;
-    int read;
-    while (buffer.hasRemaining() && (read = input.read(buffer, position)) != NO_BYTE_TO_READ) {
+    while (buffer.hasRemaining()) {
+      int read = input.read(buffer, position);
+      if (read < 0) {
+        break;
+      }
+      if (read == 0) {
+        continue;
+      }
       length += read;
       position += read;
-      input.read(buffer, position);
     }
     return length;
   }
 
   /** read util to the end of buffer. */
   public static int readAsPossible(TsFileInput input, ByteBuffer buffer) throws IOException {
     int length = 0;
-    int read;
-    while (buffer.hasRemaining() && (read = input.read(buffer)) != NO_BYTE_TO_READ) {
+    while (buffer.hasRemaining()) {
+      int read = input.read(buffer);
+      if (read < 0) {
+        break;
+      }
+      if (read == 0) {
+        continue;
+      }
       length += read;
     }
     return length;
@@ -856,10 +878,14 @@ public static int readAsPossible(TsFileInput input, ByteBuffer target, long offs
     if (target.remaining() > len) {
       target.limit(target.position() + len);
     }
-    int read;
-    while (length < len
-        && target.hasRemaining()
-        && (read = input.read(target, offset)) != NO_BYTE_TO_READ) {
+    while (length < len && target.hasRemaining()) {
+      int read = input.read(target, offset);
+      if (read < 0) {
+        break;
+      }
+      if (read == 0) {
+        continue;
+      }
       length += read;
       offset += read;
     }

java/tsfile/src/test/java/org/apache/tsfile/utils/BytesUtilsTest.java

@@ -427,4 +427,70 @@ public void testReadBool() throws IOException {
     InputStream in = new ByteArrayInputStream(bs);
     assertEquals(l, BytesUtils.readBool(in));
   }
+
+  @Test
+  public void testInputStreamReadersHandleFragmentedStream() throws IOException {
+    int intValue = 0x12345678;
+    Assert.assertEquals(
+        intValue,
+        BytesUtils.readInt(
+            new FragmentedInputStream(BytesUtils.intToBytes(intValue), Integer.BYTES - 1)));
+
+    long longValue = 0x123456789ABCDEFL;
+    Assert.assertEquals(
+        longValue,
+        BytesUtils.readLong(
+            new FragmentedInputStream(BytesUtils.longToBytes(longValue), Long.BYTES - 1)));
+
+    float floatValue = 1.25F;
+    Assert.assertEquals(
+        floatValue,
+        BytesUtils.readFloat(
+            new FragmentedInputStream(BytesUtils.floatToBytes(floatValue), Float.BYTES - 1)),
+        0.0F);
+
+    double doubleValue = 1.25D;
+    Assert.assertEquals(
+        doubleValue,
+        BytesUtils.readDouble(
+            new FragmentedInputStream(BytesUtils.doubleToBytes(doubleValue), Double.BYTES - 1)),
+        0.0D);
+
+    Assert.assertTrue(
+        BytesUtils.readBool(
+            new FragmentedInputStream(BytesUtils.boolToBytes(true), Integer.BYTES - 3)));
+  }
+
+  @Test
+  public void testInputStreamReadersRejectTruncatedInput() {
+    for (int length = 0; length < Integer.BYTES; length++) {
+      final byte[] truncated = new byte[length];
+      Assert.assertThrows(
+          IOException.class, () -> BytesUtils.readInt(new ByteArrayInputStream(truncated)));
+    }
+    for (int length = 0; length < Long.BYTES; length++) {
+      final byte[] truncated = new byte[length];
+      Assert.assertThrows(
+          IOException.class, () -> BytesUtils.readLong(new ByteArrayInputStream(truncated)));
+    }
+  }
+
+  private static final class FragmentedInputStream extends ByteArrayInputStream {
+    private final int maxReadLength;
+    private boolean returnZeroOnce = true;
+
+    private FragmentedInputStream(byte[] bytes, int maxReadLength) {
+      super(bytes);
+      this.maxReadLength = maxReadLength;
+    }
+
+    @Override
+    public synchronized int read(byte[] bytes, int offset, int length) {
+      if (returnZeroOnce) {
+        returnZeroOnce = false;
+        return 0;
+      }
+      return super.read(bytes, offset, Math.min(length, maxReadLength));
+    }
+  }
 }