fix: handle off-by-one error in DIFAT read and invalid chains (#11)

Author: secDre4mer

limits.go

@@ -4,5 +4,6 @@ const (
 	// Limit the size of the document
 	MAX_SECTORS = 1024 * 1024
 
-	MAX_SECTOR_SHIFT = 10
+	sectorShiftV3 = 0x9
+	sectorShiftV4 = 0xC
 )

oleparse.go

@@ -37,14 +37,14 @@ type OLEHeader struct {
 	Clid  [16]byte
 
 	MinorVersion    uint16
-	DllVersion      uint16
+	MajorVersion    uint16
 	ByteOrder       uint16
 	SectorShift     uint16
 	MiniSectorShift uint16
 	Reserved        uint16
 
 	Reserved1        uint32
-	Reserved2        uint32
+	CsectDir         uint32 // Count of directory sectors. Only available in version 4.
 	CsectFat         uint32
 	SectDirStart     uint32
 	Signature        uint32
@@ -89,6 +89,9 @@ func NewDirectory(data []byte, index uint32) (*Directory, error) {
 	if err != nil {
 		return nil, err
 	}
+	if self.Header.Mse == 0 { // Unallocated
+		return nil, nil
+	}
 
 	self.Name = strings.TrimRight(
 		string(utf16.Decode(self.Header.AB[:])), "\x00")
@@ -117,7 +120,7 @@ type VBAModule struct {
 }
 
 func (self *OLEFile) ReadSector(sector uint32) []byte {
-	start := 512 + self.SectorSize*int(sector)
+	start := self.SectorSize * int(sector+1)
 
 	to_read := self.SectorSize
 	if start > len(self.data) || start < 0 {
@@ -145,18 +148,18 @@ func (self *OLEFile) ReadMiniSector(sector uint32) []byte {
 	return self.ministream[start : start+to_read]
 }
 
-func (self *OLEFile) ReadFat(sector uint32) uint32 {
+func (self *OLEFile) ReadFat(sector uint32) (uint32, bool) {
 	if int(sector) >= len(self.Fat) {
-		return 0
+		return 0, false
 	}
-	return self.Fat[sector]
+	return self.Fat[sector], true
 }
 
-func (self *OLEFile) ReadMiniFat(sector uint32) uint32 {
+func (self *OLEFile) ReadMiniFat(sector uint32) (uint32, bool) {
 	if int(sector) >= len(self.MiniFat) {
-		return 0
+		return 0, false
 	}
-	return self.MiniFat[sector]
+	return self.MiniFat[sector], true
 }
 
 func (self *OLEFile) ReadChain(start uint32) []byte {
@@ -170,13 +173,18 @@ func (self *OLEFile) ReadMiniChain(start uint32) []byte {
 func (self *OLEFile) _ReadChain(
 	start uint32,
 	ReadSector func(uint32) []byte,
-	ReadFat func(sector uint32) uint32) []byte {
+	ReadFat func(sector uint32) (uint32, bool),
+) []byte {
 	check := make(map[uint32]bool)
 	result := []byte{}
 
 	for sector := start; sector != ENDOFCHAIN; {
 		result = append(result, ReadSector(sector)...)
-		next := ReadFat(sector)
+		next, ok := ReadFat(sector)
+		if !ok {
+			DebugPrintf("invalid sector %x in chain", sector)
+			return result
+		}
 		_, pres := check[next]
 		if pres {
 			DebugPrintf("infinite loop detected at %v to %v starting at %v",
@@ -225,6 +233,9 @@ func (self *OLEFile) OpenStreamByName(name string) ([]byte, error) {
 	return self.GetStream(d.Index), nil
 }
 
+// NewOLEFile creates a new OLEFile object from the given data.
+//
+// The OLE format is described in https://winprotocoldoc.z19.web.core.windows.net/MS-CFB/%5bMS-CFB%5d.pdf
 func NewOLEFile(data []byte) (*OLEFile, error) {
 	if len(data) < 8 ||
 		string(data[:8]) != OLE_SIGNATURE {
@@ -238,9 +249,21 @@ func NewOLEFile(data []byte) (*OLEFile, error) {
 		return nil, err
 	}
 
-	if self.Header.SectorShift > MAX_SECTOR_SHIFT {
-		return nil, fmt.Errorf(
-			"Sector size too large: %v", self.Header.SectorShift)
+	var expectedSectorShift uint16
+	switch self.Header.MajorVersion {
+	case 3:
+		expectedSectorShift = sectorShiftV3
+	case 4:
+		expectedSectorShift = sectorShiftV4
+	default:
+		return nil, fmt.Errorf("unsupported major version: %v", self.Header.MajorVersion)
+	}
+	if self.Header.MinorVersion != 0x3E {
+		return nil, fmt.Errorf("unsupported minor version: %v", self.Header.MinorVersion)
+	}
+
+	if self.Header.SectorShift != expectedSectorShift {
+		return nil, fmt.Errorf("unexpected sector size: %d", 1<<self.Header.SectorShift)
 	}
 
 	self.SectorSize = 1 << self.Header.SectorShift
@@ -250,11 +273,11 @@ func NewOLEFile(data []byte) (*OLEFile, error) {
 	}
 
 	self.MiniSectorSize = 1 << self.Header.MiniSectorShift
-	if (len(data)-512)%self.SectorSize != 0 {
+	if len(data)%self.SectorSize != 0 {
 		DebugPrintf("Last sector has invalid size\n")
 	}
 
-	self.SectorCount = (len(data) - 512) / self.SectorSize
+	self.SectorCount = len(data)/self.SectorSize - 1 // Subtract 1 for the header sector
 	for _, sect := range self.Header.SectFat {
 		if sect != FREESECT {
 			self.FatSectors = append(self.FatSectors, sect)
@@ -279,7 +302,7 @@ func NewOLEFile(data []byte) (*OLEFile, error) {
 		}
 
 		next := dif_values[len(dif_values)-1]
-		for _, value := range dif_values[:len(dif_values)-2] {
+		for _, value := range dif_values[:len(dif_values)-1] {
 			if value != FREESECT {
 				self.FatSectors = append(self.FatSectors, value)
 			}
@@ -319,6 +342,9 @@ func NewOLEFile(data []byte) (*OLEFile, error) {
 		if err != nil {
 			return nil, err
 		}
+		if dir_obj == nil { // Unallocated index
+			continue
+		}
 		self.Directory = append(self.Directory, dir_obj)
 	}
 

oleparse_test.go

@@ -4,6 +4,7 @@ import (
 	"archive/zip"
 	"encoding/json"
 	"io"
+	"os"
 	"strings"
 	"testing"
 
@@ -20,6 +21,20 @@ func TestMacros(t *testing.T) {
 	goldie.Assert(t, "vba_macros", serialized)
 }
 
+func TestOlev4(t *testing.T) {
+	oleData, err := os.ReadFile("test_data/vs.msi")
+	if err != nil {
+		t.Fatalf("Failed to open test file: %v", err)
+	}
+	oleFile, err := NewOLEFile(oleData)
+	if err != nil {
+		t.Fatalf("Failed to parse OLE file: %v", err)
+	}
+	if len(oleFile.Directory) != 28 {
+		t.Fatalf("Expected 28 directory entries, got %d", len(oleFile.Directory))
+	}
+}
+
 func FuzzExtractMacros(f *testing.F) {
 	r, err := zip.OpenReader("test_data/xlswithmacro.xlsm")
 	if err != nil {