From f6e716d3ff1bb29519e28cc693a490a2a0d6dbd8 Mon Sep 17 00:00:00 2001
From: Lorenz Bauer <i@lmb.io>
Date: Fri, 18 Nov 2022 13:23:44 +0000
Subject: [PATCH] btf: fix Spec.TypeByID for split BTF

TypeByID currently assumes that the first ID in s.types is 0, which
is not correct when dealing with split BTF.
---
 btf/btf.go      | 9 ++++++++-
 btf/btf_test.go | 5 +++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/btf/btf.go b/btf/btf.go
index 539baff11..ef131bf15 100644
--- a/btf/btf.go
+++ b/btf/btf.go
@@ -554,7 +554,14 @@ func (s *Spec) Add(typ Type) (TypeID, error) {
 // Returns an error wrapping ErrNotFound if a Type with the given ID
 // does not exist in the Spec.
 func (s *Spec) TypeByID(id TypeID) (Type, error) {
-	return s.types.ByID(id)
+	firstID := s.firstTypeID()
+	i := int(id - firstID)
+
+	if id < firstID || i >= len(s.types) {
+		return nil, fmt.Errorf("type ID %d: %w", id, ErrNotFound)
+	}
+
+	return s.types[i], nil
 }
 
 // TypeID returns the ID for a given Type.
diff --git a/btf/btf_test.go b/btf/btf_test.go
index a662e5aab..29eb7f459 100644
--- a/btf/btf_test.go
+++ b/btf/btf_test.go
@@ -464,6 +464,11 @@ func TestLoadSplitSpecFromReader(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+
+	typeByID, err := splitSpec.TypeByID(typeID)
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, typeByID, qt.Equals, typ)
+
 	fnType := typ.(*Func)
 	fnProto := fnType.Type.(*FuncProto)