From 1127cb194b59e7ea99f6a502c97a5e33e1de5526 Mon Sep 17 00:00:00 2001
From: Sean Baildon <sean@baildon.co>
Date: Sun, 22 Aug 2021 12:51:31 +0100
Subject: [PATCH] pkg/mesh: optionally assign external IP to node's private IP

---
 cmd/kg/main.go   | 42 ++++++++++++++++++++++--------------------
 pkg/mesh/mesh.go | 10 ++++++++--
 2 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/cmd/kg/main.go b/cmd/kg/main.go
index 6ba443c..1abc6ed 100644
--- a/cmd/kg/main.go
+++ b/cmd/kg/main.go
@@ -92,25 +92,26 @@ var cmd = &cobra.Command{
 }
 
 var (
-	backend       string
-	cleanUpIface  bool
-	createIface   bool
-	cni           bool
-	cniPath       string
-	compatibility string
-	encapsulate   string
-	granularity   string
-	hostname      string
-	kubeconfig    string
-	iface         string
-	listen        string
-	local         bool
-	master        string
-	mtu           uint
-	topologyLabel string
-	port          uint
-	subnet        string
-	resyncPeriod  time.Duration
+	backend               string
+	cleanUpIface          bool
+	createIface           bool
+	cni                   bool
+	cniPath               string
+	compatibility         string
+	encapsulate           string
+	granularity           string
+	hostname              string
+	kubeconfig            string
+	iface                 string
+	listen                string
+	local                 bool
+	master                string
+	mtu                   uint
+	topologyLabel         string
+	port                  uint
+	subnet                string
+	resyncPeriod          time.Duration
+	prioritisePrivateAddr bool
 
 	printVersion bool
 	logLevel     string
@@ -139,6 +140,7 @@ func init() {
 	cmd.Flags().UintVar(&port, "port", mesh.DefaultKiloPort, "The port over which WireGuard peers should communicate.")
 	cmd.Flags().StringVar(&subnet, "subnet", mesh.DefaultKiloSubnet.String(), "CIDR from which to allocate addresses for WireGuard interfaces.")
 	cmd.Flags().DurationVar(&resyncPeriod, "resync-period", 30*time.Second, "How often should the Kilo controllers reconcile?")
+	cmd.Flags().BoolVar(&prioritisePrivateAddr, "prioritise-private-addresses", false, "Prefer to assign a private IP address to the node's endpoint")
 
 	cmd.PersistentFlags().BoolVar(&printVersion, "version", false, "Print version and exit")
 	cmd.PersistentFlags().StringVar(&logLevel, "log-level", logLevelInfo, fmt.Sprintf("Log level to use. Possible values: %s", availableLogLevels))
@@ -234,7 +236,7 @@ func runRoot(_ *cobra.Command, _ []string) error {
 		return fmt.Errorf("backend %v unknown; possible values are: %s", backend, availableBackends)
 	}
 
-	m, err := mesh.New(b, enc, gr, hostname, uint32(port), s, local, cni, cniPath, iface, cleanUpIface, createIface, mtu, resyncPeriod, log.With(logger, "component", "kilo"))
+	m, err := mesh.New(b, enc, gr, hostname, uint32(port), s, local, cni, cniPath, iface, cleanUpIface, createIface, mtu, resyncPeriod, prioritisePrivateAddr, log.With(logger, "component", "kilo"))
 	if err != nil {
 		return fmt.Errorf("failed to create Kilo mesh: %v", err)
 	}
diff --git a/pkg/mesh/mesh.go b/pkg/mesh/mesh.go
index 18f84a5..7c63ca1 100644
--- a/pkg/mesh/mesh.go
+++ b/pkg/mesh/mesh.go
@@ -86,7 +86,7 @@ type Mesh struct {
 }
 
 // New returns a new Mesh instance.
-func New(backend Backend, enc encapsulation.Encapsulator, granularity Granularity, hostname string, port uint32, subnet *net.IPNet, local, cni bool, cniPath, iface string, cleanUpIface bool, createIface bool, mtu uint, resyncPeriod time.Duration, logger log.Logger) (*Mesh, error) {
+func New(backend Backend, enc encapsulation.Encapsulator, granularity Granularity, hostname string, port uint32, subnet *net.IPNet, local, cni bool, cniPath, iface string, cleanUpIface bool, createIface bool, mtu uint, resyncPeriod time.Duration, prioritisePrivateAddr bool, logger log.Logger) (*Mesh, error) {
 	if err := os.MkdirAll(kiloPath, 0700); err != nil {
 		return nil, fmt.Errorf("failed to create directory to store configuration: %v", err)
 	}
@@ -143,6 +143,12 @@ func New(backend Backend, enc encapsulation.Encapsulator, granularity Granularit
 		enc = encapsulation.Noop(enc.Strategy())
 		level.Debug(logger).Log("msg", "running without a private IP address")
 	}
+	var externalIP *net.IPNet
+	if prioritisePrivateAddr && privateIP != nil {
+		externalIP = privateIP
+	} else {
+		externalIP = publicIP
+	}
 	level.Debug(logger).Log("msg", fmt.Sprintf("using %s as the public IP address", publicIP.String()))
 	ipTables, err := iptables.New(iptables.WithLogger(log.With(logger, "component", "iptables")), iptables.WithResyncPeriod(resyncPeriod))
 	if err != nil {
@@ -154,7 +160,7 @@ func New(backend Backend, enc encapsulation.Encapsulator, granularity Granularit
 		cni:          cni,
 		cniPath:      cniPath,
 		enc:          enc,
-		externalIP:   publicIP,
+		externalIP:   externalIP,
 		granularity:  granularity,
 		hostname:     hostname,
 		internalIP:   privateIP,