From 7cf95b6356e55b292efdab89447335fc8d935fca Mon Sep 17 00:00:00 2001
From: Spencer Smith <rsmitty@users.noreply.github.com>
Date: Thu, 22 Oct 2020 04:05:59 -0400
Subject: [PATCH] [rtnl] support functional args for routes (#86)

This PR will add the ability to use functional args for route
add/replace with rtnl. Currently, two functions are supported,
`WithRouteSrc` and `WithRouteAttrs`. This allows the consolidation of
the route methods, as well as adding the ability to BYO-Attributes for a
given route. This layout can of course be extended to add other
route-specific args, as well as a starting point for options and args for other parts of the
rtnl package as time goes on.

Note that this removes the `RouteAddSrc` and `RouteReplaceSrc`
functions.

Signed-off-by: Spencer Smith <robertspencersmith@gmail.com>
---
 rtnl/route.go         | 55 +++++++++++++------------------------------
 rtnl/route_options.go | 47 ++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 38 deletions(-)
 create mode 100644 rtnl/route_options.go

diff --git a/rtnl/route.go b/rtnl/route.go
index 9859a45..24565af 100644
--- a/rtnl/route.go
+++ b/rtnl/route.go
@@ -9,7 +9,14 @@ import (
 )
 
 // generating route message
-func genRouteMessage(ifc *net.Interface, dst net.IPNet, src *net.IPNet, gw net.IP) (rm *rtnetlink.RouteMessage, err error) {
+func genRouteMessage(ifc *net.Interface, dst net.IPNet, gw net.IP, options ...RouteOption) (rm *rtnetlink.RouteMessage, err error) {
+
+	opts := DefaultRouteOptions(ifc, dst, gw)
+
+	for _, option := range options {
+		option(opts)
+	}
+
 	af, err := addrFamily(dst.IP)
 	if err != nil {
 		return nil, err
@@ -27,19 +34,10 @@ func genRouteMessage(ifc *net.Interface, dst net.IPNet, src *net.IPNet, gw net.I
 		scope = unix.RT_SCOPE_LINK
 	}
 
-	attr := rtnetlink.RouteAttributes{
-		Dst:      dst.IP,
-		OutIface: uint32(ifc.Index),
-	}
-
-	if gw != nil {
-		attr.Gateway = gw
-	}
-
 	var srclen int
-	if src != nil {
-		srclen, _ = src.Mask.Size()
-		attr.Src = src.IP
+	if opts.Src != nil {
+		srclen, _ = opts.Src.Mask.Size()
+		opts.Attrs.Src = opts.Src.IP
 	}
 
 	dstlen, _ := dst.Mask.Size()
@@ -52,49 +50,30 @@ func genRouteMessage(ifc *net.Interface, dst net.IPNet, src *net.IPNet, gw net.I
 		Scope:      scope,
 		DstLength:  uint8(dstlen),
 		SrcLength:  uint8(srclen),
-		Attributes: attr,
+		Attributes: opts.Attrs,
 	}
 	return tx, nil
 }
 
 // RouteAdd adds infomation about a network route.
-func (c *Conn) RouteAdd(ifc *net.Interface, dst net.IPNet, gw net.IP) (err error) {
-	rm, err := genRouteMessage(ifc, dst, nil, gw)
+func (c *Conn) RouteAdd(ifc *net.Interface, dst net.IPNet, gw net.IP, options ...RouteOption) (err error) {
+	rm, err := genRouteMessage(ifc, dst, gw, options...)
 	if err != nil {
 		return err
 	}
+
 	return c.Conn.Route.Add(rm)
 }
 
 // RouteReplace adds or replace information about a network route.
-func (c *Conn) RouteReplace(ifc *net.Interface, dst net.IPNet, gw net.IP) (err error) {
-	rm, err := genRouteMessage(ifc, dst, nil, gw)
+func (c *Conn) RouteReplace(ifc *net.Interface, dst net.IPNet, gw net.IP, options ...RouteOption) (err error) {
+	rm, err := genRouteMessage(ifc, dst, gw, options...)
 	if err != nil {
 		return err
 	}
 	return c.Conn.Route.Replace(rm)
 }
 
-// RouteReplaceSrc adds or replace infomation about a network route with the given destination
-// and source. If source is `nil` it's ignored.
-func (c *Conn) RouteReplaceSrc(ifc *net.Interface, dst net.IPNet, src *net.IPNet, gw net.IP) (err error) {
-	rm, err := genRouteMessage(ifc, dst, src, gw)
-	if err != nil {
-		return err
-	}
-	return c.Conn.Route.Replace(rm)
-}
-
-// RouteAddSrc adds infomation about a network route with the given destination
-// and source. If source is `nil` it's ignored.
-func (c *Conn) RouteAddSrc(ifc *net.Interface, dst net.IPNet, src *net.IPNet, gw net.IP) (err error) {
-	rm, err := genRouteMessage(ifc, dst, src, gw)
-	if err != nil {
-		return err
-	}
-	return c.Conn.Route.Add(rm)
-}
-
 // RouteDel deletes the route to the given destination.
 func (c *Conn) RouteDel(ifc *net.Interface, dst net.IPNet) error {
 	af, err := addrFamily(dst.IP)
diff --git a/rtnl/route_options.go b/rtnl/route_options.go
new file mode 100644
index 0000000..26ddf59
--- /dev/null
+++ b/rtnl/route_options.go
@@ -0,0 +1,47 @@
+package rtnl
+
+import (
+	"net"
+
+	"github.com/jsimonetti/rtnetlink"
+)
+
+// RouteOptions is the functional options struct
+type RouteOptions struct {
+	Src   *net.IPNet
+	Attrs rtnetlink.RouteAttributes
+}
+
+// RouteOption is the functional options func
+type RouteOption func(*RouteOptions)
+
+// DefaultRouteOptions defines the default route options.
+func DefaultRouteOptions(ifc *net.Interface, dst net.IPNet, gw net.IP) *RouteOptions {
+	ro := &RouteOptions{
+		Src: nil,
+		Attrs: rtnetlink.RouteAttributes{
+			Dst:      dst.IP,
+			OutIface: uint32(ifc.Index),
+		},
+	}
+
+	if gw != nil {
+		ro.Attrs.Gateway = gw
+	}
+
+	return ro
+}
+
+// WithRouteSrc sets the src address.
+func WithRouteSrc(src *net.IPNet) RouteOption {
+	return func(opts *RouteOptions) {
+		opts.Src = src
+	}
+}
+
+// WithRouteAttrs sets the attributes.
+func WithRouteAttrs(attrs rtnetlink.RouteAttributes) RouteOption {
+	return func(opts *RouteOptions) {
+		opts.Attrs = attrs
+	}
+}