From f4e6e3d5d507b436ba397d350c3dc0bce2c12e05 Mon Sep 17 00:00:00 2001
From: Christian Worm Mortensen <cworm@akamai.com>
Date: Mon, 8 Apr 2024 13:52:04 +0200
Subject: [PATCH] Allow a Police to be specified directly on a U32 filter

---
 filter_linux.go |  14 +++++++
 filter_test.go  | 102 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+)

diff --git a/filter_linux.go b/filter_linux.go
index 536e268..87cd18f 100644
--- a/filter_linux.go
+++ b/filter_linux.go
@@ -41,6 +41,7 @@ type U32 struct {
 	RedirIndex int
 	Sel        *TcU32Sel
 	Actions    []Action
+	Police     *PoliceAction
 }
 
 func (filter *U32) Attrs() *FilterAttrs {
@@ -331,6 +332,12 @@ func (h *Handle) filterModify(filter Filter, proto, flags int) error {
 		if filter.Link != 0 {
 			options.AddRtAttr(nl.TCA_U32_LINK, nl.Uint32Attr(filter.Link))
 		}
+		if filter.Police != nil {
+			police := options.AddRtAttr(nl.TCA_U32_POLICE, nil)
+			if err := encodePolice(police, filter.Police); err != nil {
+				return err
+			}
+		}
 		actionsAttr := options.AddRtAttr(nl.TCA_U32_ACT, nil)
 		// backwards compatibility
 		if filter.RedirIndex != 0 {
@@ -952,6 +959,13 @@ func parseU32Data(filter Filter, data []syscall.NetlinkRouteAttr) (bool, error)
 					u32.RedirIndex = int(action.Ifindex)
 				}
 			}
+		case nl.TCA_U32_POLICE:
+			var police PoliceAction
+			adata, _ := nl.ParseRouteAttr(datum.Value)
+			for _, aattr := range adata {
+				parsePolice(aattr, &police)
+			}
+			u32.Police = &police
 		case nl.TCA_U32_CLASSID:
 			u32.ClassId = native.Uint32(datum.Value)
 		case nl.TCA_U32_DIVISOR:
diff --git a/filter_test.go b/filter_test.go
index 64b7516..3a49f1b 100644
--- a/filter_test.go
+++ b/filter_test.go
@@ -2276,6 +2276,108 @@ func TestFilterU32PoliceAddDel(t *testing.T) {
 	}
 }
 
+func TestFilterU32DirectPoliceAddDel(t *testing.T) {
+	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)
+	}
+
+	const (
+		policeRate     = 0x40000000 // 1 Gbps
+		policeBurst    = 0x19000    // 100 KB
+		policePeakRate = 0x4000     // 16 Kbps
+	)
+
+	police := NewPoliceAction()
+	police.Rate = policeRate
+	police.PeakRate = policePeakRate
+	police.Burst = policeBurst
+	police.ExceedAction = TC_POLICE_SHOT
+	police.NotExceedAction = TC_POLICE_UNSPEC
+
+	classId := MakeHandle(1, 1)
+	filter := &U32{
+		FilterAttrs: FilterAttrs{
+			LinkIndex: link.Attrs().Index,
+			Parent:    MakeHandle(0xffff, 0),
+			Priority:  1,
+			Protocol:  unix.ETH_P_ALL,
+		},
+		ClassId: classId,
+		Police:  police,
+	}
+
+	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")
+	}
+	u32, ok := filters[0].(*U32)
+	if !ok {
+		t.Fatal("Filter is the wrong type")
+	}
+
+	if u32.Police == nil {
+		t.Fatalf("No police in filter")
+	}
+
+	if u32.Police.Rate != policeRate {
+		t.Fatal("Filter Rate doesn't match")
+	}
+
+	if u32.Police.PeakRate != policePeakRate {
+		t.Fatal("Filter PeakRate doesn't match")
+	}
+
+	if u32.Police.LinkLayer != nl.LINKLAYER_ETHERNET {
+		t.Fatal("Filter LinkLayer doesn't match")
+	}
+
+	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")
+	}
+}
+
 func TestFilterChainAddDel(t *testing.T) {
 	tearDown := setUpNetlinkTest(t)
 	defer tearDown()