From 0a62618e75408fc2f7aac8e716e378bb2705d4f7 Mon Sep 17 00:00:00 2001
From: Martynas Pumputis <m@lambda.lt>
Date: Fri, 18 Nov 2022 17:05:03 +0100
Subject: [PATCH] Add --output-file to log traces

Signed-off-by: Martynas Pumputis <m@lambda.lt>
---
 internal/pwru/output.go | 41 +++++++++++++++++++++++++++++------------
 internal/pwru/types.go  |  3 +++
 main.go                 |  5 ++++-
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/internal/pwru/output.go b/internal/pwru/output.go
index 2da6fb02..44092c65 100644
--- a/internal/pwru/output.go
+++ b/internal/pwru/output.go
@@ -6,7 +6,9 @@ package pwru
 
 import (
 	"fmt"
+	"io"
 	"net"
+	"os"
 	"runtime"
 	"syscall"
 
@@ -22,24 +24,38 @@ type output struct {
 	printSkbMap   *ebpf.Map
 	printStackMap *ebpf.Map
 	addr2name     Addr2Name
+	writer        io.Writer
 }
 
-func NewOutput(flags *Flags, printSkbMap *ebpf.Map, printStackMap *ebpf.Map, addr2Name Addr2Name) *output {
+func NewOutput(flags *Flags, printSkbMap *ebpf.Map, printStackMap *ebpf.Map,
+	addr2Name Addr2Name) (*output, error) {
+
+	writer := os.Stdout
+
+	if flags.OutputFile != "" {
+		file, err := os.Create(flags.OutputFile)
+		if err != nil {
+			return nil, err
+		}
+		writer = file
+	}
+
 	return &output{
 		flags:         flags,
 		lastSeenSkb:   map[uint64]uint64{},
 		printSkbMap:   printSkbMap,
 		printStackMap: printStackMap,
 		addr2name:     addr2Name,
-	}
+		writer:        writer,
+	}, nil
 }
 
 func (o *output) PrintHeader() {
-	fmt.Printf("%18s %6s %16s %24s", "SKB", "CPU", "PROCESS", "FUNC")
+	fmt.Fprintf(o.writer, "%18s %6s %16s %24s", "SKB", "CPU", "PROCESS", "FUNC")
 	if o.flags.OutputTS != "none" {
-		fmt.Printf(" %16s", "TIMESTAMP")
+		fmt.Fprintf(o.writer, " %16s", "TIMESTAMP")
 	}
-	fmt.Printf("\n")
+	fmt.Fprintf(o.writer, "\n")
 }
 
 func (o *output) Print(event *Event) {
@@ -75,18 +91,19 @@ func (o *output) Print(event *Event) {
 	} else {
 		funcName = fmt.Sprintf("0x%x", addr)
 	}
-	fmt.Printf("%18s %6s %16s %24s", fmt.Sprintf("0x%x", event.SAddr), fmt.Sprintf("%d", event.CPU), fmt.Sprintf("[%s]", execName), funcName)
+	fmt.Fprintf(o.writer, "%18s %6s %16s %24s", fmt.Sprintf("0x%x", event.SAddr),
+		fmt.Sprintf("%d", event.CPU), fmt.Sprintf("[%s]", execName), funcName)
 	if o.flags.OutputTS != "none" {
-		fmt.Printf(" %16d", ts)
+		fmt.Fprintf(o.writer, " %16d", ts)
 	}
 	o.lastSeenSkb[event.SAddr] = event.Timestamp
 
 	if o.flags.OutputMeta {
-		fmt.Printf(" netns=%d mark=0x%x ifindex=%d proto=%x mtu=%d len=%d", event.Meta.Netns, event.Meta.Mark, event.Meta.Ifindex, event.Meta.Proto, event.Meta.MTU, event.Meta.Len)
+		fmt.Fprintf(o.writer, " netns=%d mark=0x%x ifindex=%d proto=%x mtu=%d len=%d", event.Meta.Netns, event.Meta.Mark, event.Meta.Ifindex, event.Meta.Proto, event.Meta.MTU, event.Meta.Len)
 	}
 
 	if o.flags.OutputTuple {
-		fmt.Printf(" %s:%d->%s:%d(%s)",
+		fmt.Fprintf(o.writer, " %s:%d->%s:%d(%s)",
 			addrToStr(event.Tuple.L3Proto, event.Tuple.Saddr), byteorder.NetworkToHost16(event.Tuple.Sport),
 			addrToStr(event.Tuple.L3Proto, event.Tuple.Daddr), byteorder.NetworkToHost16(event.Tuple.Dport),
 			protoToStr(event.Tuple.L4Proto))
@@ -98,7 +115,7 @@ func (o *output) Print(event *Event) {
 		if err := o.printStackMap.Lookup(&id, &stack); err == nil {
 			for _, ip := range stack.IPs {
 				if ip > 0 {
-					fmt.Printf("\n%s", o.addr2name.findNearestSym(ip))
+					fmt.Fprintf(o.writer, "\n%s", o.addr2name.findNearestSym(ip))
 				}
 			}
 		}
@@ -108,11 +125,11 @@ func (o *output) Print(event *Event) {
 	if o.flags.OutputSkb {
 		id := uint32(event.PrintSkbId)
 		if str, err := o.printSkbMap.LookupBytes(&id); err == nil {
-			fmt.Printf("\n%s", string(str))
+			fmt.Fprintf(o.writer, "\n%s", string(str))
 		}
 	}
 
-	fmt.Println()
+	fmt.Fprintln(o.writer)
 }
 
 func protoToStr(proto uint8) string {
diff --git a/internal/pwru/types.go b/internal/pwru/types.go
index f7380b3c..c0f705e5 100644
--- a/internal/pwru/types.go
+++ b/internal/pwru/types.go
@@ -34,6 +34,7 @@ type Flags struct {
 	OutputSkb        bool
 	OutputStack      bool
 	OutputLimitLines uint64
+	OutputFile       string
 
 	PerCPUBuffer int
 	KMods        []string
@@ -60,6 +61,8 @@ func (f *Flags) SetFlags() {
 	flag.BoolVar(&f.OutputStack, "output-stack", false, "print stack")
 	flag.Uint64Var(&f.OutputLimitLines, "output-limit-lines", 0, "exit the program after the number of events has been received/printed")
 	flag.IntVar(&f.PerCPUBuffer, "per-cpu-buffer", os.Getpagesize(), "per CPU buffer in bytes")
+
+	flag.StringVar(&f.OutputFile, "output-file", "", "write traces to file")
 }
 
 type Tuple struct {
diff --git a/main.go b/main.go
index ded2c222..fb4739ba 100644
--- a/main.go
+++ b/main.go
@@ -207,7 +207,10 @@ func main() {
 
 	log.Println("Listening for events..")
 
-	output := pwru.NewOutput(&flags, printSkbMap, printStackMap, addr2name)
+	output, err := pwru.NewOutput(&flags, printSkbMap, printStackMap, addr2name)
+	if err != nil {
+		log.Fatalf("Failed to create outputer: %s", err)
+	}
 	output.PrintHeader()
 
 	defer func() {