From 5b6923c27277c5a73bbb56a080ef20d2b471033c Mon Sep 17 00:00:00 2001
From: kaiyou <pierre@jaury.eu>
Date: Sun, 10 May 2020 18:36:21 +0200
Subject: [PATCH] Unexporter AssignOverlayAddress

Instead, the overlay address is assigned upon creation.
Also, the wireguard state is responsible for populating
the local node object.
---
 main.go         | 12 +++++-------
 wg/wireguard.go | 14 +++++++++++---
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/main.go b/main.go
index f15fc95..2f62bda 100644
--- a/main.go
+++ b/main.go
@@ -35,14 +35,14 @@ func main() {
 	logrus.SetLevel(logLevel)
 
 	// Create the wireguard and cluster configuration
-	wgstate, err := wg.New(config.Interface, config.WireguardPort)
-	if err != nil {
-		logrus.WithError(err).Fatal("could not instantiate wireguard controller")
-	}
 	cluster, err := cluster.New(config.Init, config.ClusterKey, config.BindAddr, config.ClusterPort, config.UseIPAsName)
 	if err != nil {
 		logrus.WithError(err).Fatal("could not create cluster")
 	}
+	wgstate, err := wg.New(config.Interface, config.WireguardPort, (*net.IPNet)(config.OverlayNet), cluster.Name())
+	if err != nil {
+		logrus.WithError(err).Fatal("could not instantiate wireguard controller")
+	}
 
 	// Prepare the /etc/hosts writer
 	hostsFile := &etchosts.EtcHosts{
@@ -50,10 +50,8 @@ func main() {
 	}
 
 	// Assign a local node address and propagate it to the cluster
-	wgstate.AssignOverlayAddr((*net.IPNet)(config.OverlayNet), cluster.Name())
 	localNode := &common.Node{}
-	localNode.OverlayAddr = wgstate.OverlayAddr
-	localNode.PubKey = wgstate.PubKey.String()
+	wgstate.UpdateNode(localNode)
 	cluster.SetLocalNode(localNode)
 	cluster.Update()
 
diff --git a/wg/wireguard.go b/wg/wireguard.go
index 733d36a..d6830d6 100644
--- a/wg/wireguard.go
+++ b/wg/wireguard.go
@@ -25,7 +25,7 @@ type State struct {
 // New creates a new Wesher Wireguard state
 // The Wireguard keys are generated for every new interface
 // The interface must later be setup using SetUpInterface
-func New(iface string, port int) (*State, error) {
+func New(iface string, port int, ipnet *net.IPNet, name string) (*State, error) {
 	client, err := wgctrl.New()
 	if err != nil {
 		return nil, errors.Wrap(err, "could not instantiate wireguard client")
@@ -44,15 +44,23 @@ func New(iface string, port int) (*State, error) {
 		PrivKey: privKey,
 		PubKey:  pubKey,
 	}
+	state.assignOverlayAddr(ipnet, name)
+
 	return &state, nil
 }
 
-// AssignOverlayAddr assigns a new address to the interface
+// UpdateNode populates a node instance with wireguard specific fields
+func (s *State) UpdateNode(node *common.Node) {
+	node.OverlayAddr = s.OverlayAddr
+	node.PubKey = s.PubKey.String()
+}
+
+// assignOverlayAddr assigns a new address to the interface
 // The address is assigned inside the provided network and depends on the
 // provided name deterministically
 // Currently, the address is assigned by hashing the name and mapping that
 // hash in the target network space
-func (s *State) AssignOverlayAddr(ipnet *net.IPNet, name string) {
+func (s *State) assignOverlayAddr(ipnet *net.IPNet, name string) {
 	// TODO: this is way too brittle and opaque
 	bits, size := ipnet.Mask.Size()
 	ip := make([]byte, len(ipnet.IP))