btf: support raw BTF in LoadSpecFromReader

Allow passing BTF not wrapped in an ELF to LoadSpecFromReader. This
allows users to use external BTF for vmlinux, which is usually shipped
without the ELF.

internal/btf/btf.go

@@ -78,12 +78,18 @@ func (h *btfHeader) stringStart() int64 {
 	return int64(h.HdrLen + h.StringOff)
 }
 
-// LoadSpecFromReader reads BTF sections from an ELF.
+// LoadSpecFromReader reads from an ELF or a raw BTF blob.
 //
-// Returns ErrNotFound if the reader contains no BTF.
+// Returns ErrNotFound if reading from an ELF which contains no BTF.
 func LoadSpecFromReader(rd io.ReaderAt) (*Spec, error) {
 	file, err := internal.NewSafeELFFile(rd)
 	if err != nil {
+		if bo := guessRawBTFByteOrder(rd); bo != nil {
+			// Try to parse a naked BTF blob. This will return an error if
+			// we encounter a Datasec, since we can't fix it up.
+			return loadRawSpec(io.NewSectionReader(rd, 0, math.MaxInt64), bo, nil, nil)
+		}
+
 		return nil, err
 	}
 	defer file.Close()
@@ -162,16 +168,6 @@ func loadSpecFromELF(file *internal.SafeELFFile, variableOffsets map[variable]ui
 	return spec, nil
 }
 
-// LoadRawSpec reads a blob of BTF data that isn't wrapped in an ELF file.
-//
-// Prefer using LoadSpecFromReader, since this function only supports a subset
-// of BTF.
-func LoadRawSpec(btf io.Reader, bo binary.ByteOrder) (*Spec, error) {
-	// This will return an error if we encounter a Datasec, since we can't fix
-	// it up.
-	return loadRawSpec(btf, bo, nil, nil)
-}
-
 func loadRawSpec(btf io.Reader, bo binary.ByteOrder, sectionSizes map[string]uint32, variableOffsets map[variable]uint32) (*Spec, error) {
 	rawTypes, rawStrings, err := parseBTF(btf, bo)
 	if err != nil {
@@ -224,7 +220,7 @@ func loadKernelSpec() (*Spec, error) {
 	if err == nil {
 		defer fh.Close()
 
-		return LoadRawSpec(fh, internal.NativeEndian)
+		return loadRawSpec(fh, internal.NativeEndian, nil, nil)
 	}
 
 	var uname unix.Utsname
@@ -299,6 +295,19 @@ func parseBTFHeader(r io.Reader, bo binary.ByteOrder) (*btfHeader, error) {
 	return &header, nil
 }
 
+func guessRawBTFByteOrder(r io.ReaderAt) binary.ByteOrder {
+	for _, bo := range []binary.ByteOrder{
+		binary.LittleEndian,
+		binary.BigEndian,
+	} {
+		if _, err := parseBTFHeader(io.NewSectionReader(r, 0, math.MaxInt64), bo); err == nil {
+			return bo
+		}
+	}
+
+	return nil
+}
+
 // parseBTF reads a .BTF section into memory and parses it into a list of
 // raw types and a string table.
 func parseBTF(btf io.Reader, bo binary.ByteOrder) ([]rawType, stringTable, error) {

internal/btf/btf_test.go

@@ -97,7 +97,7 @@ func TestTypeByNameAmbiguous(t *testing.T) {
 }
 
 func TestTypeByName(t *testing.T) {
-	spec, err := loadRawSpec(readVMLinux(t), binary.LittleEndian, nil, nil)
+	spec, err := LoadSpecFromReader(readVMLinux(t))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -264,6 +264,13 @@ func TestLoadKernelSpec(t *testing.T) {
 	}
 }
 
+func TestGuessBTFByteOrder(t *testing.T) {
+	bo := guessRawBTFByteOrder(readVMLinux(t))
+	if bo != binary.LittleEndian {
+		t.Fatalf("Guessed %s instead of %s", bo, binary.LittleEndian)
+	}
+}
+
 func TestSpecCopy(t *testing.T) {
 	spec := parseELFBTF(t, "../../testdata/loader-el.elf")
 

internal/cmd/gentypes/main.go

@@ -6,7 +6,6 @@ package main
 
 import (
 	"bytes"
-	"encoding/binary"
 	"fmt"
 	"os"
 	"sort"
@@ -41,7 +40,7 @@ func run(args []string) error {
 		return err
 	}
 
-	spec, err := btf.LoadRawSpec(bytes.NewReader(raw), binary.LittleEndian)
+	spec, err := btf.LoadSpecFromReader(bytes.NewReader(raw))
 	if err != nil {
 		return err
 	}

prog_test.go

@@ -593,9 +593,11 @@ func TestProgramTypeLSM(t *testing.T) {
 }
 
 func TestProgramTargetBTF(t *testing.T) {
-	// Load a file that contains valid BTF, but doesn't contain the types
-	// we need for bpf_iter.
-	fh, err := os.Open("testdata/invalid_btf_map_init-el.elf")
+	if _, err := os.Stat("/sys/kernel/btf/vmlinux"); os.IsNotExist(err) {
+		t.Skip("/sys/kernel/btf/vmlinux not present")
+	}
+
+	fh, err := os.Open("/sys/kernel/btf/vmlinux")
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -615,12 +617,11 @@ func TestProgramTargetBTF(t *testing.T) {
 	}, ProgramOptions{
 		TargetBTF: reader,
 	})
-	if err == nil {
-		prog.Close()
-	}
-	if !errors.Is(err, ErrNotSupported) {
-		t.Error("Expected ErrNotSupported, got", err)
+	testutils.SkipIfNotSupported(t, err)
+	if err != nil {
+		t.Fatal("NewProgram with TargetBTF:", err)
 	}
+	prog.Close()
 	if !reader.read {
 		t.Error("TargetBTF is not read")
 	}