From 70ca0345eede55c73106738585d3a8acb39697d5 Mon Sep 17 00:00:00 2001
From: Ivan Kolodyazhny <e0ne@e0ne.info>
Date: Wed, 31 Aug 2022 22:14:03 +0300
Subject: [PATCH] Support dst_port and src_port attributes in flower filter

---
 filter_linux.go | 26 ++++++++++++++++++++++++++
 filter_test.go  | 12 ++++++++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/filter_linux.go b/filter_linux.go
index afa80ed..fbd434a 100644
--- a/filter_linux.go
+++ b/filter_linux.go
@@ -67,6 +67,8 @@ type Flower struct {
 	SkipHw        bool
 	SkipSw        bool
 	IPProto       *nl.IPProto
+	DestPort      uint16
+	SrcPort       uint16
 
 	Actions []Action
 }
@@ -135,6 +137,26 @@ func (filter *Flower) encode(parent *nl.RtAttr) error {
 	if filter.IPProto != nil {
 		ipproto := *filter.IPProto
 		parent.AddRtAttr(nl.TCA_FLOWER_KEY_IP_PROTO, ipproto.Serialize())
+		if filter.SrcPort != 0 {
+			switch ipproto {
+			case nl.IPPROTO_TCP:
+				parent.AddRtAttr(nl.TCA_FLOWER_KEY_TCP_SRC, htons(filter.SrcPort))
+			case nl.IPPROTO_UDP:
+				parent.AddRtAttr(nl.TCA_FLOWER_KEY_UDP_SRC, htons(filter.SrcPort))
+			case nl.IPPROTO_SCTP:
+				parent.AddRtAttr(nl.TCA_FLOWER_KEY_SCTP_SRC, htons(filter.SrcPort))
+			}
+		}
+		if filter.DestPort != 0 {
+			switch ipproto {
+			case nl.IPPROTO_TCP:
+				parent.AddRtAttr(nl.TCA_FLOWER_KEY_TCP_DST, htons(filter.DestPort))
+			case nl.IPPROTO_UDP:
+				parent.AddRtAttr(nl.TCA_FLOWER_KEY_UDP_DST, htons(filter.DestPort))
+			case nl.IPPROTO_SCTP:
+				parent.AddRtAttr(nl.TCA_FLOWER_KEY_SCTP_DST, htons(filter.DestPort))
+			}
+		}
 	}
 
 	var flags uint32 = 0
@@ -182,6 +204,10 @@ func (filter *Flower) decode(data []syscall.NetlinkRouteAttr) error {
 			val := new(nl.IPProto)
 			*val = nl.IPProto(datum.Value[0])
 			filter.IPProto = val
+		case nl.TCA_FLOWER_KEY_TCP_SRC, nl.TCA_FLOWER_KEY_UDP_SRC, nl.TCA_FLOWER_KEY_SCTP_SRC:
+			filter.SrcPort = ntohs(datum.Value)
+		case nl.TCA_FLOWER_KEY_TCP_DST, nl.TCA_FLOWER_KEY_UDP_DST, nl.TCA_FLOWER_KEY_SCTP_DST:
+			filter.DestPort = ntohs(datum.Value)
 		case nl.TCA_FLOWER_ACT:
 			tables, err := nl.ParseRouteAttr(datum.Value)
 			if err != nil {
diff --git a/filter_test.go b/filter_test.go
index a002af3..76ce866 100644
--- a/filter_test.go
+++ b/filter_test.go
@@ -1771,7 +1771,7 @@ func TestFilterFlowerAddDel(t *testing.T) {
 	testMask := net.CIDRMask(24, 32)
 
 	ipproto := new(nl.IPProto)
-	*ipproto = nl.IPPROTO_ICMP
+	*ipproto = nl.IPPROTO_TCP
 
 	filter := &Flower{
 		FilterAttrs: FilterAttrs{
@@ -1792,6 +1792,8 @@ func TestFilterFlowerAddDel(t *testing.T) {
 		EncDestPort:   8472,
 		EncKeyId:      1234,
 		IPProto:       ipproto,
+		DestPort:      1111,
+		SrcPort:       1111,
 		Actions: []Action{
 			&MirredAction{
 				ActionAttrs: ActionAttrs{
@@ -1860,9 +1862,15 @@ func TestFilterFlowerAddDel(t *testing.T) {
 	if filter.EncDestPort != flower.EncDestPort {
 		t.Fatalf("Flower EncDestPort doesn't match")
 	}
-	if flower.IPProto != nil || *filter.IPProto != *flower.IPProto {
+	if flower.IPProto == nil || *filter.IPProto != *flower.IPProto {
 		t.Fatalf("Flower IPProto doesn't match")
 	}
+	if filter.DestPort != flower.DestPort {
+		t.Fatalf("Flower DestPort doesn't match")
+	}
+	if filter.SrcPort != flower.SrcPort {
+		t.Fatalf("Flower SrcPort doesn't match")
+	}
 
 	mia, ok := flower.Actions[0].(*MirredAction)
 	if !ok {