From fe7cacd1b06f7240aca74071bd632af07b596e99 Mon Sep 17 00:00:00 2001
From: Timo Beckers <timo@isovalent.com>
Date: Fri, 14 Jun 2024 11:48:05 +0200
Subject: [PATCH] map: add Map.Handle() method for retrieving a Map's type
 information

Signed-off-by: Timo Beckers <timo@isovalent.com>
---
 info.go     | 15 +++++++++++++++
 map.go      | 18 ++++++++++++++++++
 map_test.go | 29 +++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/info.go b/info.go
index f447476cb..25adbb94e 100644
--- a/info.go
+++ b/info.go
@@ -30,6 +30,8 @@ type MapInfo struct {
 	Flags      uint32
 	// Name as supplied by user space at load time. Available from 4.15.
 	Name string
+
+	btf btf.ID
 }
 
 func newMapInfoFromFd(fd *sys.FD) (*MapInfo, error) {
@@ -50,6 +52,7 @@ func newMapInfoFromFd(fd *sys.FD) (*MapInfo, error) {
 		info.MaxEntries,
 		uint32(info.MapFlags),
 		unix.ByteSliceToString(info.Name[:]),
+		btf.ID(info.BtfId),
 	}, nil
 }
 
@@ -77,6 +80,18 @@ func (mi *MapInfo) ID() (MapID, bool) {
 	return mi.id, mi.id > 0
 }
 
+// BTFID returns the BTF ID associated with the Map.
+//
+// The ID is only valid as long as the associated Map is kept alive.
+// Available from 4.18.
+//
+// The bool return value indicates whether this optional field is available and
+// populated. (The field may be available but not populated if the kernel
+// supports the field but the Map was loaded without BTF information.)
+func (mi *MapInfo) BTFID() (btf.ID, bool) {
+	return mi.btf, mi.btf > 0
+}
+
 // programStats holds statistics of a program.
 type programStats struct {
 	// Total accumulated runtime of the program ins ns.
diff --git a/map.go b/map.go
index a3d5583dc..e48412cbf 100644
--- a/map.go
+++ b/map.go
@@ -571,6 +571,24 @@ func (m *Map) Info() (*MapInfo, error) {
 	return newMapInfoFromFd(m.fd)
 }
 
+// Handle returns a reference to the Map's type information in the kernel.
+//
+// Returns ErrNotSupported if the kernel has no BTF support, or if there is no
+// BTF associated with the Map.
+func (m *Map) Handle() (*btf.Handle, error) {
+	info, err := m.Info()
+	if err != nil {
+		return nil, err
+	}
+
+	id, ok := info.BTFID()
+	if !ok {
+		return nil, fmt.Errorf("map %s: retrieve BTF ID: %w", m, ErrNotSupported)
+	}
+
+	return btf.NewHandleFromID(id)
+}
+
 // MapLookupFlags controls the behaviour of the map lookup calls.
 type MapLookupFlags uint64
 
diff --git a/map_test.go b/map_test.go
index 388598e98..fc2418e00 100644
--- a/map_test.go
+++ b/map_test.go
@@ -11,6 +11,7 @@ import (
 	"unsafe"
 
 	"github.com/cilium/ebpf/asm"
+	"github.com/cilium/ebpf/btf"
 	"github.com/cilium/ebpf/internal"
 	"github.com/cilium/ebpf/internal/sys"
 	"github.com/cilium/ebpf/internal/testutils"
@@ -1791,6 +1792,34 @@ func TestMapPinning(t *testing.T) {
 	qt.Assert(t, qt.StringContains(err.Error(), "ValueSize"))
 }
 
+func TestMapHandle(t *testing.T) {
+	testutils.SkipOnOldKernel(t, "4.18", "btf_id in map info")
+
+	kv := &btf.Int{Size: 4}
+	m, err := NewMap(&MapSpec{
+		Type:       Hash,
+		KeySize:    kv.Size,
+		ValueSize:  kv.Size,
+		Key:        kv,
+		Value:      kv,
+		MaxEntries: 1,
+	})
+	qt.Assert(t, qt.IsNil(err))
+	defer m.Close()
+
+	h, err := m.Handle()
+	qt.Assert(t, qt.IsNil(err))
+	qt.Assert(t, qt.IsNotNil(h))
+	defer h.Close()
+
+	spec, err := h.Spec(nil)
+	qt.Assert(t, qt.IsNil(err))
+
+	typ, err := spec.TypeByID(1)
+	qt.Assert(t, qt.IsNil(err))
+	qt.Assert(t, qt.ContentEquals(typ, btf.Type(kv)))
+}
+
 func TestPerfEventArrayCompatible(t *testing.T) {
 	ms := &MapSpec{
 		Type: PerfEventArray,