Repository: cilium/ebpf
Pull request #520 (2021-12-07 16:21:09)

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.

lmb: This is good to go if someone has time for a look.

Patch:
diff --git a/internal/btf/btf.go b/internal/btf/btf.go
index 8c36f8b..4ddba6d 100644
--- a/internal/btf/btf.go
+++ b/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) {
diff --git a/internal/cmd/gentypes/main.go b/internal/cmd/gentypes/main.go
index ee71060..cef8c46 100644
--- a/internal/cmd/gentypes/main.go
+++ b/internal/cmd/gentypes/main.go
@@ -6,8 +6,6 @@ package main
 
 import (
 	"bytes"
-	"compress/gzip"
-	"encoding/binary"
 	"fmt"
 	"os"
 	"sort"
@@ -37,19 +35,12 @@ func run(args []string) error {
 		return fmt.Errorf("expect location of compressed vmlinux .BTF as argument")
 	}
 
-	fh, err := os.Open(args[0])
+	raw, err := internal.ReadAllCompressed(args[0])
 	if err != nil {
 		return err
 	}
-	defer fh.Close()
 
-	gz, err := gzip.NewReader(fh)
-	if err != nil {
-		return err
-	}
-	defer gz.Close()
-
-	spec, err := btf.LoadRawSpec(gz, binary.LittleEndian)
+	spec, err := btf.LoadSpecFromReader(bytes.NewReader(raw))
 	if err != nil {
 		return err
 	}
diff --git a/internal/io.go b/internal/io.go
index 3ed432f..7177e59 100644
--- a/internal/io.go
+++ b/internal/io.go
@@ -2,6 +2,7 @@ package internal
 
 import (
 	"bufio"
+	"compress/gzip"
 	"errors"
 	"io"
 	"os"
@@ -42,3 +43,20 @@ func (DiscardZeroes) Write(p []byte) (int, error) {
 	}
 	return len(p), nil
 }
+
+// ReadAllCompressed decompresses a gzipped file into memory.
+func ReadAllCompressed(file string) ([]byte, error) {
+	fh, err := os.Open(file)
+	if err != nil {
+		return nil, err
+	}
+	defer fh.Close()
+
+	gz, err := gzip.NewReader(fh)
+	if err != nil {
+		return nil, err
+	}
+	defer gz.Close()
+
+	return io.ReadAll(gz)
+}


Test patch:
diff --git a/internal/btf/btf_test.go b/internal/btf/btf_test.go
index 026229a..cd54431 100644
--- a/internal/btf/btf_test.go
+++ b/internal/btf/btf_test.go
@@ -2,36 +2,36 @@ package btf
 
 import (
 	"bytes"
-	"compress/gzip"
 	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
 	"os"
+	"sync"
 	"testing"
 
 	"github.com/cilium/ebpf/internal"
 	"github.com/cilium/ebpf/internal/testutils"
 )
 
-func parseVMLinuxBTF(tb testing.TB) (*Spec, error) {
-	fh, err := os.Open("testdata/vmlinux-btf.gz")
-	if err != nil {
-		tb.Fatal(err)
-	}
-	defer fh.Close()
+var vmlinux struct {
+	sync.Once
+	err error
+	raw []byte
+}
 
-	rd, err := gzip.NewReader(fh)
-	if err != nil {
-		tb.Fatal(err)
-	}
+func readVMLinux(tb testing.TB) *bytes.Reader {
+	tb.Helper()
 
-	spec, err := loadRawSpec(rd, binary.LittleEndian, nil, nil)
-	if err != nil {
-		tb.Fatal("Can't load BTF:", err)
+	vmlinux.Do(func() {
+		vmlinux.raw, vmlinux.err = internal.ReadAllCompressed("testdata/vmlinux-btf.gz")
+	})
+
+	if vmlinux.err != nil {
+		tb.Fatal(vmlinux.err)
 	}
 
-	return spec, nil
+	return bytes.NewReader(vmlinux.raw)
 }
 
 func parseELFBTF(tb testing.TB, file string) *Spec {
@@ -73,15 +73,35 @@ func TestAnyTypesByName(t *testing.T) {
 	})
 }
 
+func TestTypeByNameAmbiguous(t *testing.T) {
+	testutils.Files(t, testutils.Glob(t, "testdata/relocs-*.elf"), func(t *testing.T, file string) {
+		spec := parseELFBTF(t, file)
+
+		var typ *Struct
+		if err := spec.TypeByName("ambiguous", &typ); err != nil {
+			t.Fatal(err)
+		}
+
+		if name := typ.TypeName(); name != "ambiguous" {
+			t.Fatal("expected type name 'ambiguous', got:", name)
+		}
+
+		if err := spec.TypeByName("ambiguous___flavour", &typ); err != nil {
+			t.Fatal(err)
+		}
+
+		if name := typ.TypeName(); name != "ambiguous___flavour" {
+			t.Fatal("expected type name 'ambiguous___flavour', got:", name)
+		}
+	})
+}
+
 func TestTypeByName(t *testing.T) {
-	spec, err := parseVMLinuxBTF(t)
+	spec, err := LoadSpecFromReader(readVMLinux(t))
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	// spec.TypeByName MUST fail if typ is a nil btf.Type.
-	i := 0
-	p := &i
 	for _, typ := range []interface{}{
 		nil,
 		Struct{},
@@ -90,12 +110,15 @@ func TestTypeByName(t *testing.T) {
 		&[]Struct{},
 		map[int]Struct{},
 		&map[int]Struct{},
-		p,
-		&p,
+		int(0),
+		new(int),
 	} {
-		if err := spec.TypeByName("iphdr", typ); err == nil {
-			t.Fatalf("TypeByName does not fail with type %T", typ)
-		}
+		t.Run(fmt.Sprintf("%T", typ), func(t *testing.T) {
+			// spec.TypeByName MUST fail if typ is a nil btf.Type.
+			if err := spec.TypeByName("iphdr", typ); err == nil {
+				t.Fatalf("FindType does not fail with type %T", typ)
+			}
+		})
 	}
 
 	// spec.TypeByName MUST return the same address for multiple calls with the same type name.
@@ -110,43 +133,8 @@ func TestTypeByName(t *testing.T) {
 	if iphdr1 != iphdr2 {
 		t.Fatal("multiple TypeByName calls for `iphdr` name do not return the same addresses")
 	}
-}
-
-func TestTypeByNameAmbiguous(t *testing.T) {
-	testutils.Files(t, testutils.Glob(t, "testdata/relocs-*.elf"), func(t *testing.T, file string) {
-		spec := parseELFBTF(t, file)
-
-		var typ *Struct
-		if err := spec.TypeByName("ambiguous", &typ); err != nil {
-			t.Fatal(err)
-		}
-
-		if name := typ.TypeName(); name != "ambiguous" {
-			t.Fatal("expected type name 'ambiguous', got:", name)
-		}
-
-		if err := spec.TypeByName("ambiguous___flavour", &typ); err != nil {
-			t.Fatal(err)
-		}
-
-		if name := typ.TypeName(); name != "ambiguous___flavour" {
-			t.Fatal("expected type name 'ambiguous___flavour', got:", name)
-		}
-	})
-}
-
-func TestParseVmlinux(t *testing.T) {
-	spec, err := parseVMLinuxBTF(t)
-	if err != nil {
-		t.Fatal(err)
-	}
 
-	var iphdr *Struct
-	err = spec.TypeByName("iphdr", &iphdr)
-	if err != nil {
-		t.Fatalf("unable to find `iphdr` struct: %s", err)
-	}
-	for _, m := range iphdr.Members {
+	for _, m := range iphdr1.Members {
 		if m.Name == "version" {
 			// __u8 is a typedef
 			td, ok := m.Type.(*Typedef)
@@ -172,23 +160,7 @@ func TestParseVmlinux(t *testing.T) {
 }
 
 func BenchmarkParseVmlinux(b *testing.B) {
-	fh, err := os.Open("testdata/vmlinux-btf.gz")
-	if err != nil {
-		b.Fatal(err)
-	}
-	defer fh.Close()
-
-	gr, err := gzip.NewReader(fh)
-	if err != nil {
-		b.Fatal(err)
-	}
-
-	var buf bytes.Buffer
-	if _, err := buf.ReadFrom(gr); err != nil {
-		b.Fatal(err)
-	}
-	rd := bytes.NewReader(buf.Bytes())
-
+	rd := readVMLinux(b)
 	b.ResetTimer()
 
 	for n := 0; n < b.N; n++ {
@@ -292,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")
 
diff --git a/prog_test.go b/prog_test.go
index ac6970e..f0fd705 100644
--- a/prog_test.go
+++ b/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")
 	}