btf: export API to create BTF from scratch

Export Spec.Add which adds a Type to a Spec. It's then possible to create
BTF on-the-fly like so:

    spec := NewSpec()
    id, err := spec.Add(...)
    handle, err := NewHandle(&spec)

Spec.Add is responsible for allocating type IDs, which makes the BTF
encoder a lot simpler.

Attempting to create a Handle from a Spec that doesn't contain all
types will return an error. This can happen if a Type is modified
after it has been added to a Spec.

btf/btf.go

@@ -30,22 +30,26 @@ var (
 // ID represents the unique ID of a BTF object.
 type ID = sys.BTFID
 
-// Spec represents decoded BTF.
+// Spec allows querying a set of Types and loading the set into the
+// kernel.
 type Spec struct {
-	// String table from ELF used to decode split BTF, may be nil.
-	strings *stringTable
-
 	// All types contained by the spec, not including types from the base in
 	// case the spec was parsed from split BTF.
 	types []Type
 
 	// Type IDs indexed by type.
 	typeIDs map[Type]TypeID
 
+	// The last allocated type ID.
+	lastTypeID TypeID
+
 	// Types indexed by essential name.
 	// Includes all struct flavors and types with the same name.
 	namedTypes map[essentialName][]Type
 
+	// String table from ELF, may be nil.
+	strings *stringTable
+
 	// Byte order of the ELF we decoded the spec from, may be nil.
 	byteOrder binary.ByteOrder
 }
@@ -76,6 +80,18 @@ func (h *btfHeader) stringStart() int64 {
 	return int64(h.HdrLen + h.StringOff)
 }
 
+// NewSpec creates a Spec containing only Void.
+func NewSpec() *Spec {
+	return &Spec{
+		[]Type{(*Void)(nil)},
+		map[Type]TypeID{(*Void)(nil): 0},
+		0,
+		make(map[essentialName][]Type),
+		nil,
+		nil,
+	}
+}
+
 // LoadSpec opens file and calls LoadSpecFromReader on it.
 func LoadSpec(file string) (*Spec, error) {
 	fh, err := os.Open(file)
@@ -234,18 +250,19 @@ func loadRawSpec(btf io.ReaderAt, bo binary.ByteOrder, base *Spec) (*Spec, error
 		return nil, err
 	}
 
-	typeIDs, typesByName := indexTypes(types, TypeID(len(baseTypes)))
+	typeIDs, typesByName, lastTypeID := indexTypes(types, TypeID(len(baseTypes)))
 
 	return &Spec{
 		namedTypes: typesByName,
 		typeIDs:    typeIDs,
 		types:      types,
+		lastTypeID: lastTypeID,
 		strings:    rawStrings,
 		byteOrder:  bo,
 	}, nil
 }
 
-func indexTypes(types []Type, typeIDOffset TypeID) (map[Type]TypeID, map[essentialName][]Type) {
+func indexTypes(types []Type, typeIDOffset TypeID) (map[Type]TypeID, map[essentialName][]Type, TypeID) {
 	namedTypes := 0
 	for _, typ := range types {
 		if typ.TypeName() != "" {
@@ -259,14 +276,16 @@ func indexTypes(types []Type, typeIDOffset TypeID) (map[Type]TypeID, map[essenti
 	typeIDs := make(map[Type]TypeID, len(types))
 	typesByName := make(map[essentialName][]Type, namedTypes)
 
+	var lastTypeID TypeID
 	for i, typ := range types {
 		if name := newEssentialName(typ.TypeName()); name != "" {
 			typesByName[name] = append(typesByName[name], typ)
 		}
-		typeIDs[typ] = TypeID(i) + typeIDOffset
+		lastTypeID = TypeID(i) + typeIDOffset
+		typeIDs[typ] = lastTypeID
 	}
 
-	return typeIDs, typesByName
+	return typeIDs, typesByName, lastTypeID
 }
 
 // LoadKernelSpec returns the current kernel's BTF information.
@@ -478,15 +497,15 @@ func fixupDatasec(types []Type, sectionSizes map[string]uint32, offsets map[symb
 // Copy creates a copy of Spec.
 func (s *Spec) Copy() *Spec {
 	types := copyTypes(s.types, nil)
-
-	typeIDs, typesByName := indexTypes(types, s.firstTypeID())
+	typeIDs, typesByName, lastTypeID := indexTypes(types, s.firstTypeID())
 
 	// NB: Other parts of spec are not copied since they are immutable.
 	return &Spec{
-		s.strings,
 		types,
 		typeIDs,
+		lastTypeID,
 		typesByName,
+		s.strings,
 		s.byteOrder,
 	}
 }
@@ -501,6 +520,41 @@ func (sw sliceWriter) Write(p []byte) (int, error) {
 	return copy(sw, p), nil
 }
 
+// Add a Type.
+//
+// Adding the identical Type multiple times is valid and will return a stable ID.
+//
+// See [Type] for details on identity.
+func (s *Spec) Add(typ Type) (TypeID, error) {
+	if typ == nil {
+		return 0, fmt.Errorf("canot add nil Type")
+	}
+
+	hasID := func(t Type) (skip bool) {
+		_, isVoid := t.(*Void)
+		_, alreadyEncoded := s.typeIDs[t]
+		return isVoid || alreadyEncoded
+	}
+
+	iter := postorderTraversal(typ, hasID)
+	for iter.Next() {
+		id := s.lastTypeID + 1
+		if id < s.lastTypeID {
+			return 0, fmt.Errorf("type ID overflow")
+		}
+
+		s.typeIDs[iter.Type] = id
+		s.types = append(s.types, iter.Type)
+		s.lastTypeID = id
+
+		if name := newEssentialName(iter.Type.TypeName()); name != "" {
+			s.namedTypes[name] = append(s.namedTypes[name], iter.Type)
+		}
+	}
+
+	return s.TypeID(typ)
+}
+
 // TypeByID returns the BTF Type with the given type ID.
 //
 // Returns an error wrapping ErrNotFound if a Type with the given ID

btf/btf_test.go

@@ -218,6 +218,41 @@ func TestTypeByName(t *testing.T) {
 	}
 }
 
+func TestSpecAdd(t *testing.T) {
+	i := &Int{
+		Name:     "foo",
+		Size:     2,
+		Encoding: Signed | Char,
+	}
+
+	s := NewSpec()
+	id, err := s.Add(i)
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(1), qt.Commentf("First non-void type doesn't get id 1"))
+	id, err = s.Add(i)
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(1), qt.Commentf("Adding a type twice returns different ids"))
+
+	id, err = s.TypeID(i)
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(1))
+
+	id, err = s.Add(&Pointer{i})
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(2))
+
+	id, err = s.Add(&Typedef{"baz", i})
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(3))
+
+	typ, err := s.AnyTypeByName("foo")
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, typ, qt.Equals, i)
+
+	_, err = s.AnyTypeByName("baz")
+	qt.Assert(t, err, qt.IsNil)
+}
+
 func BenchmarkParseVmlinux(b *testing.B) {
 	rd := vmlinuxTestdataReader(b)
 	b.ReportAllocs()
@@ -329,8 +364,12 @@ func TestLoadSpecFromElf(t *testing.T) {
 }
 
 func TestVerifierError(t *testing.T) {
-	btf, _ := newEncoder(kernelEncoderOptions, nil).Encode()
-	_, err := newHandleFromRawBTF(btf)
+	var buf bytes.Buffer
+	if err := marshalSpec(&buf, NewSpec(), nil, nil); err != nil {
+		t.Fatal(err)
+	}
+
+	_, err := newHandleFromRawBTF(buf.Bytes())
 	testutils.SkipIfNotSupported(t, err)
 	var ve *internal.VerifierError
 	if !errors.As(err, &ve) {

btf/ext_info.go

@@ -8,7 +8,6 @@ import (
 	"io"
 	"math"
 	"sort"
-	"sync"
 
 	"github.com/cilium/ebpf/asm"
 	"github.com/cilium/ebpf/internal"
@@ -131,12 +130,6 @@ func (ei *ExtInfos) Assign(insns asm.Instructions, section string) {
 	}
 }
 
-var nativeEncoderPool = sync.Pool{
-	New: func() any {
-		return newEncoder(kernelEncoderOptions, nil)
-	},
-}
-
 // MarshalExtInfos encodes function and line info embedded in insns into kernel
 // wire format.
 //
@@ -157,23 +150,19 @@ func MarshalExtInfos(insns asm.Instructions) (_ *Handle, funcInfos, lineInfos []
 		}
 	}
 
-	// Avoid allocating encoder, etc. if there is no BTF at all.
 	return nil, nil, nil, nil
 
 marshal:
-	enc := nativeEncoderPool.Get().(*encoder)
-	defer nativeEncoderPool.Put(enc)
-
-	enc.Reset()
-
+	stb := newStringTableBuilder(0)
+	spec := NewSpec()
 	var fiBuf, liBuf bytes.Buffer
 	for {
 		if fn := FuncMetadata(iter.Ins); fn != nil {
 			fi := &funcInfo{
 				fn:     fn,
 				offset: iter.Offset,
 			}
-			if err := fi.marshal(&fiBuf, enc); err != nil {
+			if err := fi.marshal(&fiBuf, spec); err != nil {
 				return nil, nil, nil, fmt.Errorf("write func info: %w", err)
 			}
 		}
@@ -183,7 +172,7 @@ marshal:
 				line:   line,
 				offset: iter.Offset,
 			}
-			if err := li.marshal(&liBuf, enc.strings); err != nil {
+			if err := li.marshal(&liBuf, stb); err != nil {
 				return nil, nil, nil, fmt.Errorf("write line info: %w", err)
 			}
 		}
@@ -193,12 +182,14 @@ marshal:
 		}
 	}
 
-	btf, err := enc.Encode()
-	if err != nil {
-		return nil, nil, nil, err
+	buf := getBuffer()
+	defer putBuffer(buf)
+
+	if err := marshalSpec(buf, spec, stb, kernelMarshalOptions); err != nil {
+		return nil, nil, nil, fmt.Errorf("marshal BTF: %w", err)
 	}
 
-	handle, err := newHandleFromRawBTF(btf)
+	handle, err := newHandleFromRawBTF(buf.Bytes())
 	return handle, fiBuf.Bytes(), liBuf.Bytes(), err
 }
 
@@ -392,8 +383,8 @@ func newFuncInfos(bfis []bpfFuncInfo, ts types) ([]funcInfo, error) {
 }
 
 // marshal into the BTF wire format.
-func (fi *funcInfo) marshal(w *bytes.Buffer, enc *encoder) error {
-	id, err := enc.Add(fi.fn)
+func (fi *funcInfo) marshal(w *bytes.Buffer, spec *Spec) error {
+	id, err := spec.Add(fi.fn)
 	if err != nil {
 		return err
 	}

btf/handle.go

@@ -30,26 +30,19 @@ func NewHandle(spec *Spec) (*Handle, error) {
 		return nil, fmt.Errorf("can't load %s BTF on %s", spec.byteOrder, internal.NativeEndian)
 	}
 
-	var stb *stringTableBuilder
-	if spec.strings != nil {
-		stb = newStringTableBuilderFromTable(spec.strings)
+	if spec.firstTypeID() != 0 {
+		return nil, fmt.Errorf("split BTF can't be loaded into the kernel")
 	}
 
-	enc := newEncoder(kernelEncoderOptions, stb)
+	buf := getBuffer()
+	defer putBuffer(buf)
 
-	for _, typ := range spec.types {
-		_, err := enc.Add(typ)
-		if err != nil {
-			return nil, fmt.Errorf("add %s: %w", typ, err)
-		}
-	}
-
-	btf, err := enc.Encode()
+	err := marshalSpec(buf, spec, nil, kernelMarshalOptions)
 	if err != nil {
 		return nil, fmt.Errorf("marshal BTF: %w", err)
 	}
 
-	return newHandleFromRawBTF(btf)
+	return newHandleFromRawBTF(buf.Bytes())
 }
 
 func newHandleFromRawBTF(btf []byte) (*Handle, error) {