diff --git a/java/tsfile/src/main/java/org/apache/tsfile/utils/BytesUtils.java b/java/tsfile/src/main/java/org/apache/tsfile/utils/BytesUtils.java index 2cc68069c..f2d6e6383 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/utils/BytesUtils.java +++ b/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; } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/utils/ReadWriteIOUtils.java b/java/tsfile/src/main/java/org/apache/tsfile/utils/ReadWriteIOUtils.java index f27d35553..187cc616c 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/utils/ReadWriteIOUtils.java +++ b/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 readMap(ByteBuffer buffer) { int length = readInt(buffer); if (length == NO_BYTE_TO_READ) { @@ -829,11 +840,16 @@ 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; } @@ -841,8 +857,14 @@ public static int readAsPossible(TsFileInput input, long position, ByteBuffer bu /** 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; } diff --git a/java/tsfile/src/test/java/org/apache/tsfile/utils/BytesUtilsTest.java b/java/tsfile/src/test/java/org/apache/tsfile/utils/BytesUtilsTest.java index 76c01bee9..7ec012f4f 100644 --- a/java/tsfile/src/test/java/org/apache/tsfile/utils/BytesUtilsTest.java +++ b/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)); + } + } } diff --git a/java/tsfile/src/test/java/org/apache/tsfile/utils/ReadWriteIOUtilsTest.java b/java/tsfile/src/test/java/org/apache/tsfile/utils/ReadWriteIOUtilsTest.java index 0d1bd3667..01f747c6c 100644 --- a/java/tsfile/src/test/java/org/apache/tsfile/utils/ReadWriteIOUtilsTest.java +++ b/java/tsfile/src/test/java/org/apache/tsfile/utils/ReadWriteIOUtilsTest.java @@ -19,9 +19,11 @@ package org.apache.tsfile.utils; import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.read.reader.TsFileInput; import org.junit.Assert; import org.junit.Test; +import org.mockito.Mockito; import java.io.BufferedInputStream; import java.io.ByteArrayInputStream; @@ -74,6 +76,165 @@ public void testReadLongRejectsEveryTruncatedLength() { } } + @Test + public void testFixedWidthReadersFromFragmentedStream() throws IOException { + byte[] shortBytes = ByteBuffer.allocate(Short.BYTES).putShort((short) 0x1234).array(); + Assert.assertEquals( + (short) 0x1234, ReadWriteIOUtils.readShort(new FragmentedInputStream(shortBytes, 2))); + + byte[] intBytes = ByteBuffer.allocate(Integer.BYTES).putInt(0x12345678).array(); + Assert.assertEquals( + 0x12345678, ReadWriteIOUtils.readInt(new FragmentedInputStream(intBytes, 3))); + + byte[] floatBytes = ByteBuffer.allocate(Float.BYTES).putFloat(1.25F).array(); + Assert.assertEquals( + 1.25F, ReadWriteIOUtils.readFloat(new FragmentedInputStream(floatBytes, 3)), 0.0F); + + byte[] doubleBytes = ByteBuffer.allocate(Double.BYTES).putDouble(1.25D).array(); + Assert.assertEquals( + 1.25D, ReadWriteIOUtils.readDouble(new FragmentedInputStream(doubleBytes, 3)), 0.0D); + } + + @Test + public void testStringReadersFromFragmentedStream() throws IOException { + String expected = "fragmented-中文-value"; + + ByteArrayOutputStream fixedLengthBytes = new ByteArrayOutputStream(); + ReadWriteIOUtils.write(expected, fixedLengthBytes); + Assert.assertEquals( + expected, + ReadWriteIOUtils.readString(new FragmentedInputStream(fixedLengthBytes.toByteArray(), 2))); + + ByteArrayOutputStream varIntBytes = new ByteArrayOutputStream(); + ReadWriteIOUtils.writeVar(expected, varIntBytes); + Assert.assertEquals( + expected, + ReadWriteIOUtils.readVarIntString(new FragmentedInputStream(varIntBytes.toByteArray(), 2))); + } + + @Test + public void testReadBytesFromFragmentedStream() throws IOException { + byte[] expected = new byte[] {1, 2, 3, 4}; + Assert.assertArrayEquals( + expected, ReadWriteIOUtils.readBytes(new FragmentedInputStream(expected, 2), 4)); + } + + @Test + public void testReadAsPossibleDoesNotPerformExtraRead() throws IOException { + TsFileInput input = Mockito.mock(TsFileInput.class); + ByteBuffer buffer = ByteBuffer.allocate(Integer.BYTES); + Mockito.when(input.read(buffer, 0L)) + .thenAnswer( + invocation -> { + buffer.put(new byte[] {1, 2, 3, 4}); + return Integer.BYTES; + }); + + Assert.assertEquals(Integer.BYTES, ReadWriteIOUtils.readAsPossible(input, 0L, buffer)); + Mockito.verify(input, Mockito.times(1)).read(buffer, 0L); + } + + @Test + public void testReadAsPossibleRetriesZeroAtPosition() throws IOException { + TsFileInput input = Mockito.mock(TsFileInput.class); + ByteBuffer buffer = ByteBuffer.allocate(Integer.BYTES); + Mockito.when(input.read(buffer, 7L)) + .thenReturn(0) + .thenAnswer( + invocation -> { + buffer.put(new byte[] {1, 2, 3, 4}); + return Integer.BYTES; + }); + + Assert.assertEquals(Integer.BYTES, ReadWriteIOUtils.readAsPossible(input, 7L, buffer)); + Mockito.verify(input, Mockito.times(2)).read(buffer, 7L); + } + + @Test + public void testReadAsPossibleRetriesZeroAtCurrentPosition() throws IOException { + TsFileInput input = Mockito.mock(TsFileInput.class); + ByteBuffer buffer = ByteBuffer.allocate(Integer.BYTES); + Mockito.when(input.read(buffer)) + .thenReturn(0) + .thenAnswer( + invocation -> { + buffer.put(new byte[] {1, 2, 3, 4}); + return Integer.BYTES; + }); + + Assert.assertEquals(Integer.BYTES, ReadWriteIOUtils.readAsPossible(input, buffer)); + Mockito.verify(input, Mockito.times(2)).read(buffer); + } + + @Test + public void testReadAsPossibleRetriesZeroForLimitedRead() throws IOException { + TsFileInput input = Mockito.mock(TsFileInput.class); + ByteBuffer buffer = ByteBuffer.allocate(2 * Integer.BYTES); + Mockito.when(input.read(buffer, 11L)) + .thenReturn(0) + .thenAnswer( + invocation -> { + buffer.put(new byte[] {1, 2, 3, 4}); + return Integer.BYTES; + }); + + Assert.assertEquals(Integer.BYTES, ReadWriteIOUtils.readAsPossible(input, buffer, 11L, 4)); + Assert.assertEquals(2 * Integer.BYTES, buffer.limit()); + Mockito.verify(input, Mockito.times(2)).read(buffer, 11L); + } + + @Test + public void testFixedWidthReadersRejectEveryTruncatedLength() { + for (int length = 0; length < Integer.BYTES; length++) { + final byte[] truncated = new byte[length]; + Assert.assertThrows( + IOException.class, () -> ReadWriteIOUtils.readInt(new ByteArrayInputStream(truncated))); + } + for (int length = 0; length < Short.BYTES; length++) { + final byte[] truncated = new byte[length]; + Assert.assertThrows( + IOException.class, () -> ReadWriteIOUtils.readShort(new ByteArrayInputStream(truncated))); + } + for (int length = 0; length < Float.BYTES; length++) { + final byte[] truncated = new byte[length]; + Assert.assertThrows( + IOException.class, () -> ReadWriteIOUtils.readFloat(new ByteArrayInputStream(truncated))); + } + for (int length = 0; length < Double.BYTES; length++) { + final byte[] truncated = new byte[length]; + Assert.assertThrows( + IOException.class, + () -> ReadWriteIOUtils.readDouble(new ByteArrayInputStream(truncated))); + } + } + + @Test + public void testSingleByteReadersRejectEof() { + InputStream input = new ByteArrayInputStream(new byte[0]); + Assert.assertThrows(EOFException.class, () -> ReadWriteIOUtils.readByte(input)); + Assert.assertThrows(EOFException.class, () -> ReadWriteIOUtils.readBool(input)); + Assert.assertThrows(EOFException.class, () -> ReadWriteIOUtils.readBoolean(input)); + } + + 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)); + } + } + @Test public void testReadStringRejectsOversizedLength() { for (int length : new int[] {2, Integer.MAX_VALUE}) {