diff --git a/datalog/src/main/java/org/wpilib/datalog/DataLogIterator.java b/datalog/src/main/java/org/wpilib/datalog/DataLogIterator.java index 904f5d873f3..500887143bd 100644 --- a/datalog/src/main/java/org/wpilib/datalog/DataLogIterator.java +++ b/datalog/src/main/java/org/wpilib/datalog/DataLogIterator.java @@ -31,7 +31,7 @@ record = m_reader.getRecord(m_pos); @Override public boolean hasNext() { - return (m_pos + 16) <= m_reader.size(); + return m_reader.hasRecord(m_pos); } @Override diff --git a/datalog/src/main/java/org/wpilib/datalog/DataLogReader.java b/datalog/src/main/java/org/wpilib/datalog/DataLogReader.java index 154401ef52a..9fb12c03bad 100644 --- a/datalog/src/main/java/org/wpilib/datalog/DataLogReader.java +++ b/datalog/src/main/java/org/wpilib/datalog/DataLogReader.java @@ -113,42 +113,68 @@ private long readVarInt(int pos, int len) { } DataLogRecord getRecord(int pos) { + return getRecordInfo(pos).record; + } + + int getNextRecord(int pos) { + return getRecordInfo(pos).nextPos; + } + + boolean hasRecord(int pos) { try { + getRecordInfo(pos); + return true; + } catch (NoSuchElementException ex) { + return false; + } + } + + private RecordInfo getRecordInfo(int pos) { + try { + int remaining = m_buf.remaining(); + if (pos < 0 || pos >= remaining || remaining - pos < 4) { + throw new NoSuchElementException(); + } int lenbyte = m_buf.get(pos) & 0xff; int entryLen = (lenbyte & 0x3) + 1; int sizeLen = ((lenbyte >> 2) & 0x3) + 1; int timestampLen = ((lenbyte >> 4) & 0x7) + 1; int headerLen = 1 + entryLen + sizeLen + timestampLen; + if (headerLen > remaining - pos) { + throw new NoSuchElementException(); + } int entry = (int) readVarInt(pos + 1, entryLen); - int size = (int) readVarInt(pos + 1 + entryLen, sizeLen); + long size = readVarInt(pos + 1 + entryLen, sizeLen); + if (size > remaining - pos - headerLen) { + throw new NoSuchElementException(); + } + int checkedSize = (int) size; long timestamp = readVarInt(pos + 1 + entryLen + sizeLen, timestampLen); // build a slice of the data contents ByteBuffer data = m_buf.duplicate(); data.position(pos + headerLen); - data.limit(pos + headerLen + size); - return new DataLogRecord(entry, timestamp, data.slice()); + data.limit(pos + headerLen + checkedSize); + return new RecordInfo( + new DataLogRecord(entry, timestamp, data.slice()), + pos + headerLen + checkedSize); } catch (BufferUnderflowException | IndexOutOfBoundsException ex) { throw new NoSuchElementException(); } } - int getNextRecord(int pos) { - int lenbyte = m_buf.get(pos) & 0xff; - int entryLen = (lenbyte & 0x3) + 1; - int sizeLen = ((lenbyte >> 2) & 0x3) + 1; - int timestampLen = ((lenbyte >> 4) & 0x7) + 1; - int headerLen = 1 + entryLen + sizeLen + timestampLen; - - int size = 0; - for (int i = 0; i < sizeLen; i++) { - size |= (m_buf.get(pos + 1 + entryLen + i) & 0xff) << (i * 8); - } - return pos + headerLen + size; - } - int size() { return m_buf.remaining(); } + private static class RecordInfo { + RecordInfo(DataLogRecord record, int nextPos) { + this.record = record; + this.nextPos = nextPos; + } + + final DataLogRecord record; + final int nextPos; + } + private final ByteBuffer m_buf; } diff --git a/datalog/src/test/java/org/wpilib/datalog/DataLogReaderTest.java b/datalog/src/test/java/org/wpilib/datalog/DataLogReaderTest.java new file mode 100644 index 00000000000..e902fc44650 --- /dev/null +++ b/datalog/src/test/java/org/wpilib/datalog/DataLogReaderTest.java @@ -0,0 +1,33 @@ +// Copyright (c) FIRST and other WPILib contributors. +// Open Source Software; you can modify and/or share it under the terms of +// the WPILib BSD license file in the root directory of this project. + +package org.wpilib.datalog; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.NoSuchElementException; +import org.junit.jupiter.api.Test; + +class DataLogReaderTest { + @Test + void rejectsUnsignedRecordSizeLargerThanInput() { + ByteBuffer input = ByteBuffer.allocate(28).order(ByteOrder.LITTLE_ENDIAN); + input.put(new byte[] {'W', 'P', 'I', 'L', 'O', 'G'}); + input.putShort((short) 0x0100); + input.putInt(0); + + input.put((byte) 0x0c); // entryLen=1, sizeLen=4, timestampLen=1 + input.put((byte) 1); + input.putInt(0xfffffff9); + input.put((byte) 0); + input.rewind(); + + DataLogIterator iterator = new DataLogReader(input).iterator(); + assertFalse(iterator.hasNext()); + assertThrows(NoSuchElementException.class, iterator::next); + } +}