Protect BytesStreamOutput against overflows of the current number of written bytes. (#21174)

Closes #21159
This commit is contained in:
Adrien Grand 2016-11-02 10:11:14 +01:00 committed by GitHub
parent 9d6fac809c
commit acb12ecfa7
2 changed files with 16 additions and 12 deletions

View File

@ -69,7 +69,7 @@ public class BytesStreamOutput extends StreamOutput implements BytesStream {
@Override
public void writeByte(byte b) throws IOException {
ensureCapacity(count+1);
ensureCapacity(count + 1L);
bytes.set(count, b);
count++;
}
@ -87,7 +87,7 @@ public class BytesStreamOutput extends StreamOutput implements BytesStream {
}
// get enough pages for new size
ensureCapacity(count+length);
ensureCapacity(((long) count) + length);
// bulk copy
bytes.set(count, b, offset, length);
@ -113,18 +113,13 @@ public class BytesStreamOutput extends StreamOutput implements BytesStream {
}
@Override
public void seek(long position) throws IOException {
if (position > Integer.MAX_VALUE) {
throw new IllegalArgumentException("position " + position + " > Integer.MAX_VALUE");
}
count = (int)position;
ensureCapacity(count);
public void seek(long position) {
ensureCapacity(position);
count = (int) position;
}
public void skip(int length) {
count += length;
ensureCapacity(count);
seek(((long) count) + length);
}
@Override
@ -156,7 +151,10 @@ public class BytesStreamOutput extends StreamOutput implements BytesStream {
return bytes.ramBytesUsed();
}
private void ensureCapacity(int offset) {
private void ensureCapacity(long offset) {
if (offset > Integer.MAX_VALUE) {
throw new IllegalArgumentException(getClass().getSimpleName() + " cannot hold more than 2GB of data");
}
bytes = bigArrays.grow(bytes, offset);
}

View File

@ -238,6 +238,9 @@ public class BytesStreamsTests extends ESTestCase {
assertEquals(position, out.position());
assertEquals(position, BytesReference.toBytes(out.bytes()).length);
IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> out.seek(Integer.MAX_VALUE + 1L));
assertEquals("BytesStreamOutput cannot hold more than 2GB of data", iae.getMessage());
out.close();
}
@ -251,6 +254,9 @@ public class BytesStreamsTests extends ESTestCase {
out.skip(forward);
assertEquals(position + forward, out.position());
IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> out.skip(Integer.MAX_VALUE - 50));
assertEquals("BytesStreamOutput cannot hold more than 2GB of data", iae.getMessage());
out.close();
}