[datalog] Reject oversized Java data log records

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

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;
 }

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);
+  }
+}