btf: allow passing *Type to Spec.TypeByName

lmb · Jul 14, 2022 · f124d1c · f124d1c
1 parent bf69929
commit f124d1c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 16 deletions.
diff --git a/btf/btf.go b/btf/btf.go
@@ -609,6 +609,9 @@ func (s *Spec) AnyTypeByName(name string) (Type, error) {
 // Type exists in the Spec. If multiple candidates are found,
 // an error is returned.
 func (s *Spec) TypeByName(name string, typ interface{}) error {
+	typeInterface := reflect.TypeOf((*Type)(nil)).Elem()
+
+	// typ may be **T or *Type
 	typValue := reflect.ValueOf(typ)
 	if typValue.Kind() != reflect.Ptr {
 		return fmt.Errorf("%T is not a pointer", typ)
@@ -620,7 +623,12 @@ func (s *Spec) TypeByName(name string, typ interface{}) error {
 	}
 
 	wanted := typPtr.Type()
-	if !wanted.AssignableTo(reflect.TypeOf((*Type)(nil)).Elem()) {
+	if wanted == typeInterface {
+		// This is *Type. Unwrap the value's type.
+		wanted = typPtr.Elem().Type()
+	}
+
+	if !wanted.AssignableTo(typeInterface) {
 		return fmt.Errorf("%T does not satisfy Type interface", typ)
 	}
 
@@ -643,7 +651,7 @@ func (s *Spec) TypeByName(name string, typ interface{}) error {
 	}
 
 	if candidate == nil {
-		return fmt.Errorf("type %s: %w", name, ErrNotFound)
+		return fmt.Errorf("%s %s: %w", wanted, name, ErrNotFound)
 	}
 
 	typPtr.Set(reflect.ValueOf(candidate))

diff --git a/btf/btf_test.go b/btf/btf_test.go
@@ -130,6 +130,12 @@ func TestTypeByName(t *testing.T) {
 		t.Fatal("multiple TypeByName calls for `iphdr` name do not return the same addresses")
 	}
 
+	// It's valid to pass a *Type to TypeByName.
+	typ := Type(iphdr2)
+	if err := spec.TypeByName("iphdr", &typ); err != nil {
+		t.Fatal("Can't look up using *Type:", err)
+	}
+
 	// Excerpt from linux/ip.h, https://elixir.bootlin.com/linux/latest/A/ident/iphdr
 	//
 	// struct iphdr {

diff --git a/prog.go b/prog.go
@@ -782,29 +782,34 @@ func findTargetInKernel(spec *btf.Spec, name string, progType ProgramType, attac
 
 	var (
 		typeName, featureName string
-		isBTFTypeFunc         = true
+		target                btf.Type
 	)
 
 	switch (match{progType, attachType}) {
 	case match{LSM, AttachLSMMac}:
 		typeName = "bpf_lsm_" + name
 		featureName = name + " LSM hook"
+		target = (*btf.Func)(nil)
 	case match{Tracing, AttachTraceIter}:
 		typeName = "bpf_iter_" + name
 		featureName = name + " iterator"
+		target = (*btf.Func)(nil)
 	case match{Tracing, AttachTraceFEntry}:
 		typeName = name
 		featureName = fmt.Sprintf("fentry %s", name)
+		target = (*btf.Func)(nil)
 	case match{Tracing, AttachTraceFExit}:
 		typeName = name
 		featureName = fmt.Sprintf("fexit %s", name)
+		target = (*btf.Func)(nil)
 	case match{Tracing, AttachModifyReturn}:
 		typeName = name
 		featureName = fmt.Sprintf("fmod_ret %s", name)
+		target = (*btf.Func)(nil)
 	case match{Tracing, AttachTraceRawTp}:
 		typeName = fmt.Sprintf("btf_trace_%s", name)
 		featureName = fmt.Sprintf("raw_tp %s", name)
-		isBTFTypeFunc = false
+		target = (*btf.Typedef)(nil)
 	default:
 		return 0, errUnrecognizedAttachType
 	}
@@ -814,18 +819,7 @@ func findTargetInKernel(spec *btf.Spec, name string, progType ProgramType, attac
 		return 0, fmt.Errorf("load kernel spec: %w", err)
 	}
 
-	var target btf.Type
-	if isBTFTypeFunc {
-		var targetFunc *btf.Func
-		err = spec.TypeByName(typeName, &targetFunc)
-		target = targetFunc
-	} else {
-		var targetTypedef *btf.Typedef
-		err = spec.TypeByName(typeName, &targetTypedef)
-		target = targetTypedef
-	}
-
-	if err != nil {
+	if err := spec.TypeByName(typeName, &target); err != nil {
 		if errors.Is(err, btf.ErrNotFound) {
 			return 0, &internal.UnsupportedFeatureError{
 				Name: featureName,