From 31b8263e96e037525ca4af0769eeb7ba28db7e0b Mon Sep 17 00:00:00 2001
From: kaiyou <pierre@jaury.eu>
Date: Thu, 14 May 2020 12:03:09 +0200
Subject: [PATCH] Use the StatePath to config when loading the state

---
 cluster/cluster.go | 14 ++++++++------
 cluster/state.go   | 13 +++++++++----
 main.go            |  2 +-
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/cluster/cluster.go b/cluster/cluster.go
index 31fa96c..8258ce7 100644
--- a/cluster/cluster.go
+++ b/cluster/cluster.go
@@ -23,16 +23,17 @@ type Cluster struct {
 	mlConfig  *memberlist.Config
 	localNode *common.Node
 	LocalName string
+	statePath string
 	state     *state
 	events    chan memberlist.NodeEvent
 }
 
 // New is used to create a new Cluster instance
 // The returned instance is ready to be updated with the local node settings then joined
-func New(init bool, clusterKey []byte, bindAddr string, bindPort int, useIPAsName bool) (*Cluster, error) {
+func New(init bool, clusterKey []byte, bindAddr string, bindPort int, statePath string, useIPAsName bool) (*Cluster, error) {
 	state := &state{}
 	if !init {
-		loadState(state)
+		loadState(state, statePath)
 	}
 
 	clusterKey, err := computeClusterKey(state, clusterKey)
@@ -61,8 +62,9 @@ func New(init bool, clusterKey []byte, bindAddr string, bindPort int, useIPAsNam
 		LocalName: ml.LocalNode().Name,
 		// The big channel buffer is a work-around for https://github.com/hashicorp/memberlist/issues/23
 		// More than this many simultaneous events will deadlock cluster.members()
-		events: make(chan memberlist.NodeEvent, 100),
-		state:  state,
+		events:    make(chan memberlist.NodeEvent, 100),
+		state:     state,
+		statePath: statePath,
 	}
 	return &cluster, nil
 }
@@ -94,7 +96,7 @@ func (c *Cluster) Join(addrs []string) error {
 
 // Leave saves the current state before leaving, then leaves the cluster
 func (c *Cluster) Leave() {
-	c.state.save()
+	c.state.save(c.statePath)
 	c.ml.Leave(10 * time.Second)
 	c.ml.Shutdown() //nolint: errcheck
 }
@@ -144,7 +146,7 @@ func (c *Cluster) Members() <-chan []common.Node {
 			}
 			c.state.Nodes = nodes
 			changes <- nodes
-			c.state.save()
+			c.state.save(c.statePath)
 		}
 	}()
 	return changes
diff --git a/cluster/state.go b/cluster/state.go
index c6a8fa3..1904b5a 100644
--- a/cluster/state.go
+++ b/cluster/state.go
@@ -16,10 +16,12 @@ type state struct {
 	Nodes      []common.Node
 }
 
-// TODO: this should be replaced by a configurable path later
-var statePath = "/var/lib/wesher/state.json"
+const defaultStatePath = "/var/lib/wesher/state.json"
 
-func (s *state) save() error {
+func (s *state) save(statePath string) error {
+	if statePath == "" {
+		statePath = defaultStatePath
+	}
 	if err := os.MkdirAll(path.Dir(statePath), 0700); err != nil {
 		return err
 	}
@@ -32,7 +34,10 @@ func (s *state) save() error {
 	return ioutil.WriteFile(statePath, stateOut, 0600)
 }
 
-func loadState(cs *state) {
+func loadState(cs *state, statePath string) {
+	if statePath == "" {
+		statePath = defaultStatePath
+	}
 	content, err := ioutil.ReadFile(statePath)
 	if err != nil {
 		if !os.IsNotExist(err) {
diff --git a/main.go b/main.go
index 781966d..0a23024 100644
--- a/main.go
+++ b/main.go
@@ -35,7 +35,7 @@ func main() {
 	logrus.SetLevel(logLevel)
 
 	// Create the wireguard and cluster configuration
-	cluster, err := cluster.New(config.Init, config.ClusterKey, config.BindAddr, config.ClusterPort, config.UseIPAsName)
+	cluster, err := cluster.New(config.Init, config.ClusterKey, config.BindAddr, config.ClusterPort, config.StatePath, config.UseIPAsName)
 	if err != nil {
 		logrus.WithError(err).Fatal("could not create cluster")
 	}