Skip to content

Commit

Permalink
btf: Add spec types iterator
Browse files Browse the repository at this point in the history
This commit introduces the TypesIterator which is used to iterate over
all types of a given spec.

One notable user is "pwru" [1] which needs to iterate over all vmlinux
BTF types to find functions which accept SKB as a param.

The iterator is very similar to the existing InstructionIterator, as
both export the Next() method and set the values and indices in the
public fields.

[1]: https://github.com/cilium/pwru

Suggested-by: Timo Beckers <timo@isovalent.com>
Signed-off-by: Martynas Pumputis <m@lambda.lt>
  • Loading branch information
brb authored and lmb committed May 23, 2022
1 parent 00ae3f2 commit c4f6259
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 0 deletions.
24 changes: 24 additions & 0 deletions btf/btf.go
Original file line number Diff line number Diff line change
Expand Up @@ -633,6 +633,30 @@ func (s *Spec) TypeByName(name string, typ interface{}) error {
return nil
}

// TypesIterator iterates over types of a given spec.
type TypesIterator struct {
spec *Spec
index int
// The last visited type in the spec.
Type Type
}

// Iterate returns the types iterator.
func (s *Spec) Iterate() *TypesIterator {
return &TypesIterator{spec: s, index: 0}
}

// Next returns true as long as there are any remaining types.
func (iter *TypesIterator) Next() bool {
if len(iter.spec.types) <= iter.index {
return false
}

iter.Type = iter.spec.types[iter.index]
iter.index++
return true
}

// Handle is a reference to BTF loaded into the kernel.
type Handle struct {
spec *Spec
Expand Down
35 changes: 35 additions & 0 deletions btf/btf_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -342,3 +342,38 @@ func ExampleSpec_TypeByName() {
// We've found struct foo
fmt.Println(foo.Name)
}

func TestTypesIterator(t *testing.T) {
spec, err := LoadSpecFromReader(readVMLinux(t))
if err != nil {
t.Fatal(err)
}

if len(spec.types) < 1 {
t.Fatal("Not enough types")
}

// Assertion that 'iphdr' type exists within the spec
_, err = spec.AnyTypeByName("iphdr")
if err != nil {
t.Fatalf("Failed to find 'iphdr' type by name: %s", err)
}

found := false
count := 0

iter := spec.Iterate()
for iter.Next() {
if !found && iter.Type.TypeName() == "iphdr" {
found = true
}
count += 1
}

if l := len(spec.types); l != count {
t.Fatalf("Failed to iterate over all types (%d vs %d)", l, count)
}
if !found {
t.Fatal("Cannot find 'iphdr' type")
}
}

0 comments on commit c4f6259

Please sign in to comment.