From c3d72edb67609c3ebea13a36fbd80b8050de9684 Mon Sep 17 00:00:00 2001
From: Willi Ballenthin <willi.ballenthin@gmail.com>
Date: Fri, 21 Jul 2023 14:53:31 +0200
Subject: [PATCH] pe: cache DataAfterSection for performance

---
 objfile/pe.go | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/objfile/pe.go b/objfile/pe.go
index 7f24fc8..59a3659 100644
--- a/objfile/pe.go
+++ b/objfile/pe.go
@@ -20,7 +20,8 @@ import (
 )
 
 type peFile struct {
-	pe *pe.File
+	pe                    *pe.File
+	dataAfterSectionCache map[*pe.Section][]byte
 }
 
 func openPE(r io.ReaderAt) (rawFile, error) {
@@ -28,7 +29,7 @@ func openPE(r io.ReaderAt) (rawFile, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &peFile{f}, nil
+	return &peFile{f, make(map[*pe.Section][]byte)}, nil
 }
 
 func (f *peFile) read_memory(VA uint64, size uint64) (data []byte, err error) {
@@ -127,6 +128,15 @@ func (f *peFile) symbols() ([]Sym, error) {
 	return syms, nil
 }
 
+func (f *peFile) dataAfterSection(sec *pe.Section) []byte {
+	if data, ok := f.dataAfterSectionCache[sec]; ok {
+		return data
+	}
+	data, _ := sec.Data()
+	f.dataAfterSectionCache[sec] = data
+	return data
+}
+
 func (f *peFile) pcln_scan() (candidates []PclntabCandidate, err error) {
 	var imageBase uint64
 	switch oh := f.pe.OptionalHeader.(type) {
@@ -176,7 +186,7 @@ func (f *peFile) pcln_scan() (candidates []PclntabCandidate, err error) {
 	// 2) if not found, byte scan for it
 	for _, sec := range f.pe.Sections {
 		// malware can split the pclntab across multiple sections, re-merge
-		data := f.pe.DataAfterSection(sec)
+		data := f.dataAfterSection(sec)
 
 		if !foundpcln {
 			matches := findAllOccurrences(data, pclntab_sigs)
@@ -288,7 +298,7 @@ func (f *peFile) pcln_scan() (candidates []PclntabCandidate, err error) {
 	if len(stompedmagic_candidates) != 0 {
 		for _, sec := range f.pe.Sections {
 			// malware can split the pclntab across multiple sections, re-merge
-			data := f.pe.DataAfterSection(sec)
+			data := f.dataAfterSection(sec)
 			for _, stompedMagicCandidate := range stompedmagic_candidates {
 				pclntab_va_candidate := stompedMagicCandidate.PclntabVa
 
@@ -391,7 +401,7 @@ func (f *peFile) moduledata_scan(pclntabVA uint64, is64bit bool, littleendian bo
 scan:
 	for _, sec := range f.pe.Sections {
 		// malware can split the pclntab across multiple sections, re-merge
-		data := f.pe.DataAfterSection(sec)
+		data := f.dataAfterSection(sec)
 
 		if !foundmodule {
 			// fall back to scanning for structure using address of pclntab, which is first value in struc