link: add Iterator

Add an iterator for links that loops over all links in the system.

Signed-off-by: Charalampos (Babis) Stylianopoulos <charalampos.stylianopoulos@emnify.com>

internal/cmd/gentypes/main.go

@@ -363,6 +363,14 @@ import (
 				truncateAfter("next_id"),
 			},
 		},
+		{
+			"LinkGetNextId", retError, "obj_next_id", "BPF_LINK_GET_NEXT_ID",
+			[]patch{
+				choose(0, "start_id"), rename("start_id", "id"),
+				replace(linkID, "id", "next_id"),
+				truncateAfter("next_id"),
+			},
+		},
 		// These piggy back on the obj_next_id decl, but only support the
 		// first field...
 		{
@@ -377,6 +385,10 @@ import (
 			"ProgGetFdById", retFd, "obj_next_id", "BPF_PROG_GET_FD_BY_ID",
 			[]patch{choose(0, "start_id"), rename("start_id", "id"), truncateAfter("id")},
 		},
+		{
+			"LinkGetFdById", retFd, "obj_next_id", "BPF_LINK_GET_FD_BY_ID",
+			[]patch{choose(0, "start_id"), rename("start_id", "id"), replace(linkID, "id"), truncateAfter("id")},
+		},
 		{
 			"ObjGetInfoByFd", retError, "info_by_fd", "BPF_OBJ_GET_INFO_BY_FD",
 			[]patch{replace(pointer, "info")},

internal/sys/types.go

@@ -837,6 +837,26 @@ func LinkCreateUprobeMulti(attr *LinkCreateUprobeMultiAttr) (*FD, error) {
 	return NewFD(int(fd))
 }
 
+type LinkGetFdByIdAttr struct{ Id LinkID }
+
+func LinkGetFdById(attr *LinkGetFdByIdAttr) (*FD, error) {
+	fd, err := BPF(BPF_LINK_GET_FD_BY_ID, unsafe.Pointer(attr), unsafe.Sizeof(*attr))
+	if err != nil {
+		return nil, err
+	}
+	return NewFD(int(fd))
+}
+
+type LinkGetNextIdAttr struct {
+	Id     LinkID
+	NextId LinkID
+}
+
+func LinkGetNextId(attr *LinkGetNextIdAttr) error {
+	_, err := BPF(BPF_LINK_GET_NEXT_ID, unsafe.Pointer(attr), unsafe.Sizeof(*attr))
+	return err
+}
+
 type LinkUpdateAttr struct {
 	LinkFd    uint32
 	NewProgFd uint32

link/link.go

@@ -1,7 +1,9 @@
 package link
 
 import (
+	"errors"
 	"fmt"
+	"os"
 
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/btf"
@@ -46,8 +48,15 @@ type Link interface {
 
 // NewLinkFromFD creates a link from a raw fd.
 //
-// You should not use fd after calling this function.
+// Deprecated: use [NewFromFD] instead.
 func NewLinkFromFD(fd int) (Link, error) {
+	return NewFromFD(fd)
+}
+
+// NewFromFD creates a link from a raw fd.
+//
+// You should not use fd after calling this function.
+func NewFromFD(fd int) (Link, error) {
 	sysFD, err := sys.NewFD(fd)
 	if err != nil {
 		return nil, err
@@ -56,6 +65,19 @@ func NewLinkFromFD(fd int) (Link, error) {
 	return wrapRawLink(&RawLink{fd: sysFD})
 }
 
+// NewFromID returns the link associated with the given id.
+//
+// Returns ErrNotExist if there is no link with the given id.
+func NewFromID(id ID) (Link, error) {
+	getFdAttr := &sys.LinkGetFdByIdAttr{Id: id}
+	fd, err := sys.LinkGetFdById(getFdAttr)
+	if err != nil {
+		return nil, fmt.Errorf("get link fd from ID %d: %w", id, err)
+	}
+
+	return wrapRawLink(&RawLink{fd, ""})
+}
+
 // LoadPinnedLink loads a link that was persisted into a bpffs.
 func LoadPinnedLink(fileName string, opts *ebpf.LoadPinOptions) (Link, error) {
 	raw, err := loadPinnedRawLink(fileName, opts)
@@ -446,3 +468,74 @@ func (l *RawLink) Info() (*Info, error) {
 		extra,
 	}, nil
 }
+
+// Iterator allows iterating over links attached into the kernel.
+type Iterator struct {
+	// The ID of the current link. Only valid after a call to Next
+	ID ID
+	// The current link. Only valid until a call to Next.
+	// See Take if you want to retain the link.
+	Link Link
+	err  error
+}
+
+// Next retrieves the next link.
+//
+// Returns true if another link was found. Call [Iterator.Err] after the function returns false.
+func (it *Iterator) Next() bool {
+	id := it.ID
+	for {
+		getIdAttr := &sys.LinkGetNextIdAttr{Id: id}
+		err := sys.LinkGetNextId(getIdAttr)
+		if errors.Is(err, os.ErrNotExist) {
+			// There are no more links.
+			break
+		} else if err != nil {
+			it.err = fmt.Errorf("get next link ID: %w", err)
+			break
+		}
+
+		id = getIdAttr.NextId
+		l, err := NewFromID(id)
+		if errors.Is(err, os.ErrNotExist) {
+			// Couldn't load the link fast enough. Try next ID.
+			continue
+		} else if err != nil {
+			it.err = fmt.Errorf("get link for ID %d: %w", id, err)
+			break
+		}
+
+		if it.Link != nil {
+			it.Link.Close()
+		}
+		it.ID, it.Link = id, l
+		return true
+	}
+
+	// No more links or we encountered an error.
+	if it.Link != nil {
+		it.Link.Close()
+	}
+	it.Link = nil
+	return false
+}
+
+// Take the ownership of the current link.
+//
+// It's the callers responsibility to close the link.
+func (it *Iterator) Take() Link {
+	l := it.Link
+	it.Link = nil
+	return l
+}
+
+// Err returns an error if iteration failed for some reason.
+func (it *Iterator) Err() error {
+	return it.err
+}
+
+func (it *Iterator) Close() {
+	if it.Link != nil {
+		it.Link.Close()
+	}
+}

link/link_test.go

@@ -88,6 +88,67 @@ func TestRawLinkLoadPinnedWithOptions(t *testing.T) {
 	}
 }
 
+func TestIterator(t *testing.T) {
+	cgroup, prog := mustCgroupFixtures(t)
+
+	tLink, err := AttachRawLink(RawLinkOptions{
+		Target:  int(cgroup.Fd()),
+		Program: prog,
+		Attach:  ebpf.AttachCGroupInetEgress,
+	})
+	testutils.SkipIfNotSupported(t, err)
+	if err != nil {
+		t.Fatal("Can't create original raw link:", err)
+	}
+	defer tLink.Close()
+	tLinkInfo, err := tLink.Info()
+	testutils.SkipIfNotSupported(t, err)
+	if err != nil {
+		t.Fatal("Can't get original link info:", err)
+	}
+
+	it := new(Iterator)
+	defer it.Close()
+
+	prev := it.ID
+	var foundLink Link
+	for it.Next() {
+		// Iterate all loaded links.
+		if it.Link == nil {
+			t.Fatal("Next doesn't assign link")
+		}
+		if it.ID == prev {
+			t.Fatal("Iterator doesn't advance ID")
+		}
+		prev = it.ID
+		if it.ID == tLinkInfo.ID {
+			foundLink = it.Take()
+		}
+	}
+	if err := it.Err(); err != nil {
+		t.Fatal("Iteration returned an error:", err)
+	}
+	if it.Link != nil {
+		t.Fatal("Next doesn't clean up link on last iteration")
+	}
+	if prev != it.ID {
+		t.Fatal("Next changes ID on last iteration")
+	}
+	if foundLink == nil {
+		t.Fatal("Original link not found")
+	}
+	defer foundLink.Close()
+	// Confirm that we found the original link.
+	info, err := foundLink.Info()
+	if err != nil {
+		t.Fatal("Can't get link info:", err)
+	}
+	if info.ID != tLinkInfo.ID {
+		t.Fatal("Found link has wrong ID")
+	}
+
+}
+
 func newPinnedRawLink(t *testing.T, cgroup *os.File, prog *ebpf.Program) (*RawLink, string) {
 	t.Helper()
 
@@ -235,7 +296,7 @@ func testLink(t *testing.T, link Link, prog *ebpf.Program) {
 		}
 		defer unix.Close(dupFD)
 
-		newLink, err := NewLinkFromFD(dupFD)
+		newLink, err := NewFromFD(dupFD)
 		testutils.SkipIfNotSupported(t, err)
 		if err != nil {
 			t.Fatal("Can't create new link from dup link FD:", err)