[POC] bpf2go/struct_ops: implement StructOps support

cmd/bpf2go/gen/output.go

@@ -79,6 +79,10 @@ func (n templateName) Programs() string {
 	return string(n) + "Programs"
 }
 
+func (n templateName) StructOps() string {
+	return string(n) + "StructOps"
+}
+
 func (n templateName) CloseHelper() string {
 	return "_" + toUpperFirst(string(n)) + "Close"
 }
@@ -104,6 +108,62 @@ type GenerateArgs struct {
 	Output io.Writer
 	// Function which transforms the input into a valid go identifier. Uses the default behaviour if nil
 	Identifier func(string) string
+	// StructOps to be emitted
+	StructOps []StructOpsSpec
+}
+
+type StructOpsSpec struct {
+	Name string
+	Type *btf.Struct
+}
+
+// generateStructOpsShadowType produces a Go struct definition that mirrors the memory
+// layout of a BTF struct, specifically for use with struct_ops maps.
+func generateStructOpsShadowType(goTypeName string, st *btf.Struct, gf *btf.GoFormatter) (string, error) {
+	var sb strings.Builder
+
+	sb.WriteString(fmt.Sprintf("// %s is a struct type for the struct_ops map.\n", goTypeName))
+	sb.WriteString(fmt.Sprintf("type %s struct {\n", goTypeName))
+	sb.WriteString("\t_ structs.HostLayout\n")
+
+	prevOffset := uint32(0)
+	for _, m := range st.Members {
+		offset := m.Offset.Bytes()
+		if padding := offset - prevOffset; padding > 0 {
+			sb.WriteString(fmt.Sprintf("\t_ [%d]byte\n", padding))
+		}
+
+		fieldName := gf.Identifier(m.Name)
+		if fieldName == "" {
+			continue
+		}
+
+		var fieldType string
+		if ptr, ok := btf.As[*btf.Pointer](btf.UnderlyingType(m.Type)); ok {
+			if _, ok := btf.As[*btf.FuncProto](btf.UnderlyingType(ptr.Target)); ok {
+				fieldType = "*ebpf.Program"
+			}
+		}
+
+		if fieldType == "" {
+			decl, err := gf.TypeDeclaration("T", m.Type)
+			if err != nil {
+				return "", fmt.Errorf("field %s: %w", m.Name, err)
+			}
+			fieldType = strings.TrimPrefix(decl, "type T ")
+			fieldType = strings.Split(fieldType, ";")[0]
+		}
+
+		sb.WriteString(fmt.Sprintf("\t%s %s `ebpf:\"%s\"`\n", fieldName, fieldType, m.Name))
+		size, err := btf.Sizeof(m.Type)
+		if err != nil {
+			return "", fmt.Errorf("field %s: %w", m.Name, err)
+		}
+		prevOffset = offset + uint32(size)
+	}
+
+	sb.WriteString("}\n")
+	return sb.String(), nil
 }
 
 // Generate bindings for a BPF ELF file.
@@ -149,6 +209,12 @@ func Generate(args GenerateArgs) error {
 		typeNames[typ] = args.Stem + args.Identifier(typ.TypeName())
 	}
 
+	structOps := make(map[string]string)
+	for _, stOps := range args.StructOps {
+		goTypeName := "StructOps" + args.Identifier(stOps.Name)
+		structOps[stOps.Name] = goTypeName
+	}
+
 	// Ensure we don't have conflicting names and generate a sorted list of
 	// named types so that the output is stable.
 	types, err := sortTypes(typeNames)
@@ -174,6 +240,16 @@ func Generate(args GenerateArgs) error {
 		typeDecls = append(typeDecls, decl)
 	}
 
