1 files changed, 65 insertions, 40 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)
+}