[RFC] pwru: Support --output-skb-metadata

.github/workflows/release.yaml

@@ -18,7 +18,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@3041bf56c941b39c61721a86cd11f3bb1338122a
         with:
-          go-version: 1.23.1
+          go-version: 1.24.1
 
       - name: Generate the artifacts
         run: make release

.github/workflows/test.yml

@@ -16,7 +16,7 @@ jobs:
     - name: Set up Go
       uses: actions/setup-go@3041bf56c941b39c61721a86cd11f3bb1338122a
       with:
-        go-version: 1.23.1
+        go-version: 1.24.1
 
     - name: Check module vendoring
       run: |

Dockerfile

@@ -1,4 +1,4 @@
-FROM --platform=$BUILDPLATFORM docker.io/library/golang:1.23.1 AS build
+FROM --platform=$BUILDPLATFORM docker.io/library/golang:1.24.1 AS build
 ARG TARGETPLATFORM
 ARG BUILDPLATFORM
 ARG TARGETARCH

bpf/kprobe_pwru.c

@@ -89,6 +89,7 @@ struct event_t {
 	u64 param_second;
 	u64 param_third;
 	u32 cpu_id;
+	u64 skb_metadata[4]; /* output up-to-4 metadata */
 } __attribute__((packed));
 
 #define MAX_QUEUE_ENTRIES 10000
