refactor Nodes

2 files changed, 81 insertions, 52 deletions

diff --git a/internal/config/config.go b/internal/config/config.go
index a882268..b5030ee 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -11,19 +11,26 @@ import (
 	"codeberg.org/snonux/gorum/internal/utils"
 )
 
+type Node struct {
+	Port       int
+	Priority   int
+	Alias      string `json:",omitempty"`
+	Hostname   string
+	originalId string
+}
+
 type Config struct {
-	StateDir        string
-	ScoreFile       string
-	WinnerFile      string
-	Address         string
-	Nodes           []string
-	LoopIntervalS   int64  `json:"LoopIntervalS,omitempty"`
-	MyID            string `json:"MyID,omitempty"`
-	RelaxedMode     bool   `json:"RelaxedMode,omitempty"`
-	nodeNumberCache map[string]int
-	EmailTo         string `json:"EmailTo,omitempty"`
-	EmailFrom       string `json:"EmailFrom,omitempty"`
-	SMTPServer      string `json:"SMTPServer,omitempty"`
+	StateDir      string
+	ScoreFile     string
+	WinnerFile    string
+	Address       string
+	Nodes         map[string]Node
+	LoopIntervalS int64  `json:"LoopIntervalS,omitempty"`
+	MyID          string `json:"MyID,omitempty"`
+	RelaxedMode   bool   `json:"RelaxedMode,omitempty"`
+	EmailTo       string `json:"EmailTo,omitempty"`
+	EmailFrom     string `json:"EmailFrom,omitempty"`
+	SMTPServer    string `json:"SMTPServer,omitempty"`
 }
 
 func NewFromConfigFile(configFile string) (Config, error) {
@@ -45,6 +52,15 @@ func NewFromConfigFile(configFile string) (Config, error) {
 		return conf, err
 	}
 
+	// Make it so that the key is the Hostname for internal lookup.
+	newNodes := make(map[string]Node, len(conf.Nodes))
+	for id, node := range conf.Nodes {
+		log.Printf("adding node %s: %v", id, node)
+		node.originalId = id
+		newNodes[node.Hostname] = node
+	}
+	conf.Nodes = newNodes
+
 	return conf.setDefaults()
 }
 
@@ -70,9 +86,14 @@ func WithAddress(address string) func(*Config) {
 	}
 }
 
-func WithNodes(nodes ...string) func(*Config) {
+func WithNodes(nodes ...Node) func(*Config) {
 	return func(conf *Config) {
-		conf.Nodes = append(conf.Nodes, nodes...)
+		if conf.Nodes == nil {
+			conf.Nodes = make(map[string]Node)
+		}
+		for _, node := range nodes {
+			conf.Nodes[node.Hostname] = node
+		}
 	}
 }
 
@@ -106,35 +127,20 @@ func (conf Config) setDefaults() (Config, error) {
 		conf.MyID = hostname
 	}
 
-	conf.nodeNumberCache = make(map[string]int, len(conf.Nodes))
-	for i, node := range conf.Nodes {
-		conf.nodeNumberCache[utils.StripPort(node)] = i
-	}
 	return conf, nil
 }
 
