Repository: cilium/ebpf
Pull request #124 (2020-08-07 08:27:28)

collection: add Assign helpers
Interacting with a CollectionSpec or Collection involves a lot
of boilerplate, since checking that all required programs are
present is cumbersome. It becomes really unwieldy if you're
trying to correctly clean up Programs and Maps on error:

   prog := coll.DetachProgram("foo")
   if prog == nil {
       // error!
   }
   defer prog.Close() // or maybe only close on error?

   map := coll.DetachMap("bar")
   // ...

It also means that the names of programs and maps are spread
all over the code base.

Add helpers to CollectionSpec and Collection to simplify to
something like this:

   var objs struct{
       Program *ebpf.Program `ebpf:"foo"`
       Map     *ebpf.Map     `ebpf:"map"`
   }

   if err := coll.Assign(&objs); err != nil {
       // Either foo or bar doesn't exist!
   }

These functions will be used as the basis to build a code
generator similar to libbpf skeletons.


Patch:
diff --git a/collection.go b/collection.go
index 0c8b65d..1ed34c0 100644
--- a/collection.go
+++ b/collection.go
@@ -4,6 +4,8 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"reflect"
+	"strings"
 
 	"github.com/cilium/ebpf/asm"
 	"github.com/cilium/ebpf/internal"
@@ -126,6 +128,65 @@ func (cs *CollectionSpec) RewriteConstants(consts map[string]interface{}) error
 	return nil
 }
 
