From 5aceb9e681cd6c82a2eccc25e1452d72d991c613 Mon Sep 17 00:00:00 2001
From: Mateusz Zalega <mateusz@monogon.tech>
Date: Wed, 25 Jan 2023 11:20:06 +0000
Subject: [PATCH] Support "sample" filter action

This change adds support for packet sampling using "psample" kernel
module.
---
 filter.go       |  23 +++++++++
 filter_linux.go |  25 +++++++++
 filter_test.go  | 132 ++++++++++++++++++++++++++++++++++++++++++++++++
 nl/tc_linux.go  |  11 ++++
 4 files changed, 191 insertions(+)

diff --git a/filter.go b/filter.go
index 84e1ca7..e4f3167 100644
--- a/filter.go
+++ b/filter.go
@@ -369,6 +369,29 @@ func NewPoliceAction() *PoliceAction {
 	}
 }
 
+type SampleAction struct {
+	ActionAttrs
+	Group     uint32
+	Rate      uint32
+	TruncSize uint32
+}
+
+func (action *SampleAction) Type() string {
+	return "sample"
+}
+
+func (action *SampleAction) Attrs() *ActionAttrs {
+	return &action.ActionAttrs
+}
+
+func NewSampleAction() *SampleAction {
+	return &SampleAction{
+		ActionAttrs: ActionAttrs{
+			Action: TC_ACT_PIPE,
+		},
+	}
+}
+
 // MatchAll filters match all packets
 type MatchAll struct {
 	FilterAttrs
diff --git a/filter_linux.go b/filter_linux.go
index 1930661..d61e357 100644
--- a/filter_linux.go
+++ b/filter_linux.go
@@ -705,6 +705,17 @@ func EncodeActions(attr *nl.RtAttr, actions []Action) error {
 			aopts.AddRtAttr(nl.TCA_ACT_BPF_PARMS, gen.Serialize())
 			aopts.AddRtAttr(nl.TCA_ACT_BPF_FD, nl.Uint32Attr(uint32(action.Fd)))
 			aopts.AddRtAttr(nl.TCA_ACT_BPF_NAME, nl.ZeroTerminated(action.Name))
+		case *SampleAction:
+			table := attr.AddRtAttr(tabIndex, nil)
+			tabIndex++
+			table.AddRtAttr(nl.TCA_ACT_KIND, nl.ZeroTerminated("sample"))
+			aopts := table.AddRtAttr(nl.TCA_ACT_OPTIONS, nil)
+			gen := nl.TcGen{}
+			toTcGen(action.Attrs(), &gen)
+			aopts.AddRtAttr(nl.TCA_ACT_SAMPLE_PARMS, gen.Serialize())
+			aopts.AddRtAttr(nl.TCA_ACT_SAMPLE_RATE, nl.Uint32Attr(action.Rate))
+			aopts.AddRtAttr(nl.TCA_ACT_SAMPLE_PSAMPLE_GROUP, nl.Uint32Attr(action.Group))
+			aopts.AddRtAttr(nl.TCA_ACT_SAMPLE_TRUNC_SIZE, nl.Uint32Attr(action.TruncSize))
 		case *GenericAction:
 			table := attr.AddRtAttr(tabIndex, nil)
 			tabIndex++
@@ -790,6 +801,8 @@ func parseActions(tables []syscall.NetlinkRouteAttr) ([]Action, error) {
 					action = &ConnmarkAction{}
 				case "csum":
 					action = &CsumAction{}
+				case "sample":
+					action = &SampleAction{}
 				case "gact":
 					action = &GenericAction{}
 				case "tunnel_key":
@@ -902,6 +915,18 @@ func parseActions(tables []syscall.NetlinkRouteAttr) ([]Action, error) {
 							tcTs := nl.DeserializeTcf(adatum.Value)
 							actionTimestamp = toTimeStamp(tcTs)
 						}
+					case "sample":
+						switch adatum.Attr.Type {
+						case nl.TCA_ACT_SAMPLE_PARMS:
+							gen := *nl.DeserializeTcGen(adatum.Value)
+							toAttrs(&gen, action.Attrs())
+						case nl.TCA_ACT_SAMPLE_RATE:
+							action.(*SampleAction).Rate = native.Uint32(adatum.Value[0:4])
+						case nl.TCA_ACT_SAMPLE_PSAMPLE_GROUP:
+							action.(*SampleAction).Group = native.Uint32(adatum.Value[0:4])
+						case nl.TCA_ACT_SAMPLE_TRUNC_SIZE:
+							action.(*SampleAction).TruncSize = native.Uint32(adatum.Value[0:4])
+						}
 					case "gact":
 						switch adatum.Attr.Type {
 						case nl.TCA_GACT_PARMS:
diff --git a/filter_test.go b/filter_test.go
index 3a49f1b..774e7d6 100644
--- a/filter_test.go
+++ b/filter_test.go
@@ -2471,3 +2471,135 @@ func TestFilterChainAddDel(t *testing.T) {
 		t.Fatal("Failed to remove qdisc")
 	}
 }
+
+func TestFilterSampleAddDel(t *testing.T) {
+	minKernelRequired(t, 4, 11)
+	if _, err := GenlFamilyGet("psample"); err != nil {
+		t.Skip("psample genetlink family unavailable - is CONFIG_PSAMPLE enabled?")
+	}
+
+	tearDown := setUpNetlinkTest(t)
+	defer tearDown()
+	if err := LinkAdd(&Ifb{LinkAttrs{Name: "foo"}}); err != nil {
+		t.Fatal(err)
+	}
+	link, err := LinkByName("foo")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := LinkSetUp(link); err != nil {
+		t.Fatal(err)
+	}
+
+	qdisc := &Ingress{
+		QdiscAttrs: QdiscAttrs{
+			LinkIndex: link.Attrs().Index,
+			Handle:    MakeHandle(0xffff, 0),
+			Parent:    HANDLE_INGRESS,
+		},
+	}
+	if err := QdiscAdd(qdisc); err != nil {
+		t.Fatal(err)
+	}
+	qdiscs, err := SafeQdiscList(link)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	found := false
+	for _, v := range qdiscs {
+		if _, ok := v.(*Ingress); ok {
+			found = true
+			break
+		}
+	}
+	if !found {
+		t.Fatal("Qdisc is the wrong type")
+	}
+
+	sample := NewSampleAction()
+	sample.Group = 7
+	sample.Rate = 12
+	sample.TruncSize = 200
+
+	classId := MakeHandle(1, 1)
+	filter := &MatchAll{
+		FilterAttrs: FilterAttrs{
+			LinkIndex: link.Attrs().Index,
+			Parent:    MakeHandle(0xffff, 0),
+			Priority:  1,
+			Protocol:  unix.ETH_P_ALL,
+		},
+		ClassId: classId,
+		Actions: []Action{
+			sample,
+		},
+	}
+
+	if err := FilterAdd(filter); err != nil {
+		t.Fatal(err)
+	}
+
+	filters, err := FilterList(link, MakeHandle(0xffff, 0))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(filters) != 1 {
+		t.Fatal("Failed to add filter")
+	}
+	mf, ok := filters[0].(*MatchAll)
+	if !ok {
+		t.Fatal("Filter is the wrong type")
+	}
+
+	if len(mf.Actions) < 1 {
+		t.Fatalf("Too few Actions in filter")
+	}
+	if mf.ClassId != classId {
+		t.Fatalf("ClassId of the filter is the wrong value")
+	}
+
+	lsample, ok := mf.Actions[0].(*SampleAction)
+	if !ok {
+		t.Fatal("Unable to find sample action")
+	}
+	if lsample.Group != sample.Group {
+		t.Fatalf("Inconsistent sample action group")
+	}
+	if lsample.Rate != sample.Rate {
+		t.Fatalf("Inconsistent sample action rate")
+	}
+	if lsample.TruncSize != sample.TruncSize {
+		t.Fatalf("Inconsistent sample truncation size")
+	}
+
+	if err := FilterDel(filter); err != nil {
+		t.Fatal(err)
+	}
+	filters, err = FilterList(link, MakeHandle(0xffff, 0))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(filters) != 0 {
+		t.Fatal("Failed to remove filter")
+	}
+
+	if err := QdiscDel(qdisc); err != nil {
+		t.Fatal(err)
+	}
+	qdiscs, err = SafeQdiscList(link)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	found = false
+	for _, v := range qdiscs {
+		if _, ok := v.(*Ingress); ok {
+			found = true
+			break
+		}
+	}
+	if found {
+		t.Fatal("Failed to remove qdisc")
+	}
+}
diff --git a/nl/tc_linux.go b/nl/tc_linux.go
index 0720729..db3ca1c 100644
--- a/nl/tc_linux.go
+++ b/nl/tc_linux.go
@@ -77,6 +77,17 @@ const (
 	TCA_ACT_MAX
 )
 
+const (
+	TCA_ACT_SAMPLE_UNSPEC = iota
+	TCA_ACT_SAMPLE_TM
+	TCA_ACT_SAMPLE_PARMS
+	TCA_ACT_SAMPLE_RATE
+	TCA_ACT_SAMPLE_TRUNC_SIZE
+	TCA_ACT_SAMPLE_PSAMPLE_GROUP
+	TCA_ACT_SAMPLE_PAD
+	TCA_ACT_SAMPLE_MAX
+)
+
 const (
 	TCA_PRIO_UNSPEC = iota
 	TCA_PRIO_MQ
-- 
2.47.0