-func (conf Config) NodeNumber(node string) int {
-	node = utils.StripPort(node)
-	nodeNumber, ok := conf.nodeNumberCache[node]
-	if ok {
-		return nodeNumber
+func (conf Config) NodePriority(id string) (int, error) {
+	node, err := conf.findNode(id)
+	if err != nil {
+		return 0, err
 	}
-
-	log.Println("config:",
-		fmt.Errorf("node %s not found - it will affect it's score!", node))
-	return 0
+	return node.Priority, nil
 }
 
 func (conf Config) IsNode(remoteAddr string) bool {
-	remoteAddr = utils.StripPort(remoteAddr)
-
-	for _, node := range conf.Nodes {
-		if remoteAddr == utils.StripPort(node) {
-			return true
-		}
-	}
-
-	return false
+	_, err := conf.findNode(utils.StripPort(remoteAddr))
+	return err == nil
 }
 
 func (conf Config) IsNodeWithLookup(remoteAddr string,
@@ -142,19 +148,38 @@ func (conf Config) IsNodeWithLookup(remoteAddr string,
 
 	remoteAddr = utils.StripPort(remoteAddr)
 
-	for _, node := range conf.Nodes {
-		ips, err := lookupIP(utils.StripPort(node))
+	compare := func(hostname string) bool {
+		ips, err := lookupIP(hostname)
 		if err != nil {
 			log.Println("config:", err)
-			continue
+			return false
 		}
-
 		for _, ip := range ips {
 			if remoteAddr == ip.String() {
 				return true
 			}
 		}
+		return false
 	}
 
+	for _, node := range conf.Nodes {
+		if compare(node.Hostname) || compare(node.Alias) {
+			return true
+		}
+	}
 	return false
 }
+
+func (conf Config) findNode(hostname string) (Node, error) {
+	if node, ok := conf.Nodes[hostname]; ok {
+		return node, nil
+	}
+
+	for _, node := range conf.Nodes {
+		if hostname == node.Alias {
+			return node, nil
+		}
+	}
+
+	return Node{}, fmt.Errorf("node %s not found", hostname)
+}
diff --git a/internal/config/config_test.go b/internal/config/config_test.go
index e24c048..900f045 100644
--- a/internal/config/config_test.go
+++ b/internal/config/config_test.go
@@ -6,30 +6,32 @@ import (
 	"testing"
 )
 
-func TestNodeNumber(t *testing.T) {
+func TestNodePriority(t *testing.T) {
 	t.Parallel()
-	conf, _ := New(WithNodes("localhost:1234", "hamburger:4321"))
+	conf, _ := New(WithNodes(
+		Node{Hostname: "localhost", Port: 1234, Priority: 200},
+		Node{Hostname: "hamburger", Port: 4321, Priority: 100}))
 
-	num := conf.NodeNumber("localhost")
-	if num != 0 {
-		t.Errorf("localhost should be node number 0 but is %d", num)
+	num, _ := conf.NodePriority("localhost")
+	if num != 200 {
+		t.Errorf("localhost should have priority 200 but has %d", num)
 	}
-
-	num = conf.NodeNumber("hamburger")
-	if num != 1 {
-		t.Errorf("hamburger should be node number 1 but is %d", num)
+	num, _ = conf.NodePriority("hamburger")
+	if num != 100 {
+		t.Errorf("hamburger should have priority 100 but has %d", num)
 	}
 }
 
 func TestIsNode(t *testing.T) {
 	t.Parallel()
-	conf, _ := New(WithNodes("localhost:1234", "hamburger:4321"))
+	conf, _ := New(WithNodes(
+		Node{Hostname: "localhost", Port: 1234, Priority: 200},
+		Node{Hostname: "hamburger", Port: 4321, Priority: 100}))
 
 	remoteAddr := "localhost:323232"
 	if !conf.IsNode(remoteAddr) {
 		t.Errorf("%s should be node of %v", remoteAddr, conf.Nodes)
 	}
-
 	remoteAddr = "foo.zone:2345"
 	if conf.IsNode(remoteAddr) {
 		t.Errorf("%s should not be node of %v", remoteAddr, conf.Nodes)
@@ -38,7 +40,9 @@ func TestIsNode(t *testing.T) {
 
 func TestIsNodeWithLookup(t *testing.T) {
 	t.Parallel()
-	conf, _ := New(WithNodes("localhost:1234", "hamburger:4321"))
+	conf, _ := New(WithNodes(
+		Node{Hostname: "localhost", Port: 1234, Priority: 200},
+		Node{Hostname: "hamburger", Port: 4321, Priority: 100}))
 
 	lookupIP := func(addr string) ([]net.IP, error) {
 		switch addr {