From d3b7a6faddf7e9d17cea229d00692852409a3575 Mon Sep 17 00:00:00 2001
From: Steffen Vogel <post@steffenvogel.de>
Date: Sun, 9 Oct 2022 10:00:52 +0200
Subject: [PATCH] Add firewall mark option to RouteGetOptions

This option allows performing FIB lookups for a particular
firewall mark. It is equivalent to iproute2's 'ip route get mark'
option.
---
 route_linux.go |  9 +++++
 route_test.go  | 95 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+)

diff --git a/route_linux.go b/route_linux.go
index 5db8f1d..62bb2d4 100644
--- a/route_linux.go
+++ b/route_linux.go
@@ -1301,6 +1301,7 @@ type RouteGetOptions struct {
 	VrfName string
 	SrcAddr net.IP
 	UID     *uint32
+	Mark    int
 }
 
 // RouteGetWithOptions gets a route to a specific destination from the host system.
@@ -1392,8 +1393,16 @@ func (h *Handle) RouteGetWithOptions(destination net.IP, options *RouteGetOption
 			uid := *options.UID
 			b := make([]byte, 4)
 			native.PutUint32(b, uid)
+
 			req.AddData(nl.NewRtAttr(unix.RTA_UID, b))
 		}
+
+		if options.Mark > 0 {
+			b := make([]byte, 4)
+			native.PutUint32(b, uint32(options.Mark))
+
+			req.AddData(nl.NewRtAttr(unix.RTA_MARK, b))
+		}
 	}
 
 	msgs, err := req.Execute(unix.NETLINK_ROUTE, unix.RTM_NEWROUTE)
diff --git a/route_test.go b/route_test.go
index 649f8a0..cd76ce7 100644
--- a/route_test.go
+++ b/route_test.go
@@ -1748,3 +1748,98 @@ func TestRouteUIDOption(t *testing.T) {
 		t.Fatal(routes)
 	}
 }
+
+func TestRouteFWMarkOption(t *testing.T) {
+	tearDown := setUpNetlinkTest(t)
+	defer tearDown()
+
+	// setup eth0 so that network is reachable
+	err := LinkAdd(&Dummy{LinkAttrs{Name: "eth0"}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	link, err := LinkByName("eth0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err = LinkSetUp(link); err != nil {
+		t.Fatal(err)
+	}
+	addr := &Addr{
+		IPNet: &net.IPNet{
+			IP:   net.IPv4(192, 168, 1, 1),
+			Mask: net.CIDRMask(16, 32),
+		},
+	}
+	if err = AddrAdd(link, addr); err != nil {
+		t.Fatal(err)
+	}
+
+	// a table different than unix.RT_TABLE_MAIN
+	testtable := 1000
+
+	gw1 := net.IPv4(192, 168, 1, 254)
+	gw2 := net.IPv4(192, 168, 2, 254)
+
+	// add default route via gw1 (in main route table by default)
+	defaultRouteMain := Route{
+		Dst: nil,
+		Gw:  gw1,
+	}
+	if err := RouteAdd(&defaultRouteMain); err != nil {
+		t.Fatal(err)
+	}
+
+	// add default route via gw2 in test route table
+	defaultRouteTest := Route{
+		Dst:   nil,
+		Gw:    gw2,
+		Table: testtable,
+	}
+	if err := RouteAdd(&defaultRouteTest); err != nil {
+		t.Fatal(err)
+	}
+
+	// check the routes are in different tables
+	routes, err := RouteListFiltered(FAMILY_V4, &Route{
+		Dst:   nil,
+		Table: unix.RT_TABLE_UNSPEC,
+	}, RT_FILTER_DST|RT_FILTER_TABLE)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(routes) != 2 || routes[0].Table == routes[1].Table {
+		t.Fatal("Routes not added properly")
+	}
+
+	// add a rule that fwmark match should result in route lookup of test table
+	fwmark := 1000
+
+	rule := NewRule()
+	rule.Mark = fwmark
+	rule.Mask = 0xFFFFFFFF
+	rule.Table = testtable
+	if err := RuleAdd(rule); err != nil {
+		t.Fatal(err)
+	}
+
+	dstIP := net.IPv4(10, 1, 1, 1)
+
+	// check getting route without FWMark option
+	routes, err = RouteGetWithOptions(dstIP, &RouteGetOptions{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(routes) != 1 || !routes[0].Gw.Equal(gw1) {
+		t.Fatal(routes)
+	}
+
+	// check getting route with FWMark option
+	routes, err = RouteGetWithOptions(dstIP, &RouteGetOptions{Mark: fwmark})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(routes) != 1 || !routes[0].Gw.Equal(gw2) {
+		t.Fatal(routes)
+	}
+}