From c105dbbeb751825a3e2f0322218a01937bd69983 Mon Sep 17 00:00:00 2001
From: Leon Hwang <hffilwlqm@gmail.com>
Date: Sat, 9 Nov 2024 13:10:30 +0800
Subject: [PATCH] bpflbr: Enable LBR perf event

Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>
---
 internal/bpflbr/lbr_perf_event.go       | 49 +++++++++++++++++++++++++
 internal/bpflbr/lbr_perf_event_linux.go | 26 +++++++++++++
 internal/bpflbr/lbr_perf_event_other.go | 14 +++++++
 main.go                                 |  9 +++++
 4 files changed, 98 insertions(+)
 create mode 100644 internal/bpflbr/lbr_perf_event.go
 create mode 100644 internal/bpflbr/lbr_perf_event_linux.go
 create mode 100644 internal/bpflbr/lbr_perf_event_other.go

diff --git a/internal/bpflbr/lbr_perf_event.go b/internal/bpflbr/lbr_perf_event.go
new file mode 100644
index 0000000..2e8bb24
--- /dev/null
+++ b/internal/bpflbr/lbr_perf_event.go
@@ -0,0 +1,49 @@
+// Copyright 2024 Leon Hwang.
+// SPDX-License-Identifier: Apache-2.0
+
+package bpflbr
+
+import (
+	"fmt"
+
+	"github.com/cilium/ebpf"
+	"golang.org/x/sys/unix"
+)
+
+type LbrPerfEvent struct {
+	fds []int
+}
+
+func OpenLbrPerfEvent() (*LbrPerfEvent, error) {
+	var p LbrPerfEvent
+	var err error
+
+	defer func() {
+		if err != nil {
+			p.Close()
+		}
+	}()
+
+	numCPU, err := ebpf.PossibleCPU()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get number of CPUs: %w", err)
+	}
+
+	p.fds = make([]int, 0, numCPU)
+	for i := 0; i < numCPU; i++ {
+		fd, err := openLbrPerfEvent(i)
+		if err != nil {
+			return nil, fmt.Errorf("failed to open LBR perf event: %w", err)
+		}
+
+		p.fds = append(p.fds, fd)
+	}
+
+	return &p, nil
+}
+
+func (p *LbrPerfEvent) Close() {
+	for _, fd := range p.fds {
+		_ = unix.Close(fd)
+	}
+}
diff --git a/internal/bpflbr/lbr_perf_event_linux.go b/internal/bpflbr/lbr_perf_event_linux.go
new file mode 100644
index 0000000..5e78235
--- /dev/null
+++ b/internal/bpflbr/lbr_perf_event_linux.go
@@ -0,0 +1,26 @@
+// Copyright 2024 Leon Hwang.
+// SPDX-License-Identifier: Apache-2.0
+
+//go:build linux
+
+package bpflbr
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+)
+
+func openLbrPerfEvent(cpu int) (int, error) {
+	var attr unix.PerfEventAttr
+	attr.Size = uint32(unsafe.Sizeof(attr))
+	attr.Type = unix.PERF_TYPE_HARDWARE
+	attr.Config = unix.PERF_COUNT_HW_CPU_CYCLES
+	attr.Sample = 4000
+	attr.Bits |= unix.PerfBitFreq
+	attr.Sample_type = unix.PERF_SAMPLE_BRANCH_STACK
+	attr.Branch_sample_type = unix.PERF_SAMPLE_BRANCH_KERNEL |
+		unix.PERF_SAMPLE_BRANCH_ANY
+
+	return unix.PerfEventOpen(&attr, -1, cpu, -1, unix.PERF_FLAG_FD_CLOEXEC)
+}
diff --git a/internal/bpflbr/lbr_perf_event_other.go b/internal/bpflbr/lbr_perf_event_other.go
new file mode 100644
index 0000000..5ed0713
--- /dev/null
+++ b/internal/bpflbr/lbr_perf_event_other.go
@@ -0,0 +1,14 @@
+// Copyright 2024 Leon Hwang.
+// SPDX-License-Identifier: Apache-2.0
+
+//go:build !linux
+
+package bpflbr
+
+import (
+	"golang.org/x/sys/unix"
+)
+
+func openLbrPerfEvent(cpu int) (int, error) {
+	return 0, unix.EOPNOTSUPP
+}
diff --git a/main.go b/main.go
index 97f8b5d..d11eb8e 100644
--- a/main.go
+++ b/main.go
@@ -5,6 +5,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"log"
 	"os"
 	"os/signal"
@@ -14,6 +15,7 @@ import (
 	"github.com/cilium/ebpf/ringbuf"
 	"github.com/knightsc/gapstone"
 	"golang.org/x/sync/errgroup"
+	"golang.org/x/sys/unix"
 
 	"github.com/Asphaltt/bpflbr/internal/assert"
 	"github.com/Asphaltt/bpflbr/internal/bpflbr"
@@ -41,6 +43,13 @@ func main() {
 	err = bpflbr.DetectBPFFeatures(featBPFSpec)
 	assert.NoErr(err, "Failed to detect bpf features: %v")
 
+	lbrPerfEvents, err := bpflbr.OpenLbrPerfEvent()
+	if err != nil && errors.Is(err, unix.ENOENT) {
+		log.Fatalln("LBR is not supported on current system")
+	}
+	assert.NoErr(err, "Failed to open LBR perf event: %v")
+	defer lbrPerfEvents.Close()
+
 	kallsyms, err := bpflbr.NewKallsyms()
 	assert.NoErr(err, "Failed to read /proc/kallsyms: %v")