Merge pull request #28 from mandiant/optimize/willi-4

pe: cache DataAfterSection for performance

objfile/pe.go

@@ -20,15 +20,16 @@ import (
 )
 
 type peFile struct {
-	pe *pe.File
+	pe                    *pe.File
+	dataAfterSectionCache map[*pe.Section][]byte
 }
 
 func openPE(r io.ReaderAt) (rawFile, error) {
 	f, err := pe.NewFile(r)
 	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