Skip to content

Commit

Permalink
MINOR: Add more validation during KRPC deserialization
Browse files Browse the repository at this point in the history
When deserializing KRPC (which is used for RPCs sent to Kafka, Kafka Metadata records, and some
other things), check that we have at least N bytes remaining before allocating an array of size N.

Remove DataInputStreamReadable since it was hard to make this class aware of how many bytes were
remaining. Instead, when reading an individual record in the Raft layer, simply create a
ByteBufferAccessor with a ByteBuffer containing just the bytes we're interested in.

Add SimpleArraysMessageTest and ByteBufferAccessorTest. Also add some additional tests in
RequestResponseTest.

Reviewers: Tom Bentley <[email protected]>, Mickael Maison <[email protected]>, Colin McCabe <[email protected]>

Co-authored-by: Colin McCabe <[email protected]>
Co-authored-by: Manikumar Reddy <[email protected]>
Co-authored-by: Mickael Maison <[email protected]>
  • Loading branch information
4 people committed Aug 17, 2022
1 parent 967b89f commit aaceb6b
Show file tree
Hide file tree
Showing 16 changed files with 434 additions and 187 deletions.
4 changes: 4 additions & 0 deletions checkstyle/suppressions.xml
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,10 @@
<suppress checks="JavaNCSS"
files="DistributedHerderTest.java"/>

<!-- Raft -->
<suppress checks="NPathComplexity"
files="RecordsIterator.java"/>

<!-- Streams -->
<suppress checks="ClassFanOutComplexity"
files="(KafkaStreams|KStreamImpl|KTableImpl|StreamsPartitionAssignor).java"/>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,8 +54,15 @@ public double readDouble() {
}

@Override
public void readArray(byte[] arr) {
public byte[] readArray(int size) {
int remaining = buf.remaining();
if (size > remaining) {
throw new RuntimeException("Error reading byte array of " + size + " byte(s): only " + remaining +
" byte(s) available");
}
byte[] arr = new byte[size];
buf.get(arr);
return arr;
}

@Override
Expand Down

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -32,25 +32,23 @@ public interface Readable {
int readInt();
long readLong();
double readDouble();
void readArray(byte[] arr);
byte[] readArray(int length);
int readUnsignedVarint();
ByteBuffer readByteBuffer(int length);
int readVarint();
long readVarlong();
int remaining();

default String readString(int length) {
byte[] arr = new byte[length];
readArray(arr);
byte[] arr = readArray(length);
return new String(arr, StandardCharsets.UTF_8);
}

default List<RawTaggedField> readUnknownTaggedField(List<RawTaggedField> unknowns, int tag, int size) {
if (unknowns == null) {
unknowns = new ArrayList<>();
}
byte[] data = new byte[size];
readArray(data);
byte[] data = readArray(size);
unknowns.add(new RawTaggedField(tag, data));
return unknowns;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,8 @@ private static DefaultRecord readFrom(ByteBuffer buffer,
int numHeaders = ByteUtils.readVarint(buffer);
if (numHeaders < 0)
throw new InvalidRecordException("Found invalid number of record headers " + numHeaders);
if (numHeaders > buffer.remaining())
throw new InvalidRecordException("Found invalid number of record headers. " + numHeaders + " is larger than the remaining size of the buffer");

final Header[] headers;
if (numHeaders == 0)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
/*
* 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.kafka.common.message;

import org.apache.kafka.common.protocol.ByteBufferAccessor;
import org.junit.jupiter.api.Test;

import java.nio.ByteBuffer;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;

public class SimpleArraysMessageTest {
@Test
public void testArrayBoundsChecking() {
// SimpleArraysMessageData takes 2 arrays
final ByteBuffer buf = ByteBuffer.wrap(new byte[] {
(byte) 0x7f, // Set size of first array to 126 which is larger than the size of this buffer
(byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00
});
final SimpleArraysMessageData out = new SimpleArraysMessageData();
ByteBufferAccessor accessor = new ByteBufferAccessor(buf);
assertEquals("Tried to allocate a collection of size 126, but there are only 7 bytes remaining.",
assertThrows(RuntimeException.class, () -> out.read(accessor, (short) 2)).getMessage());
}

@Test
public void testArrayBoundsCheckingOtherArray() {
// SimpleArraysMessageData takes 2 arrays
final ByteBuffer buf = ByteBuffer.wrap(new byte[] {
(byte) 0x01, // Set size of first array to 0
(byte) 0x7e, // Set size of second array to 125 which is larger than the size of this buffer
(byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00
});
final SimpleArraysMessageData out = new SimpleArraysMessageData();
ByteBufferAccessor accessor = new ByteBufferAccessor(buf);
assertEquals("Tried to allocate a collection of size 125, but there are only 6 bytes remaining.",
assertThrows(RuntimeException.class, () -> out.read(accessor, (short) 2)).getMessage());
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
/*
* 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.kafka.common.protocol;

import org.junit.jupiter.api.Test;

import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;

public class ByteBufferAccessorTest {
@Test
public void testReadArray() {
ByteBuffer buf = ByteBuffer.allocate(1024);
ByteBufferAccessor accessor = new ByteBufferAccessor(buf);
final byte[] testArray = new byte[] {0x4b, 0x61, 0x46};
accessor.writeByteArray(testArray);
accessor.writeInt(12345);
accessor.flip();
final byte[] testArray2 = accessor.readArray(3);
assertArrayEquals(testArray, testArray2);
assertEquals(12345, accessor.readInt());
assertEquals("Error reading byte array of 3 byte(s): only 0 byte(s) available",
assertThrows(RuntimeException.class,
() -> accessor.readArray(3)).getMessage());
}

@Test
public void testReadString() {
ByteBuffer buf = ByteBuffer.allocate(1024);
ByteBufferAccessor accessor = new ByteBufferAccessor(buf);
String testString = "ABC";
final byte[] testArray = testString.getBytes(StandardCharsets.UTF_8);
accessor.writeByteArray(testArray);
accessor.flip();
assertEquals("ABC", accessor.readString(3));
assertEquals("Error reading byte array of 2 byte(s): only 0 byte(s) available",
assertThrows(RuntimeException.class,
() -> accessor.readString(2)).getMessage());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -247,6 +247,20 @@ public void testInvalidNumHeaders() {
buf.flip();
assertThrows(InvalidRecordException.class,
() -> DefaultRecord.readFrom(buf, 0L, 0L, RecordBatch.NO_SEQUENCE, null));

ByteBuffer buf2 = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
ByteUtils.writeVarint(sizeOfBodyInBytes, buf2);
buf2.put(attributes);
ByteUtils.writeVarlong(timestampDelta, buf2);
ByteUtils.writeVarint(offsetDelta, buf2);
ByteUtils.writeVarint(-1, buf2); // null key
ByteUtils.writeVarint(-1, buf2); // null value
ByteUtils.writeVarint(sizeOfBodyInBytes, buf2); // more headers than remaining buffer size, not allowed
buf2.position(buf2.limit());

buf2.flip();
assertThrows(InvalidRecordException.class,
() -> DefaultRecord.readFrom(buf2, 0L, 0L, RecordBatch.NO_SEQUENCE, null));
}

@Test
Expand Down
Loading

0 comments on commit aaceb6b

Please sign in to comment.