Improved bootloader parsing (#356) (#358)

* Improved bootloader parsing (#356)

* Add unit tests for inheritBootloaderKindBySHA function

---------

Co-authored-by: Rodage, Alpesh Ramesh <alpesh.ramesh.rodage@intel.com>

Author: magerstam

internal/image/imageinspect/bootloader_pe.go

@@ -8,17 +8,15 @@ import (
 	"strings"
 )
 
-// ParsePEFromBytes parses a PE (Portable Executable) binary from the given byte slice
 func ParsePEFromBytes(p string, blob []byte) (EFIBinaryEvidence, error) {
 	ev := EFIBinaryEvidence{
 		Path:            p,
 		Size:            int64(len(blob)),
 		SectionSHA256:   map[string]string{},
 		OSReleaseSorted: []KeyValue{},
-		Kind:            classifyBootloaderKind(p, nil), // refine after we parse sections
+		Kind:            BootloaderUnknown, // set after we have more evidence
 	}
 
-	// whole-file hash
 	ev.SHA256 = sha256Hex(blob)
 
 	r := bytes.NewReader(blob)
@@ -30,31 +28,26 @@ func ParsePEFromBytes(p string, blob []byte) (EFIBinaryEvidence, error) {
 
 	ev.Arch = peMachineToArch(f.FileHeader.Machine)
 
-	// Sections
 	for _, s := range f.Sections {
 		name := strings.TrimRight(s.Name, "\x00")
 		ev.Sections = append(ev.Sections, name)
 	}
 
-	// Signed evidence: presence of Authenticode blob
 	signed, sigSize, sigNote := peSignatureInfo(f)
 	ev.Signed = signed
 	ev.SignatureSize = sigSize
 	if sigNote != "" {
 		ev.Notes = append(ev.Notes, sigNote)
 	}
 
-	// SBAT section presence
 	ev.HasSBAT = hasSection(ev.Sections, ".sbat")
 
-	// UKI detection: these sections are highly indicative
 	isUKI := hasSection(ev.Sections, ".linux") &&
 		(hasSection(ev.Sections, ".cmdline") || hasSection(ev.Sections, ".osrel") || hasSection(ev.Sections, ".uname"))
 	ev.IsUKI = isUKI
 	if isUKI {
 		ev.Kind = BootloaderUKI
 	} else {
-		// reclassify based on name/path/sections
 		ev.Kind = classifyBootloaderKind(p, ev.Sections)
 	}
 
@@ -121,18 +114,21 @@ func peSignatureInfo(f *pe.File) (signed bool, sigSize int, note string) {
 	return false, 0, ""
 }
 
-// classifyBootloaderKind classifies the bootloader kind based on path and sections
+// classifyBootloaderKind classifies the bootloader kind based on path and sections.
+// It intentionally avoids content-string heuristics for stability.
+// For BOOTX64.EFI copies/aliases, rely on SHA-inheritance post-pass.
 func classifyBootloaderKind(p string, sections []string) BootloaderKind {
 	lp := strings.ToLower(p)
 
-	// Most deterministic first:
+	// Deterministic first:
 	if sections != nil && hasSection(sections, ".linux") {
-		// likely UKI; caller can override with stricter check
 		return BootloaderUKI
 	}
 
 	// Path / filename heuristics:
-	// shim often includes "shim" and/or has .sbat too
+	if strings.Contains(lp, "mmx64.efi") || strings.Contains(lp, "mmia32.efi") {
+		return BootloaderMokManager
+	}
 	if strings.Contains(lp, "shim") {
 		return BootloaderShim
 	}
@@ -142,10 +138,7 @@ func classifyBootloaderKind(p string, sections []string) BootloaderKind {
 	if strings.Contains(lp, "grub") {
 		return BootloaderGrub
 	}
-	if strings.Contains(lp, "mmx64.efi") || strings.Contains(lp, "mmia32.efi") {
-		return BootloaderMokManager
-	}
-	// fallback
+
 	return BootloaderUnknown
 }
 
@@ -160,6 +153,34 @@ func hasSection(secs []string, want string) bool {
 	return false
 }
 
+// inheritBootloaderKindBySHA assigns a kind to "unknown" EFI binaries when they
+// are byte-identical to another EFI binary already classified as a known kind.
+// This reliably handles fallback paths like EFI/BOOT/BOOTX64.EFI.
+func inheritBootloaderKindBySHA(evs []EFIBinaryEvidence) {
+	known := make(map[string]BootloaderKind) // sha256 -> kind
+
+	// First pass: record known kinds by hash.
+	for _, ev := range evs {
+		if ev.SHA256 == "" || ev.Kind == BootloaderUnknown {
+			continue
+		}
+		if _, ok := known[ev.SHA256]; !ok {
+			known[ev.SHA256] = ev.Kind
+		}
+	}
+
+	// Second pass: upgrade unknowns when a known hash exists.
+	for i := range evs {
+		if evs[i].Kind != BootloaderUnknown || evs[i].SHA256 == "" {
+			continue
+		}
+		if k, ok := known[evs[i].SHA256]; ok {
+			evs[i].Kind = k
+			evs[i].Notes = append(evs[i].Notes, "bootloader kind inherited from identical EFI binary (sha256 match)")
+		}
+	}
+}
+
 // peMachineToArch maps PE machine types to architecture strings
 func peMachineToArch(m uint16) string {
 	switch m {

internal/image/imageinspect/fs_inspect.go

@@ -181,7 +181,6 @@ func readExtSuperblock(r io.ReaderAt, partOff int64, out *FilesystemSummary) err
 	return nil
 }
 
-// readFATBootSector reads the FAT boot sector and fills in details.
 func readFATBootSector(r io.ReaderAt, partOff int64, out *FilesystemSummary) error {
 	bs := make([]byte, 512)
 	if _, err := r.ReadAt(bs, partOff); err != nil && err != io.EOF {
@@ -200,73 +199,87 @@ func readFATBootSector(r io.ReaderAt, partOff int64, out *FilesystemSummary) err
 	totSec16 := binary.LittleEndian.Uint16(bs[19:21])
 	fatSz16 := binary.LittleEndian.Uint16(bs[22:24])
 	totSec32 := binary.LittleEndian.Uint32(bs[32:36])
+	fatSz32 := binary.LittleEndian.Uint32(bs[36:40]) // BPB_FATSz32 (FAT32)
 
 	out.Type = "vfat"
 	out.BytesPerSector = bytesPerSec
 	out.SectorsPerCluster = secPerClus
 
+	// Basic sanity checks to avoid bogus classification
+	switch bytesPerSec {
+	case 512, 1024, 2048, 4096:
+		// ok
+	default:
+		return fmt.Errorf("invalid BPB: bytesPerSec=%d", bytesPerSec)
+	}
+	if secPerClus == 0 {
+		return fmt.Errorf("invalid BPB: sectorsPerCluster=0")
+	}
+	if numFATs == 0 {
+		return fmt.Errorf("invalid BPB: numFATs=0")
+	}
+
 	// Total sectors is either TotSec16 or TotSec32
 	totalSectors := uint32(totSec16)
 	if totalSectors == 0 {
 		totalSectors = totSec32
 	}
+	if totalSectors == 0 {
+		return fmt.Errorf("invalid BPB: totalSectors=0")
+	}
 
-	// FAT32 detection (canonical)
-	fatSz32 := binary.LittleEndian.Uint32(bs[36:40]) // only meaningful if FAT32
+	// Canonical FAT32 detection:
+	// - RootEntCnt must be 0 for FAT32
+	// - FATSz16 must be 0 for FAT32
+	// - FATSz32 should be non-zero for FAT32 (but we won't *require* it to avoid false negatives on odd images)
 	isFAT32 := (rootEntCnt == 0) && (fatSz16 == 0) && (fatSz32 != 0)
 
 	if isFAT32 {
 		out.FATType = "FAT32"
-
 		out.UUID = fmt.Sprintf("%08x", binary.LittleEndian.Uint32(bs[67:71]))
 		out.Label = strings.TrimRight(string(bs[71:82]), " \x00")
 
-		// cluster count for FAT32
+		// cluster count for FAT32 (root dir is in data area)
 		rootDirSectors := uint32(0)
 		fatSectors := fatSz32
 		dataSectors := totalSectors - (uint32(rsvdSecCnt) + (numFATs * fatSectors) + rootDirSectors)
-		if secPerClus == 0 {
-			return fmt.Errorf("invalid BPB: sectorsPerCluster=0")
-		}
 		out.ClusterCount = dataSectors / uint32(secPerClus)
-
 		return nil
 	}
 
 	// FAT12/16-style: classify via cluster count
-	// Root dir sectors:
 	rootDirSectors := ((uint32(rootEntCnt) * 32) + (uint32(bytesPerSec) - 1)) / uint32(bytesPerSec)
-
-	// FAT size sectors (FAT12/16 uses fatSz16)
 	fatSectors := uint32(fatSz16)
 
-	// Data sectors:
-	dataSectors := totalSectors - (uint32(rsvdSecCnt) + (numFATs * fatSectors) + rootDirSectors)
-
-	// Count of clusters:
-	if secPerClus == 0 {
-		return fmt.Errorf("invalid BPB: sectorsPerCluster=0")
+	// If FAT16 fields suggest nonsense, note it (helps debug wrong offsets/sector)
+	if fatSectors == 0 {
+		out.Notes = append(out.Notes, fmt.Sprintf("BPB_FATSz16=0 but not detected as FAT32 (RootEntCnt=%d FATSz32=%d)", rootEntCnt, fatSz32))
 	}
+
+	dataSectors := totalSectors - (uint32(rsvdSecCnt) + (numFATs * fatSectors) + rootDirSectors)
 	clusterCount := dataSectors / uint32(secPerClus)
 
-	// Standard FAT thresholds
+	out.ClusterCount = clusterCount
+
 	switch {
 	case clusterCount < 4085:
 		out.FATType = "FAT12"
-		out.ClusterCount = clusterCount
 	case clusterCount < 65525:
 		out.FATType = "FAT16"
-		out.ClusterCount = clusterCount
 	default:
-		// It’s possible to encounter FAT32-like cluster counts without the FAT32 BPB layout,
-		// but for an ESP this is unlikely. Still, classify as FAT32 if huge.
+		// If cluster count is huge, it's overwhelmingly likely FAT32.
 		out.FATType = "FAT32"
 	}
 
 	// FAT12/16 Extended BPB: VolID @ 39..43, Label @ 43..54
 	out.UUID = fmt.Sprintf("%08x", binary.LittleEndian.Uint32(bs[39:43]))
 	out.Label = strings.TrimRight(string(bs[43:54]), " \x00")
 
+	out.Notes = append(out.Notes, fmt.Sprintf(
+		"FAT BPB: BytsPerSec=%d SecPerClus=%d Rsvd=%d NumFATs=%d RootEntCnt=%d TotSec16=%d TotSec32=%d FATSz16=%d FATSz32=%d",
+		bytesPerSec, secPerClus, rsvdSecCnt, numFATs, rootEntCnt, totSec16, totSec32, fatSz16, fatSz32,
+	))
+
 	return nil
 }
 

internal/image/imageinspect/fs_raw.go

@@ -138,6 +138,7 @@ func scanAndHashEFIFromRawFAT(r io.ReaderAt, partOff int64, out *FilesystemSumma
 		}
 	}
 
+	inheritBootloaderKindBySHA(out.EFIBinaries)
 	sort.Slice(out.EFIBinaries, func(i, j int) bool { return out.EFIBinaries[i].Path < out.EFIBinaries[j].Path })
 
 	out.HasShim = out.HasShim || hasShim

internal/image/imageinspect/imageinspect_core_test.go

@@ -807,3 +807,108 @@ func TestInspectCore_PropagatesFilesystemError_WhenCalled(t *testing.T) {
 		t.Fatalf("expected GetFilesystem to be called at least once")
 	}
 }
+
+func TestInheritBootloaderKindBySHA_InheritsFromKnown(t *testing.T) {
+	evs := []EFIBinaryEvidence{
+		{Path: "/EFI/ubuntu/shimx64.efi", SHA256: "abc123def456", Kind: BootloaderShim},
+		{Path: "/EFI/BOOT/BOOTX64.EFI", SHA256: "abc123def456", Kind: BootloaderUnknown},
+	}
+
+	inheritBootloaderKindBySHA(evs)
+
+	if evs[1].Kind != BootloaderShim {
+		t.Errorf("expected BOOTX64.EFI to inherit kind=shim, got %q", evs[1].Kind)
+	}
+	if len(evs[1].Notes) == 0 || !strings.Contains(evs[1].Notes[0], "sha256 match") {
+		t.Errorf("expected note about sha256 inheritance, got %v", evs[1].Notes)
+	}
+	// Original should remain unchanged
+	if evs[0].Kind != BootloaderShim {
+		t.Errorf("original should stay shim, got %q", evs[0].Kind)
+	}
+}
+
+func TestInheritBootloaderKindBySHA_NoMatchLeavesUnknown(t *testing.T) {
+	evs := []EFIBinaryEvidence{
+		{Path: "/EFI/ubuntu/shimx64.efi", SHA256: "abc123", Kind: BootloaderShim},
+		{Path: "/EFI/BOOT/BOOTX64.EFI", SHA256: "different456", Kind: BootloaderUnknown},
+	}
+
+	inheritBootloaderKindBySHA(evs)
+
+	if evs[1].Kind != BootloaderUnknown {
+		t.Errorf("expected BOOTX64.EFI to remain unknown when hash differs, got %q", evs[1].Kind)
+	}
+	if len(evs[1].Notes) != 0 {
+		t.Errorf("expected no notes when no match, got %v", evs[1].Notes)
+	}
+}
+
+func TestInheritBootloaderKindBySHA_EmptySHA256Ignored(t *testing.T) {
+	evs := []EFIBinaryEvidence{
+		{Path: "/EFI/ubuntu/shimx64.efi", SHA256: "", Kind: BootloaderShim},
+		{Path: "/EFI/BOOT/BOOTX64.EFI", SHA256: "", Kind: BootloaderUnknown},
+	}
+
+	inheritBootloaderKindBySHA(evs)
+
+	// Both should remain unchanged - empty SHA256 entries are skipped
+	if evs[0].Kind != BootloaderShim {
+		t.Errorf("expected shim to remain shim, got %q", evs[0].Kind)
+	}
+	if evs[1].Kind != BootloaderUnknown {
+		t.Errorf("expected unknown to remain unknown when SHA256 empty, got %q", evs[1].Kind)
+	}
+}
+
+func TestInheritBootloaderKindBySHA_MultipleInheritances(t *testing.T) {
+	evs := []EFIBinaryEvidence{
+		{Path: "/EFI/ubuntu/shimx64.efi", SHA256: "shimhash", Kind: BootloaderShim},
+		{Path: "/EFI/fedora/grubx64.efi", SHA256: "grubhash", Kind: BootloaderGrub},
+		{Path: "/EFI/BOOT/BOOTX64.EFI", SHA256: "shimhash", Kind: BootloaderUnknown},
+		{Path: "/EFI/BOOT/grubx64.efi", SHA256: "grubhash", Kind: BootloaderUnknown},
+		{Path: "/EFI/unknown/mystery.efi", SHA256: "otherhash", Kind: BootloaderUnknown},
+	}
+
+	inheritBootloaderKindBySHA(evs)
+
+	if evs[2].Kind != BootloaderShim {
+		t.Errorf("BOOTX64.EFI should inherit shim, got %q", evs[2].Kind)
+	}
+	if evs[3].Kind != BootloaderGrub {
+		t.Errorf("grubx64.efi copy should inherit grub, got %q", evs[3].Kind)
+	}
+	if evs[4].Kind != BootloaderUnknown {
+		t.Errorf("mystery.efi should remain unknown, got %q", evs[4].Kind)
+	}
+}
+
+func TestInheritBootloaderKindBySHA_FirstKnownWins(t *testing.T) {
+	// If the same hash appears with different kinds, first one wins
+	evs := []EFIBinaryEvidence{
+		{Path: "/EFI/first/shimx64.efi", SHA256: "samehash", Kind: BootloaderShim},
+		{Path: "/EFI/second/grubx64.efi", SHA256: "samehash", Kind: BootloaderGrub},
+		{Path: "/EFI/BOOT/BOOTX64.EFI", SHA256: "samehash", Kind: BootloaderUnknown},
+	}
+
+	inheritBootloaderKindBySHA(evs)
+
+	// The unknown should inherit from the first known (shim)
+	if evs[2].Kind != BootloaderShim {
+		t.Errorf("expected first known kind (shim) to win, got %q", evs[2].Kind)
+	}
+}
+
+func TestInheritBootloaderKindBySHA_AlreadyClassifiedNotOverwritten(t *testing.T) {
+	evs := []EFIBinaryEvidence{
+		{Path: "/EFI/ubuntu/shimx64.efi", SHA256: "abc123", Kind: BootloaderShim},
+		{Path: "/EFI/BOOT/BOOTX64.EFI", SHA256: "abc123", Kind: BootloaderGrub}, // Already classified differently
+	}
+
+	inheritBootloaderKindBySHA(evs)
+
+	// Should NOT overwrite an already-classified binary
+	if evs[1].Kind != BootloaderGrub {
+		t.Errorf("already classified binary should not be overwritten, got %q", evs[1].Kind)
+	}
+}