From 961fee4859d60f8f0164b7bce965858b5aa7c5ba Mon Sep 17 00:00:00 2001
From: Leon Hwang <hffilwlqm@gmail.com>
Date: Wed, 18 Dec 2024 20:51:29 +0800
Subject: [PATCH] disasm: Remove -B requirement

If `-B` is not provided, it will end early if meet `retq` or `int3` or
limit to 4096.

Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>
---
 internal/bpflbr/disasm.go | 30 +++++++++++++++++++++++-------
 internal/bpflbr/flags.go  |  2 +-
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/internal/bpflbr/disasm.go b/internal/bpflbr/disasm.go
index e640aa3..3bbbcde 100644
--- a/internal/bpflbr/disasm.go
+++ b/internal/bpflbr/disasm.go
@@ -7,6 +7,7 @@ import (
 	"debug/elf"
 	"errors"
 	"fmt"
+	"log"
 	"os"
 	"strconv"
 	"strings"
@@ -33,7 +34,6 @@ func Disasm(f *Flags) {
 
 	if len(f.kfuncs) != 0 {
 		assert.SliceLen(f.kfuncs, 1, "Only one --kfunc is allowed for --disasm")
-		assert.True(f.disasmBytes != 0, "--disasm-bytes must be set for --disasm to disasm %v", f.kfuncs)
 
 		dumpKfunc(f.kfuncs[0], f.disasmBytes)
 		return
@@ -48,9 +48,13 @@ func readKcore(kaddr uint64, bytes uint) ([]byte, bool) {
 	kcoreElf, err := elf.NewFile(fd)
 	assert.NoErr(err, "Failed to read %s: %v", kcorePath)
 
-	data := make([]byte, bytes)
 	for _, prog := range kcoreElf.Progs {
 		if prog.Vaddr <= kaddr && kaddr < prog.Vaddr+prog.Memsz {
+			remain := uint(prog.Memsz + prog.Vaddr - kaddr)
+			if bytes == 0 {
+				bytes = 4096 // limit to 4KiB
+			}
+			data := make([]byte, min(bytes, remain))
 			n, err := fd.ReadAt(data, int64(prog.Off+kaddr-prog.Vaddr))
 			assert.NoErr(err, "Failed to read %s: %v", kcorePath)
 			return data[:n], true
@@ -136,12 +140,16 @@ func dumpKfunc(kfunc string, bytes uint) {
 
 	b, pc := data[:], uint64(kaddr)
 	for len(b) != 0 {
-		inst, err := engine.Disasm(b, pc, 1)
+		insts, err := engine.Disasm(b, pc, 1)
 		if err != nil && len(b) <= 10 {
 			break
 		}
 		if err != nil {
 			fmt.Print(sb.String())
+			if errors.Is(err, gapstone.ErrOK) {
+				log.Println("Finish disassembling early, pls try again")
+				break
+			}
 			assert.NoErr(err, "Failed to disasm: %v")
 		}
 
@@ -159,13 +167,15 @@ func dumpKfunc(kfunc string, bytes uint) {
 			prev = li
 		}
 
+		inst := insts[0]
+
 		var opcodes []string
-		for _, insn := range inst[0].Bytes {
+		for _, insn := range inst.Bytes {
 			opcodes = append(opcodes, fmt.Sprintf("%02x", insn))
 		}
 		opcode := strings.Join(opcodes, " ")
-		opstr := inst[0].OpStr
-		fmt.Fprintf(&sb, "%#x: %-19s\t%s\t%s", pc, opcode, inst[0].Mnemonic, opstr)
+		opstr := inst.OpStr
+		fmt.Fprintf(&sb, "%#x: %-19s\t%s\t%s", pc, opcode, inst.Mnemonic, opstr)
 
 		var endpoint *branchEndpoint
 		if strings.HasPrefix(opstr, "0x") {
@@ -185,7 +195,13 @@ func dumpKfunc(kfunc string, bytes uint) {
 		}
 		fmt.Fprintln(&sb)
 
-		insnSize := uint64(inst[0].Size)
+		if bytes == 0 && len(inst.Bytes) == 1 &&
+			(inst.Bytes[0] == 0xc3 /* retq */ ||
+				inst.Bytes[0] == 0xcc /* int3 */) {
+			break
+		}
+
+		insnSize := uint64(inst.Size)
 		pc += insnSize
 		b = b[insnSize:]
 	}
diff --git a/internal/bpflbr/flags.go b/internal/bpflbr/flags.go
index 7be6fd8..503520c 100644
--- a/internal/bpflbr/flags.go
+++ b/internal/bpflbr/flags.go
@@ -141,7 +141,7 @@ func ParseFlags() (*Flags, error) {
 	f.BoolVar(&kfuncAllKmods, "kfunc-all-kmods", false, "filter functions in all kernel modules")
 	f.StringVarP(&flags.outputFile, "output", "o", "", "output file for the result, default is stdout")
 	f.BoolVarP(&flags.disasm, "disasm", "d", false, "disasm bpf prog or kernel function")
-	f.UintVarP(&flags.disasmBytes, "disasm-bytes", "B", 0, "disasm bytes of kernel function, must not 0")
+	f.UintVarP(&flags.disasmBytes, "disasm-bytes", "B", 0, "disasm bytes of kernel function, end at the very first retq/int3 insn or limit to 4096 if not provided")
 	f.BoolVar(&disasmIntelSyntax, "disasm-intel-syntax", false, "use Intel asm syntax for disasm, ATT asm syntax by default")
 	f.BoolVarP(&verbose, "verbose", "v", false, "output verbose log")
 	f.StringVarP(&mode, "mode", "m", TracingModeExit, "mode of lbr tracing, exit or entry")