From 36154a62d1b538244e28b7fb16237b284094b46c Mon Sep 17 00:00:00 2001
From: Peter Haag <peter@people.ops-trust.net>
Date: Tue, 16 Apr 2024 16:58:53 +0200
Subject: [PATCH] Add boundary check to fix #11

---
 nffile.go |  6 +++++-
 record.go | 14 +++++++++++---
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/nffile.go b/nffile.go
index 40ae5f3..e9f6a2c 100644
--- a/nffile.go
+++ b/nffile.go
@@ -267,6 +267,7 @@ func (nfFile *NfFile) AllRecords() (chan *FlowRecordV3, error) {
 	recordChannel := make(chan *FlowRecordV3, 32)
 	go func() {
 		blockChannel, _ := nfFile.ReadDataBlocks()
+		recordCnt := 0
 		for dataBlock := range blockChannel {
 			// fmt.Printf("Next block - type: %d, records: %d\n", dataBlock.Header.Type, dataBlock.Header.NumRecords)
 			offset := 0
@@ -277,9 +278,12 @@ func (nfFile *NfFile) AllRecords() (chan *FlowRecordV3, error) {
 				// fmt.Printf("Record %d type: %d, length: %d, numElementS: %d\n", i, recordType, recordSize, numElementS)
 				switch recordType {
 				case V3Record:
-					if record := NewRecord(dataBlock.Data[offset : offset+int(recordSize)]); record != nil {
+					if record, err := NewRecord(dataBlock.Data[offset : offset+int(recordSize)]); record != nil {
 						record.GetSamplerInfo(nfFile)
 						recordChannel <- record
+						recordCnt++
+					} else {
+						fmt.Fprintf(os.Stderr, "Record %d: decoding error: %v\n", recordCnt, err)
 					}
 				case ExporterInfoRecordType:
 					nfFile.addExporterInfo(dataBlock.Data[offset : offset+int(recordSize)])
diff --git a/record.go b/record.go
index 2db98ee..865798c 100644
--- a/record.go
+++ b/record.go
@@ -10,6 +10,7 @@ package nfdump
 
 import (
 	"encoding/binary"
+	"fmt"
 	"net"
 	"unsafe"
 )
@@ -45,7 +46,7 @@ type FlowRecordV3 struct {
 }
 
 // Extract next flow record from []byte stream
-func NewRecord(record []byte) *FlowRecordV3 {
+func NewRecord(record []byte) (*FlowRecordV3, error) {
 
 	offset := 0
 	recordType := binary.LittleEndian.Uint16(record[offset : offset+2])
@@ -53,7 +54,7 @@ func NewRecord(record []byte) *FlowRecordV3 {
 	numElements := binary.LittleEndian.Uint16(record[offset+4 : offset+6])
 
 	if recordType != V3Record {
-		return nil
+		return nil, fmt.Errorf("Not a v3 record")
 	}
 
 	flowRecord := new(FlowRecordV3)
@@ -64,8 +65,15 @@ func NewRecord(record []byte) *FlowRecordV3 {
 	flowRecord.recordHeader = (*recordHeaderV3)(unsafe.Pointer(&raw[0]))
 	offset = 12
 	for i := 0; i < int(numElements); i++ {
+		// fmt.Printf(" . next Element at offset: %d\n", offset)
+		if (offset + 4) > int(recordSize) {
+			return nil, fmt.Errorf("Record header boundary check error")
+		}
 		elementType := binary.LittleEndian.Uint16(raw[offset : offset+2])
 		elementSize := binary.LittleEndian.Uint16(raw[offset+2 : offset+4])
+		if (offset + int(elementSize)) > int(recordSize) {
+			return nil, fmt.Errorf("Record body boundary check error")
+		}
 		// fmt.Printf(" . Element type: %d, length: %d\n", elementType, elementSize)
 		exOffset := offset + 4
 		if elementType < MAXEXTENSIONS {
@@ -96,7 +104,7 @@ func NewRecord(record []byte) *FlowRecordV3 {
 	flowRecord.packetInterval = 1
 	flowRecord.spaceInterval = 0
 
-	return flowRecord
+	return flowRecord, nil
 }
 
 // Return generic extension