+// Assign the contents of a collection spec to a struct.
+//
+// This function is a short-cut to manually checking the presence
+// of maps and programs in a collection spec.
+//
+// The argument to must be a pointer to a struct. A field of the
+// struct is updated with values from Programs or Maps if it
+// has an `ebpf` tag and its type is *ProgramSpec or *MapSpec.
+// The tag gives the name of the program or map as found in
+// the CollectionSpec.
+//
+//    struct {
+//        Foo     *ebpf.ProgramSpec `ebpf:"xdp_foo"`
+//        Bar     *ebpf.MapSpec     `ebpf:"bar_map"`
+//        Ignored int
+//    }
+//
+// Returns an error if any of the fields can't be found, or
+// if the same map or program is assigned multiple times.
+func (cs *CollectionSpec) Assign(to interface{}) error {
+	valueOf := func(typ reflect.Type, name string) (reflect.Value, error) {
+		switch typ {
+		case reflect.TypeOf((*ProgramSpec)(nil)):
+			p := cs.Programs[name]
+			if p == nil {
+				return reflect.Value{}, fmt.Errorf("missing program %q", name)
+			}
+			return reflect.ValueOf(p), nil
+		case reflect.TypeOf((*MapSpec)(nil)):
+			m := cs.Maps[name]
+			if m == nil {
+				return reflect.Value{}, fmt.Errorf("missing map %q", name)
+			}
+			return reflect.ValueOf(m), nil
+		default:
+			return reflect.Value{}, fmt.Errorf("unsupported type %s", typ)
+		}
+	}
+
+	return assignValues(to, valueOf)
+}
+
+// LoadAndAssign creates a collection from a spec, and assigns it to a struct.
+//
+// See Collection.Assign for details.
+func (cs *CollectionSpec) LoadAndAssign(to interface{}, opts *CollectionOptions) error {
+	if opts == nil {
+		opts = &CollectionOptions{}
+	}
+
+	coll, err := NewCollectionWithOptions(cs, *opts)
+	if err != nil {
+		return err
+	}
+	defer coll.Close()
+
+	return coll.Assign(to)
+}
+
 // Collection is a collection of Programs and Maps associated
 // with their symbols
 type Collection struct {
@@ -292,3 +353,105 @@ func (coll *Collection) DetachProgram(name string) *Program {
 	delete(coll.Programs, name)
 	return p
 }
+
+// Assign the contents of a collection to a struct.
+//
+// `to` must be a pointer to a struct like the following:
+//
+//    struct {
+//        Foo     *ebpf.Program `ebpf:"xdp_foo"`
+//        Bar     *ebpf.Map     `ebpf:"bar_map"`
+//        Ignored int
+//    }
+//
+// See CollectionSpec.Assign for the semantics of this function.
+//
+// DetachMap and DetachProgram is invoked for all assigned elements
+// if the function is successful.
+func (coll *Collection) Assign(to interface{}) error {
+	assignedMaps := make(map[string]struct{})
+	assignedPrograms := make(map[string]struct{})
+	valueOf := func(typ reflect.Type, name string) (reflect.Value, error) {
+		switch typ {
+		case reflect.TypeOf((*Program)(nil)):
+			p := coll.Programs[name]
+			if p == nil {
+				return reflect.Value{}, fmt.Errorf("missing program %q", name)
+			}
+			assignedPrograms[name] = struct{}{}
+			return reflect.ValueOf(p), nil
+		case reflect.TypeOf((*Map)(nil)):
+			m := coll.Maps[name]
+			if m == nil {
+				return reflect.Value{}, fmt.Errorf("missing map %q", name)
+			}
+			assignedMaps[name] = struct{}{}
+			return reflect.ValueOf(m), nil
+		default:
+			return reflect.Value{}, fmt.Errorf("unsupported type %s", typ)
+		}
+	}
+
+	if err := assignValues(to, valueOf); err != nil {
+		return err
+	}
+
+	for name := range assignedPrograms {
+		coll.DetachProgram(name)
+	}
+
+	for name := range assignedMaps {
+		coll.DetachMap(name)
+	}
+
+	return nil
+}
+
+func assignValues(to interface{}, valueOf func(reflect.Type, string) (reflect.Value, error)) error {
+	v := reflect.ValueOf(to)
+	if v.Kind() != reflect.Ptr || v.Elem().Kind() != reflect.Struct {
+		return fmt.Errorf("%T is not a pointer to a struct", to)
+	}
+
+	type elem struct {
+		typ  reflect.Type
+		name string
+	}
+
+	var (
+		s          = v.Elem()
+		sT         = s.Type()
+		assignedTo = make(map[elem]string)
+	)
+	for i := 0; i < sT.NumField(); i++ {
+		field := sT.Field(i)
+
+		name := field.Tag.Get("ebpf")
+		if name == "" {
+			continue
+		}
+		if strings.Contains(name, ",") {
+			return fmt.Errorf("field %s: ebpf tag contains a comma", field.Name)
+		}
+
+		e := elem{field.Type, name}
+		if assignedField := assignedTo[e]; assignedField != "" {
+			return fmt.Errorf("field %s: %q was already assigned to %s", field.Name, name, assignedField)
+		}
+
+		value, err := valueOf(field.Type, name)
+		if err != nil {
+			return fmt.Errorf("field %s: %w", field.Name, err)
+		}
+
+		fieldValue := s.Field(i)
+		if !fieldValue.CanSet() {
+			return fmt.Errorf("can't set value of field %s", field.Name)
+		}
+
+		fieldValue.Set(value)
+		assignedTo[e] = field.Name
+	}
+
+	return nil
+}


Test patch:
diff --git a/collection_test.go b/collection_test.go
index e1b9a17..484e84e 100644
--- a/collection_test.go
+++ b/collection_test.go
@@ -1,6 +1,7 @@
 package ebpf
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/cilium/ebpf/asm"
@@ -153,3 +154,204 @@ func TestCollectionSpecRewriteMaps(t *testing.T) {
 		t.Fatal("new / override map not used")
 	}
 }
