fix: prevent consumer crash when reading incomplete record batches

src/apis/consumer/fetch-v15.ts

@@ -186,7 +186,12 @@ export function parseResponse (
             const recordsBatchesReader = Reader.from(r.buffer.subarray(r.position, r.position + recordsSize))
             partition.records = []
             do {
-              partition.records.push(readRecordsBatch(recordsBatchesReader))
+              const batch = readRecordsBatch(recordsBatchesReader)
+              if (batch) {
+                partition.records.push(batch)
+              } else {
+                break
+              }
             } while (recordsBatchesReader.position < recordsSize)
 
             r.skip(recordsSize)

src/apis/consumer/fetch-v16.ts

@@ -186,7 +186,12 @@ export function parseResponse (
             const recordsBatchesReader = Reader.from(r.buffer.subarray(r.position, r.position + recordsSize))
             partition.records = []
             do {
-              partition.records.push(readRecordsBatch(recordsBatchesReader))
+              const batch = readRecordsBatch(recordsBatchesReader)
+              if (batch) {
+                partition.records.push(batch)
+              } else {
+                break
+              }
             } while (recordsBatchesReader.position < recordsSize)
 
             r.skip(recordsSize)

src/apis/consumer/fetch-v17.ts

@@ -186,7 +186,12 @@ export function parseResponse (
             const recordsBatchesReader = Reader.from(r.buffer.subarray(r.position, r.position + recordsSize))
             partition.records = []
             do {
-              partition.records.push(readRecordsBatch(recordsBatchesReader))
+              const batch = readRecordsBatch(recordsBatchesReader)
+              if (batch) {
+                partition.records.push(batch)
+              } else {
+                break
+              }
             } while (recordsBatchesReader.position < recordsSize)
 
             r.skip(recordsSize)

src/protocol/reader.ts

@@ -409,4 +409,8 @@ export class Reader {
       this.skip(length)
     }
   }
+
+  canRead (length: number): boolean {
+    return this.buffer.length - this.position >= length
+  }
 }

src/protocol/records.ts

@@ -12,6 +12,20 @@ import { DynamicBuffer } from './dynamic-buffer.ts'
 import { Reader } from './reader.ts'
 import { Writer } from './writer.ts'
 
+const BATCH_HEADER_SIZE =
+  /* FirstOffset */ 8 +
+  /* Length */ 4 +
+  /* PartitionLeaderEpoch */ 4 +
+  /* Magic */ 1 +
+  /* CRC */ 4 +
+  /* Attributes */ 2 +
+  /* LastOffsetDelta */ 4 +
+  /* FirstTimestamp */ 8 +
+  /* MaxTimestamp */ 8 +
+  /* ProducerId */ 8 +
+  /* ProducerEpoch */ 2 +
+  /* FirstSequence */ 4 +
+  /* RecordsLength */ 4
 const CURRENT_RECORD_VERSION = 2
 const IS_TRANSACTIONAL = 0b10000 // Bit 4 set
 const IS_COMPRESSED = 0b111 // Bits 0, 1 and/or 2 set
@@ -247,10 +261,23 @@ export function createRecordsBatch (
   )
 }
 
-export function readRecordsBatch (reader: Reader): RecordsBatch {
+export function readRecordsBatch (reader: Reader): RecordsBatch | undefined {
+  // Check if we have enough data to read the batch header
+  if (!reader.canRead(BATCH_HEADER_SIZE)) {
+    return
+  }
+
+  const firstOffset = reader.readInt64()
+  const length = reader.readInt32()
+
+  // Check if we have enough data to read the complete batch, including the rest of the header
+  if (!reader.canRead(length)) {
+    return
+  }
+
   const batch = {
-    firstOffset: reader.readInt64(),
-    length: reader.readInt32(),
+    firstOffset,
+    length,
     partitionLeaderEpoch: reader.readInt32(),
     magic: reader.readInt8(),
     crc: reader.readUnsignedInt32(),

test/protocol/records.test.ts

@@ -229,8 +229,8 @@
   const reader = Reader.from(batchBuffer)
   const batch = readRecordsBatch(reader)
 
   strictEqual(batch.records.length, 1, 'Should have one record')
   strictEqual(batch.records[0].value.toString(), 'value', 'Record value should match')
 })
 
 test('createRecordsBatch with compression should compress the records', () => {
@@ -257,10 +257,10 @@
   const reader = Reader.from(batchBuffer)
   const batch = readRecordsBatch(reader)
 
   strictEqual(batch.attributes & 0b111, 1, 'Attributes should have gzip compression bit set')
   strictEqual(batch.records.length, 2, 'Should have two records after decompression')
   strictEqual(batch.records[0].value.toString(), 'value1'.repeat(100), 'First record value should match')
   strictEqual(batch.records[1].value.toString(), 'value2'.repeat(100), 'Second record value should match')
 })
 
 test('createRecordsBatch should throw on unsupported compression', () => {
@@ -379,17 +379,17 @@
   const reader = Reader.from(batchBuffer)
   const batch = readRecordsBatch(reader)
 
   strictEqual(batch.attributes & 0b111, 1, 'Attributes should have gzip compression bit set')
   strictEqual(batch.records.length, 2, 'Should have two records')
   strictEqual(batch.records[0].value.toString(), 'compressed1', 'First record value should match')
   strictEqual(batch.records[1].value.toString(), 'compressed2', 'Second record value should match')
 })
 
 test('readRecordsBatch should throw on unsupported compression bitmask', () => {
   // Create a mock batch with an invalid compression bitmask
   const writer = Writer.create()
     .appendInt64(0n) // firstOffset
-    .appendInt32(100) // length
+    .appendInt32(49) // length
     .appendInt32(0) // partitionLeaderEpoch
     .appendInt8(2) // magic
     .appendUnsignedInt32(0) // crc