bpflbr: Add --output

Signed-off-by: Leon Hwang <[email protected]>
Asphaltt · Dec 5, 2024 · 15eeae0 · 15eeae0
1 parent 03e7c09
commit 15eeae0
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 3 deletions.
diff --git a/internal/bpflbr/flags.go b/internal/bpflbr/flags.go
@@ -97,6 +97,8 @@ func parseProgsFlag(progs []string) ([]ProgFlag, error) {
 type Flags struct {
 	progs []string
 
+	outputFile string
+
 	dumpProg bool
 }
 
@@ -105,6 +107,7 @@ func ParseFlags() (*Flags, error) {
 
 	f := flag.NewFlagSet("bpflbr", flag.ExitOnError)
 	f.StringSliceVarP(&flags.progs, "prog", "p", nil, "bpf prog info for bpflbr in format PROG[,PROG,..], PROG: PROGID[:<prog function name>], PROGID: <prog ID> or 'i/id:<prog ID>' or 'p/pinned:<pinned file>' or 't/tag:<prog tag>' or 'n/name:<prog full name>'; all bpf progs will be traced by default")
+	f.StringVarP(&flags.outputFile, "output", "o", "", "output file for the result, default is stdout")
 	f.BoolVar(&flags.dumpProg, "dump-jited", false, "dump native insn info of bpf prog, the one bpf prog must be provided by --prog (its function name will be ignored)")
 
 	return &flags, f.Parse(os.Args)
@@ -114,6 +117,10 @@ func (f *Flags) ParseProgs() ([]ProgFlag, error) {
 	return parseProgsFlag(f.progs)
 }
 
+func (f *Flags) OutputFile() string {
+	return f.outputFile
+}
+
 func (f *Flags) DumpProg() bool {
 	return f.dumpProg
 }
diff --git a/internal/bpflbr/lbr.go b/internal/bpflbr/lbr.go
@@ -6,14 +6,15 @@ package bpflbr
 import (
 	"errors"
 	"fmt"
+	"io"
 	"strings"
 	"unsafe"
 
 	"github.com/cilium/ebpf/ringbuf"
 	"golang.org/x/sys/unix"
 )
 
-func Run(reader *ringbuf.Reader, progs *bpfProgs, addr2line *Addr2Line, ksyms *Kallsyms) error {
+func Run(reader *ringbuf.Reader, progs *bpfProgs, addr2line *Addr2Line, ksyms *Kallsyms, w io.Writer) error {
 	type LbrEntry struct {
 		From  uintptr
 		To    uintptr
@@ -85,7 +86,7 @@ func Run(reader *ringbuf.Reader, progs *bpfProgs, addr2line *Addr2Line, ksyms *K
 
 		fmt.Fprintf(&sb, "Recv a record for %s with retval=%d :\n", progName, event.Retval)
 		stack.output(&sb)
-		fmt.Println(sb.String())
+		fmt.Fprintln(w, sb.String())
 
 		sb.Reset()
 		stack.reset()

diff --git a/main.go b/main.go
@@ -84,6 +84,14 @@ func main() {
 	assert.NoErr(err, "Failed to create ringbuf reader: %v")
 	defer reader.Close()
 
+	w := os.Stdout
+	if flags.OutputFile() != "" {
+		f, err := os.OpenFile(flags.OutputFile(), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
+		assert.NoErr(err, "Failed to create output file: %v")
+		defer f.Close()
+		w = f
+	}
+
 	log.Print("bpflbr is running..")
 
 	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
@@ -98,7 +106,7 @@ func main() {
 	})
 
 	errg.Go(func() error {
-		return bpflbr.Run(reader, bpfProgs, addr2line, kallsyms)
+		return bpflbr.Run(reader, bpfProgs, addr2line, kallsyms, w)
 	})
 
 	err = errg.Wait()