+
+func TestCollectionAssign(t *testing.T) {
+	var specs struct {
+		Program *ProgramSpec `ebpf:"prog1"`
+		Map     *MapSpec     `ebpf:"map1"`
+	}
+
+	mapSpec := &MapSpec{
+		Type:       Array,
+		KeySize:    4,
+		ValueSize:  4,
+		MaxEntries: 1,
+	}
+	progSpec := &ProgramSpec{
+		Type: SocketFilter,
+		Instructions: asm.Instructions{
+			asm.LoadImm(asm.R0, 0, asm.DWord),
+			asm.Return(),
+		},
+		License: "MIT",
+	}
+
+	cs := &CollectionSpec{
+		Maps: map[string]*MapSpec{
+			"map1": mapSpec,
+		},
+		Programs: map[string]*ProgramSpec{
+			"prog1": progSpec,
+		},
+	}
+
+	if err := cs.Assign(&specs); err != nil {
+		t.Fatal("Can't assign spec:", err)
+	}
+
+	if specs.Program != progSpec {
+		t.Fatalf("Expected Program to be %p, got %p", progSpec, specs.Program)
+	}
+
+	if specs.Map != mapSpec {
+		t.Fatalf("Expected Map to be %p, got %p", mapSpec, specs.Map)
+	}
+
+	if err := cs.Assign(new(int)); err == nil {
+		t.Fatal("Assign allows to besides *struct")
+	}
+
+	if err := cs.Assign(new(struct{ Foo int })); err != nil {
+		t.Fatal("Assign doesn't ignore untagged fields")
+	}
+
+	unexported := new(struct {
+		foo *MapSpec `ebpf:"map1"`
+	})
+
+	if err := cs.Assign(unexported); err == nil {
+		t.Error("Assign should return an error on unexported fields")
+	}
+
+	coll, err := NewCollection(cs)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer coll.Close()
+
+	var objs struct {
+		Program *Program `ebpf:"prog1"`
+		Map     *Map     `ebpf:"map1"`
+	}
+
+	if err := coll.Assign(&objs); err != nil {
+		t.Fatal("Can't Assign objects:", err)
+	}
+	objs.Program.Close()
+	objs.Map.Close()
+
+	if coll.Programs["prog1"] != nil {
+		t.Fatal("Assign doesn't detach Program")
+	}
+
+	if coll.Maps["map1"] != nil {
+		t.Fatal("Assign doesn't detach Map")
+	}
+}
+
+func ExampleCollectionSpec_Assign() {
+	spec := &CollectionSpec{
+		Maps: map[string]*MapSpec{
+			"map1": {
+				Type:       Array,
+				KeySize:    4,
+				ValueSize:  4,
+				MaxEntries: 1,
+			},
+		},
+		Programs: map[string]*ProgramSpec{
+			"prog1": {
+				Type: SocketFilter,
+				Instructions: asm.Instructions{
+					asm.LoadImm(asm.R0, 0, asm.DWord),
+					asm.Return(),
+				},
+				License: "MIT",
+			},
+		},
+	}
+
+	var specs struct {
+		Program *ProgramSpec `ebpf:"prog1"`
+		Map     *MapSpec     `ebpf:"map1"`
+	}
+
+	if err := spec.Assign(&specs); err != nil {
+		panic(err)
+	}
+
+	fmt.Println(specs.Program.Type)
+	fmt.Println(specs.Map.Type)
+
+	// Output: SocketFilter
+	// Array
+}
+
+func ExampleCollectionSpec_LoadAndAssign() {
+	spec := &CollectionSpec{
+		Maps: map[string]*MapSpec{
+			"map1": {
+				Type:       Array,
+				KeySize:    4,
+				ValueSize:  4,
+				MaxEntries: 1,
+			},
+		},
+		Programs: map[string]*ProgramSpec{
+			"prog1": {
+				Type: SocketFilter,
+				Instructions: asm.Instructions{
+					asm.LoadImm(asm.R0, 0, asm.DWord),
+					asm.Return(),
+				},
+				License: "MIT",
+			},
+		},
+	}
+
+	var objs struct {
+		Program *Program `ebpf:"prog1"`
+		Map     *Map     `ebpf:"map1"`
+	}
+
+	if err := spec.LoadAndAssign(&objs, nil); err != nil {
+		panic(err)
+	}
+
+	fmt.Println(objs.Program.ABI().Type)
+	fmt.Println(objs.Map.ABI().Type)
+
+	// Output: SocketFilter
+	// Array
+}
+
+func ExampleCollection_Assign() {
+	coll, err := NewCollection(&CollectionSpec{
+		Maps: map[string]*MapSpec{
+			"map1": {
+				Type:       Array,
+				KeySize:    4,
+				ValueSize:  4,
+				MaxEntries: 1,
+			},
+		},
+		Programs: map[string]*ProgramSpec{
+			"prog1": {
+				Type: SocketFilter,
+				Instructions: asm.Instructions{
+					asm.LoadImm(asm.R0, 0, asm.DWord),
+					asm.Return(),
+				},
+				License: "MIT",
+			},
+		},
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	var objs struct {
+		Program *Program `ebpf:"prog1"`
+		Map     *Map     `ebpf:"map1"`
+	}
+
+	if err := coll.Assign(&objs); err != nil {
+		panic(err)
+	}
+
+	fmt.Println(objs.Program.ABI().Type)
+	fmt.Println(objs.Map.ABI().Type)
+
+	// Output: SocketFilter
+	// Array
+}
diff --git a/example_program_test.go b/example_program_test.go
index debac07..5da90a2 100644
--- a/example_program_test.go
+++ b/example_program_test.go
@@ -28,10 +28,10 @@ func getTracepointID() (uint64, error) {
 	return strconv.ParseUint(tid, 10, 64)
 }
 
-// Example_program demonstrates how to attach an eBPF program to a tracepoint.
+// This demonstrates how to attach an eBPF program to a tracepoint.
 // The program will be attached to the sys_enter_open syscall and print out the integer
 // 123 everytime the sycall is used.
-func Example_program() {
+func Example_tracepoint() {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	events, err := ebpf.NewMap(&ebpf.MapSpec{
diff --git a/example_sock_elf_test.go b/example_sock_elf_test.go
index 0075a04..924c4be 100644
--- a/example_sock_elf_test.go
+++ b/example_sock_elf_test.go
@@ -90,11 +90,16 @@ func Example_socketELF() {
 		panic(err)
 	}
 
-	coll, err := ebpf.NewCollection(spec)
-	if err != nil {
+	var objs struct {
+		Prog  *ebpf.Program `ebpf:"bpf_prog1"`
+		Stats *ebpf.Map     `ebpf:"my_map"`
+	}
+
+	if err := spec.LoadAndAssign(&objs, nil); err != nil {
 		panic(err)
 	}
-	defer coll.Close()
+	defer objs.Prog.Close()
+	defer objs.Stats.Close()
 
 	sock, err := openRawSock(*index)
 	if err != nil {
@@ -102,25 +107,13 @@ func Example_socketELF() {
 	}
 	defer syscall.Close(sock)
 
-	prog := coll.DetachProgram("bpf_prog1")
-	if prog == nil {
-		panic("no program named bpf_prog1 found")
-	}
-	defer prog.Close()
-
-	if err := syscall.SetsockoptInt(sock, syscall.SOL_SOCKET, SO_ATTACH_BPF, prog.FD()); err != nil {
+	if err := syscall.SetsockoptInt(sock, syscall.SOL_SOCKET, SO_ATTACH_BPF, objs.Prog.FD()); err != nil {
 		panic(err)
 	}
 
 	fmt.Printf("Filtering on eth index: %d\n", *index)
 	fmt.Println("Packet stats:")
 
-	protoStats := coll.DetachMap("my_map")
-	if protoStats == nil {
-		panic(fmt.Errorf("no map named my_map found"))
-	}
-	defer protoStats.Close()
-
 	for {
 		const (
 			ICMP = 0x01
@@ -132,15 +125,15 @@ func Example_socketELF() {
 		var icmp uint64
 		var tcp uint64
 		var udp uint64
-		err := protoStats.Lookup(uint32(ICMP), &icmp)
+		err := objs.Stats.Lookup(uint32(ICMP), &icmp)
 		if err != nil {
 			panic(err)
 		}
-		err = protoStats.Lookup(uint32(TCP), &tcp)
+		err = objs.Stats.Lookup(uint32(TCP), &tcp)
 		if err != nil {
 			panic(err)
 		}
-		err = protoStats.Lookup(uint32(UDP), &udp)
+		err = objs.Stats.Lookup(uint32(UDP), &udp)
 		if err != nil {
 			panic(err)
 		}