From edf7c328835273e2bc6dd23c646091e6a03aa2e9 Mon Sep 17 00:00:00 2001
From: Noel Georgi <git@frezbo.dev>
Date: Mon, 27 Jan 2025 18:22:56 +0530
Subject: [PATCH] fix: pe uki extract

The PE UKI extract was returning a reader that reads upto the section
`Size` which is the size on disk which could be padded, we need a limit
reader reading upto `VirtualSize`.

Signed-off-by: Noel Georgi <git@frezbo.dev>
---
 internal/pkg/uki/internal/pe/extract.go      |  9 ++-
 internal/pkg/uki/internal/pe/extract_test.go | 74 ++++++++++++++++++++
 2 files changed, 80 insertions(+), 3 deletions(-)
 create mode 100644 internal/pkg/uki/internal/pe/extract_test.go

diff --git a/internal/pkg/uki/internal/pe/extract.go b/internal/pkg/uki/internal/pe/extract.go
index ee85cdb16..570df9bf9 100644
--- a/internal/pkg/uki/internal/pe/extract.go
+++ b/internal/pkg/uki/internal/pe/extract.go
@@ -33,13 +33,16 @@ func Extract(ukiPath string) (assetInfo AssetInfo, err error) {
 	assetInfo.fileCloser = peFile
 
 	for _, section := range peFile.Sections {
+		// read upto section.VirtualSize bytes
+		sectionReader := io.NewSectionReader(section, 0, int64(section.VirtualSize))
+
 		switch section.Name {
 		case ".initrd":
-			assetInfo.Initrd = section.Open()
+			assetInfo.Initrd = sectionReader
 		case ".cmdline":
-			assetInfo.Cmdline = section.Open()
+			assetInfo.Cmdline = sectionReader
 		case ".linux":
-			assetInfo.Kernel = section.Open()
+			assetInfo.Kernel = sectionReader
 		}
 	}
 
diff --git a/internal/pkg/uki/internal/pe/extract_test.go b/internal/pkg/uki/internal/pe/extract_test.go
new file mode 100644
index 000000000..b5baa5832
--- /dev/null
+++ b/internal/pkg/uki/internal/pe/extract_test.go
@@ -0,0 +1,74 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package pe_test
+
+import (
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/siderolabs/talos/internal/pkg/uki/internal/pe"
+)
+
+func TestUKIExtract(t *testing.T) {
+	srcFile := "testdata/sd-stub-amd64.efi"
+
+	destDir := t.TempDir()
+
+	destFile := filepath.Join(destDir, "vmlinuz.efi")
+
+	for _, section := range []string{"linux", "initrd", "cmdline"} {
+		assert.NoError(t, os.WriteFile(filepath.Join(destDir, section), []byte(section), 0o644))
+	}
+
+	assert.NoError(t, pe.AssembleNative(srcFile, destFile, []pe.Section{
+		{
+			Name:    ".linux",
+			Path:    filepath.Join(destDir, "linux"),
+			Measure: false,
+			Append:  true,
+		},
+		{
+			Name:    ".initrd",
+			Path:    filepath.Join(destDir, "initrd"),
+			Measure: false,
+			Append:  true,
+		},
+		{
+			Name:    ".cmdline",
+			Path:    filepath.Join(destDir, "cmdline"),
+			Measure: false,
+			Append:  true,
+		},
+	}))
+
+	ukiData, err := pe.Extract(destFile)
+	assert.NoError(t, err)
+
+	t.Cleanup(func() {
+		assert.NoError(t, ukiData.Close())
+	})
+
+	var kernel, initrd, cmdline strings.Builder
+
+	_, err = io.Copy(&kernel, ukiData.Kernel)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "linux", kernel.String())
+
+	_, err = io.Copy(&initrd, ukiData.Initrd)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "initrd", initrd.String())
+
+	_, err = io.Copy(&cmdline, ukiData.Cmdline)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "cmdline", cmdline.String())
+}