+	for _, st := range args.StructOps {
+		goTypeName := args.Stem + "StructOps" + args.Identifier(st.Name)
+		decl, err := generateStructOpsShadowType(goTypeName, st.Type, gf)
+		if err != nil {
+			return err
+		}
+		typeDecls = append(typeDecls, decl)
+		needsStructsPkg = true
+	}
+
 	ctx := struct {
 		Module           string
 		Package          string
@@ -185,6 +261,7 @@ func Generate(args GenerateArgs) error {
 		TypeDeclarations []string
 		File             string
 		NeedsStructsPkg  bool
+		StructOps        map[string]string
 	}{
 		b2gInt.CurrentModule,
 		args.Package,
@@ -196,6 +273,7 @@ func Generate(args GenerateArgs) error {
 		typeDecls,
 		args.ObjectFile,
 		needsStructsPkg,
+		structOps,
 	}
 
 	var buf bytes.Buffer

cmd/bpf2go/gen/output.tpl

@@ -94,6 +94,14 @@ type {{ .Name.Objects }} struct {
 	{{ .Name.Programs }}
 	{{ .Name.Maps }}
 	{{ .Name.Variables }}
+	{{ .Name.StructOps }}
+}
+
+// {{ .Name.StructOps }} contains all struct_ops types.
+type {{ .Name.StructOps }} struct {
+{{- range $name, $id := .StructOps }}
+	{{ $id }} {{ $.Name }}{{ $id }} `ebpf:"{{ $name }}"`
+{{- end }}
 }
 
 func (o *{{ .Name.Objects }}) Close() error {

cmd/bpf2go/gen/types.go

@@ -29,6 +29,10 @@ func CollectGlobalTypes(spec *ebpf.CollectionSpec) []btf.Type {
 // collectMapTypes collects all types used by MapSpecs.
 func collectMapTypes(types []btf.Type, maps map[string]*ebpf.MapSpec) []btf.Type {
 	for _, m := range maps {
+		if m.Type == ebpf.StructOpsMap {
+			continue
+		}
+
 		if m.Key != nil && m.Key.TypeName() != "" {
 			types = addType(types, m.Key)
 		}

cmd/bpf2go/main.go

@@ -404,11 +404,19 @@ func (b2g *bpf2go) convert(tgt gen.Target, goarches gen.GoArches) (err error) {
 	}
 
 	var maps []string
+	var structOps []gen.StructOpsSpec
 	for name := range spec.Maps {
 		// Skip .rodata, .data, .bss, etc. sections
 		if !strings.HasPrefix(name, ".") {
 			maps = append(maps, name)
 		}
+
+		ms := spec.Maps[name]
+		if ms.Type == ebpf.StructOpsMap {
+			userSt, _ := btf.As[*btf.Struct](ms.Value)
+			stOpsSpec := gen.StructOpsSpec{Name: name, Type: userSt}
+			structOps = append(structOps, stOpsSpec)
+		}
 	}
 
 	var variables []string
@@ -448,6 +456,7 @@ func (b2g *bpf2go) convert(tgt gen.Target, goarches gen.GoArches) (err error) {
 		Types:       types,
 		ObjectFile:  filepath.Base(objFileName),
 		Output:      goFile,
+		StructOps:   structOps,
 	})
 	if err != nil {
 		return fmt.Errorf("can't write %s: %s", goFileName, err)

collection.go

@@ -116,7 +116,7 @@ func copyMapOfSpecs[T interface{ Copy() T }](m map[string]T) map[string]T {
 // Returns an error if any of the eBPF objects can't be found, or
 // if the same Spec is assigned multiple times.
 func (cs *CollectionSpec) Assign(to interface{}) error {
-	getValue := func(typ reflect.Type, name string) (interface{}, error) {
+	getValue := func(typ reflect.Type, name string, fieldVal reflect.Value) (interface{}, error) {
 		switch typ {
 		case reflect.TypeOf((*ProgramSpec)(nil)):
 			if p := cs.Programs[name]; p != nil {
@@ -183,7 +183,19 @@ func (cs *CollectionSpec) LoadAndAssign(to interface{}, opts *CollectionOptions)
 	assignedProgs := make(map[string]bool)
 	assignedVars := make(map[string]bool)
 
-	getValue := func(typ reflect.Type, name string) (interface{}, error) {
+	getValue := func(typ reflect.Type, name string, val reflect.Value) (interface{}, error) {
+		ms := cs.Maps[name]
+		if ms != nil && ms.Type == StructOpsMap {
+			vType := typ
+			if vType.Kind() == reflect.Ptr {
+				vType = vType.Elem()
+			}
+
+			if vType.Kind() == reflect.Struct && typ != reflect.TypeOf((*Map)(nil)) {
+				return loader.loadStructOps(name, val)
+			}
+		}
+
 		switch typ {
 
 		case reflect.TypeOf((*Program)(nil)):
@@ -544,6 +556,38 @@ func (cl *collectionLoader) loadVariable(varName string) (*Variable, error) {
 	return v, nil
 }
 
+// loadStructOps synchronizes a user-provided shadow struct with its MapSpec
+// before the map is created in the kernel.
+//
+// If the field is already instantiated by the user, its values are patched
+// into the MapSpec's raw data buffer. This allows setting initial flags,
+// parameters, and function pointers before the kernel performs
+// struct_ops validation.
+func (cl *collectionLoader) loadStructOps(name string, field reflect.Value) (interface{}, error) {
+	ms := cl.coll.Maps[name]
+	if ms == nil {
+		return nil, fmt.Errorf("map %s: not found in loader", name)
+	}
+
+	if field.IsValid() && field.Kind() == reflect.Struct {
+		userType, ok := btf.As[*btf.Struct](ms.Value)
+		if !ok {
+			return nil, fmt.Errorf("map %s: value type is not a Struct", name)
+		}
+
+		rawData, ok := ms.Contents[0].Value.([]byte)
+		if !ok {
+			return nil, fmt.Errorf("map %s: spec contents are not a byte slice", name)
+		}
+
+		if err := structOpsPatchValue(field, userType, rawData); err != nil {
+			return nil, fmt.Errorf("patching spec from field for %s: %q", name, err)
+		}
+	}
+
+	return field.Interface(), nil
+}
+
 // populateDeferredMaps iterates maps holding programs or other maps and loads
 // any dependencies. Populates all maps in cl and freezes them if specified.
 func (cl *collectionLoader) populateDeferredMaps() error {
@@ -845,7 +889,7 @@ func (coll *Collection) Assign(to interface{}) error {
 
 	// Assign() only transfers already-loaded Maps and Programs. No extra
 	// loading is done.
-	getValue := func(typ reflect.Type, name string) (interface{}, error) {
+	getValue := func(typ reflect.Type, name string, fieldVal reflect.Value) (interface{}, error) {
 		switch typ {
 
 		case reflect.TypeOf((*Program)(nil)):
@@ -999,7 +1043,7 @@ func ebpfFields(structVal reflect.Value, visited map[reflect.Type]bool) ([]struc
 // getValue is called for every tagged field of 'to' and must return the value
 // to be assigned to the field with the given typ and name.
 func assignValues(to interface{},
-	getValue func(typ reflect.Type, name string) (interface{}, error)) error {
+	getValue func(typ reflect.Type, name string, val reflect.Value) (interface{}, error)) error {
 
 	toValue := reflect.ValueOf(to)
 	if toValue.Type().Kind() != reflect.Ptr {
@@ -1037,7 +1081,7 @@ func assignValues(to interface{},
 		}
 
 		// Get the eBPF object referred to by the tag.
-		value, err := getValue(field.Type, tag)
+		value, err := getValue(field.Type, tag, field.value)
 		if err != nil {
 			return fmt.Errorf("field %s: %w", field.Name, err)
 		}

collection_test.go

@@ -388,7 +388,7 @@ func TestNewCollectionFdLeak(t *testing.T) {
 }
 
 func TestAssignValues(t *testing.T) {
-	zero := func(t reflect.Type, name string) (interface{}, error) {
+	zero := func(t reflect.Type, name string, val reflect.Value) (interface{}, error) {
 		return reflect.Zero(t).Interface(), nil
 	}
 

examples/sched_ext/main.go

@@ -12,7 +12,9 @@ import (
 	"github.com/cilium/ebpf/rlimit"
 )
 
-//go:generate go tool bpf2go -no-global-types -tags linux bpf sched_ext.c -- -I../headers/
+const ScxOpsSwitchPartial = 1 << 3
+
+//go:generate go tool bpf2go -tags linux bpf sched_ext.c -- -I../headers/
 
 // Load a minimal defining sched_ext_ops map
 //
@@ -28,6 +30,8 @@ func main() {
 	}
 
 	objs := bpfObjects{}
+	objs.bpfStructOps.StructOpsMinimalSched.Flags = ScxOpsSwitchPartial
+	objs.bpfStructOps.StructOpsMinimalSched.TimeoutMs = 1000
 	if err := loadBpfObjects(&objs, nil); err != nil {
 		log.Fatalf("loading objects: %v", err)
 	}

examples/sched_ext/sched_ext.c

@@ -6,10 +6,14 @@
 char __license[] SEC("license") = "Dual MIT/GPL";
 
 struct sched_ext_ops {
-	char name[128];
+  s32 (*init)();
+  u64 flags;
+  u32 timeout_ms;
+  char name[128];
 };
 
 SEC(".struct_ops.link")
 struct sched_ext_ops minimal_sched = {
-	.name = "minimal",
+    .name = "minimal",
+    .timeout_ms = 5000,
 };