@@ -146,7 +147,7 @@ struct config {
 	u8 output_stack: 1;
 	u8 output_caller: 1;
 	u8 output_cb: 1;
-	u8 output_unused: 1;
+	u8 output_skb_metadata: 1;
 	u8 is_set: 1;
 	u8 track_skb: 1;
 	u8 track_skb_by_stackid: 1;
@@ -271,6 +272,12 @@ filter(struct sk_buff *skb) {
 	return filter_pcap(skb) && filter_meta(skb);
 }
 
+static __noinline void
+set_skb_metadata(struct sk_buff *skb, u64 *metadata) {
+	/* This func will be rewrote by Go. */
+	metadata[0] = (u64)(void *) skb;
+}
+
 static __always_inline void
 set_meta(struct sk_buff *skb, struct skb_meta *meta) {
 	meta->netns = get_netns(skb);
@@ -452,6 +459,12 @@ set_output(void *ctx, struct sk_buff *skb, struct event_t *event) {
 		struct qdisc_skb_cb *cb = (struct qdisc_skb_cb *)&skb->cb;
 		bpf_probe_read_kernel(&event->meta.cb, sizeof(event->meta.cb), (void *)&cb->data);
 	}
+
+	if (cfg->output_skb_metadata) {
+		u64 data[4];
+		set_skb_metadata(skb, data);
+		__builtin_memcpy(event->skb_metadata, data, sizeof(data));
+	}
 }
 
 static __noinline bool

go.mod

@@ -1,12 +1,16 @@
 module github.com/cilium/pwru
 
-go 1.23.1
+go 1.23.4
+
+toolchain go1.24.0
 
 require (
+	github.com/Asphaltt/mybtf v0.0.0-20250315135407-f9d09086616b
 	github.com/cheggaaa/pb/v3 v3.1.6
 	github.com/cilium/ebpf v0.17.3
 	github.com/cloudflare/cbpfc v0.0.0-20221017140110-11acb56438a2
 	github.com/jsimonetti/rtnetlink v1.4.2
+	github.com/leonhwangprojects/bice v0.1.1
 	github.com/spf13/pflag v1.0.6
 	github.com/tklauser/ps v0.0.3
 	github.com/vishvananda/netns v0.0.5
@@ -28,4 +32,5 @@ require (
 	github.com/mdlayher/socket v0.4.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/rivo/uniseg v0.2.0 // indirect
+	rsc.io/c2go v0.0.0-20170620140410-520c22818a08 // indirect
 )

go.sum

@@ -1,3 +1,5 @@
+github.com/Asphaltt/mybtf v0.0.0-20250315135407-f9d09086616b h1:LRS0ckDlNnX2Bux8k0HqLVf/2PqvvykoB8HSxXf7XjA=
+github.com/Asphaltt/mybtf v0.0.0-20250315135407-f9d09086616b/go.mod h1:ZxswciFofS8TXGsc2j5CWCY8O3XwKEysCnU7hoI1fnI=
 github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow=
 github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
 github.com/cheggaaa/pb/v3 v3.1.6 h1:h0x+vd7EiUohAJ29DJtJy+SNAc55t/elW3jCD086EXk=
@@ -30,6 +32,8 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/leonhwangprojects/bice v0.1.1 h1:gHJebJcWS8xRxJtiv3ko+HI34PRtlBpY5QlwFqgLLYg=
+github.com/leonhwangprojects/bice v0.1.1/go.mod h1:6Lsun8A3o0pPfjuCfpUR71WBuOLO64oQhXWdHdlW4SE=
 github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
 github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
@@ -73,3 +77,5 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
+rsc.io/c2go v0.0.0-20170620140410-520c22818a08 h1:AAIN5uzUq20OU2cNoPTVljUm7JDaD0Z/+Xl/uMBHJgU=
+rsc.io/c2go v0.0.0-20170620140410-520c22818a08/go.mod h1:hE73EMCiJ3lIDPgEDEK781LZrVhFDxT+I+ziQRoLVNE=

internal/pwru/config.go

@@ -23,6 +23,7 @@ const (
 	OutputStackMask
 	OutputCallerMask
 	OutputCbMask
+	OutputSkbMetadataMask
 )
 
 const (
@@ -72,6 +73,9 @@ func GetConfig(flags *Flags) (cfg FilterCfg, err error) {
 	if flags.OutputCaller {
 		cfg.OutputFlags |= OutputCallerMask
 	}
+	if len(flags.OutputSkbMetadata) > 0 {
+		cfg.OutputFlags |= OutputSkbMetadataMask
+	}
 	if flags.FilterTraceTc || flags.OutputSkbCB {
 		cfg.OutputFlags |= OutputCbMask
 	}

internal/pwru/output.go

@@ -44,6 +44,7 @@ type output struct {
 	printShinfoMap *ebpf.Map
 	printStackMap  *ebpf.Map
 	addr2name      Addr2Name
+	skbMetadata    []*SkbMetadata
 	writer         *os.File
 	kprobeMulti    bool
 	kfreeReasons   map[uint64]string
@@ -89,7 +90,7 @@ func centerAlignString(s string, width int) string {
 	return fmt.Sprintf("%s%s%s", strings.Repeat(" ", leftPadding), s, strings.Repeat(" ", rightPadding))
 }
 
-func NewOutput(flags *Flags, printSkbMap, printShinfoMap, printStackMap *ebpf.Map, addr2Name Addr2Name, kprobeMulti bool, btfSpec *btf.Spec) (*output, error) {
+func NewOutput(flags *Flags, printSkbMap, printShinfoMap, printStackMap *ebpf.Map, addr2Name Addr2Name, mds []*SkbMetadata, kprobeMulti bool, btfSpec *btf.Spec) (*output, error) {
 	writer := os.Stdout
 
 	if flags.OutputFile != "" {
@@ -120,6 +121,7 @@ func NewOutput(flags *Flags, printSkbMap, printShinfoMap, printStackMap *ebpf.Ma
 		printShinfoMap: printShinfoMap,
 		printStackMap:  printStackMap,
 		addr2name:      addr2Name,
+		skbMetadata:    mds,
 		writer:         writer,
 		kprobeMulti:    kprobeMulti,
 		kfreeReasons:   reasons,
@@ -450,6 +452,10 @@ func (o *output) Print(event *Event) {
 		fmt.Fprintf(o.writer, " %s", outFuncName)
 	}
 
+	if len(o.skbMetadata) != 0 {
+		outputSkbMetadata(o.writer, o.skbMetadata, event.SkbMetadata[:])
+	}
+
 	if o.flags.OutputStack && event.PrintStackId > 0 {
 		fmt.Fprintf(o.writer, "%s", getStackData(event, o))
 	}

internal/pwru/skb_metadata.go

@@ -0,0 +1,191 @@
+// SPDX-License-Identifier: Apache-2.0
+/* Copyright Leon Hwang */
+/* Copyright Authors of Cilium */
+
+package pwru
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"strings"
+
+	"github.com/Asphaltt/mybtf"
+	"github.com/cilium/ebpf"
+	"github.com/cilium/ebpf/asm"
+	"github.com/cilium/ebpf/btf"
+	"github.com/leonhwangprojects/bice"
+)
+
+const labelSetSkbMetadataExit = "__set_skb_metadata_exit"
+
+const setSkbMetadataStub = "set_skb_metadata"
+
+const maxSkbMetadata = 4
+
+type SkbMetadata struct {
+	expr string
+	last string
+	t    btf.Type // type of last field
+	insn asm.Instructions
+}
+
+func isChar(c byte) bool {
+	return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
+}
+
+func isDigit(c byte) bool {
+	return c >= '0' && c <= '9'
+}
+
+func isValidChar(c byte) bool {
+	return isChar(c) || c == '_' || isDigit(c)
+}
+
+func newSkbMetadata(expr string, spec *btf.Spec) (*SkbMetadata, error) {
+	var md SkbMetadata
+	md.expr = expr
+
+	if !strings.HasPrefix(expr, "skb") {
+		return nil, errors.New("--output-skb-metadata must starts with skb")
+	}
+
+	for i := len(expr) - 1; i >= 0; i-- {
+		if !isValidChar(expr[i]) {
+			md.last = expr[i+1:]
+			break
+		}
+	}
+
+	types, err := spec.AnyTypesByName("sk_buff")
+	if err != nil {
+		return nil, err
+	}
+
+	skb := types[0]
+	ptr := &btf.Pointer{Target: skb}
+
+	res, err := bice.Access(bice.AccessOptions{
+		Expr:      expr,
+		Type:      ptr,
+		Src:       asm.R1,
+		Dst:       asm.R3,
+		LabelExit: labelSetSkbMetadataExit,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse expr %s: %w", expr, err)
+	}
+
+	size, err := btf.Sizeof(res.LastField)
+	if err != nil {
+		return nil, err
+	}
+	if size > 8 {
+		return nil, fmt.Errorf("skb metadata field '%s' (type %v) is too large, max 8 bytes", md.last, res.LastField)
+	}
+
+	md.t = res.LastField
+	md.insn = res.Insns
+	return &md, nil
+}
+
+func ParseSkbMetadataExprs(exprs []string, spec *btf.Spec) ([]*SkbMetadata, error) {
+	var md []*SkbMetadata
+
+	for _, e := range exprs {
+		if e == "" {
+			continue
+		}
+
+		m, err := newSkbMetadata(e, spec)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse skb metadata expr %s: %w", e, err)
+		}
+
+		md = append(md, m)
+	}
+
+	if len(md) > maxSkbMetadata {
+		return nil, fmt.Errorf("too many skb metadata exprs, max %d", maxSkbMetadata)
+	}
+
+	return md, nil
+}
+
+func InjectSetSkbMetadata(prog *ebpf.ProgramSpec, md []*SkbMetadata) error {
+	if len(md) == 0 {
+		return nil
+	}
+
+	var insns asm.Instructions
+	if len(md) > 1 {
+		insns = append(insns,
+			asm.Mov.Reg(asm.R7, asm.R1), // r7 = skb
+		)
+	}
+	insns = append(insns,
+		asm.Mov.Reg(asm.R6, asm.R2), // r6 = metadata
+	)
+
+	insns = append(insns, md[0].insn...)
+	insns = append(insns, asm.StoreMem(asm.R6, 0, asm.R3, asm.DWord)) // metadata[0] = skb->field
+
+	offset := 0
+	for _, m := range md[1:] {
+		offset += 8
+		insns = append(insns,
+			asm.Mov.Reg(asm.R1, asm.R7), // r1 = skb
+		)
+		insns = append(insns, m.insn...)
+		insns = append(insns, asm.StoreMem(asm.R6, int16(offset), asm.R3, asm.DWord)) // metadata[i] = skb->field
+	}
+
+	insns = append(insns,
+		asm.Return().WithSymbol(labelSetSkbMetadataExit),
+	)
+
+	startIdx := -1
+	for i, ins := range prog.Instructions {
+		if ins.Symbol() == setSkbMetadataStub {
+			startIdx = i
+			break
+		}
+	}
+	if startIdx == -1 {
+		return errors.New("failed to find set_skb_metadata stub")
+	}
+
+	retIdx := -1
+	retOpCode := asm.Return().OpCode
+	for i := startIdx + 1; i < len(prog.Instructions); i++ {
+		if prog.Instructions[i].OpCode == retOpCode {
+			retIdx = i
+			break
+		}
+	}
+	if retIdx == -1 {
+		return errors.New("failed to find ret instruction")
+	}
+
+	// replace the stub's instructions
+	insns[0] = insns[0].WithMetadata(prog.Instructions[startIdx].Metadata)
+	prog.Instructions = append(prog.Instructions[:startIdx],
+		append(insns, prog.Instructions[retIdx+1:]...)...)
+
+	return nil
+}
+
+func outputSkbMetadata(w io.Writer, md []*SkbMetadata, data []uint64) {
+	for i, m := range md {
+		var b [8]byte
+		binary.NativeEndian.PutUint64(b[:], data[i])
+
+		s, err := mybtf.DumpData(m.t, b[:])
+		if err != nil {
+			fmt.Fprintf(w, " %s=..ERR..", m.last)
+		} else {
+			fmt.Fprintf(w, " %s=%s", m.last, s)
+		}
+	}
+}

internal/pwru/types.go

@@ -53,6 +53,8 @@ type Flags struct {
 	OutputJson       bool
 	OutputTCPFlags   bool
 
+	OutputSkbMetadata []string
+
 	KMods    []string
 	AllKMods bool
 
@@ -88,6 +90,7 @@ func (f *Flags) SetFlags() {
 	flag.Uint64Var(&f.OutputLimitLines, "output-limit-lines", 0, "exit the program after the number of events has been received/printed")
 	flag.BoolVar(&f.OutputSkbCB, "output-skb-cb", false, "print skb->cb")
 	flag.BoolVar(&f.OutputTCPFlags, "output-tcp-flags", false, "print TCP flags")
+	flag.StringSliceVar(&f.OutputSkbMetadata, "output-skb-metadata", nil, "print skb metadata (e.g., \"skb->mark\", \"skb->hash\")")
 
 	flag.StringVar(&f.OutputFile, "output-file", "", "write traces to file")
 
@@ -182,6 +185,7 @@ type Event struct {
 	ParamSecond   uint64
 	ParamThird    uint64
 	CPU           uint32
+	SkbMetadata   [4]uint64
 }
 
 type markFlagValue struct {