Allow plugins to override server selection when a client joins

Closes #129.

5 files changed, 104 insertions, 13 deletions

diff --git a/config.go b/config.go
index 9b6f1e3..176126d 100644
--- a/config.go
+++ b/config.go
@@ -50,6 +50,7 @@ type Config struct {
 	TelnetAddr       string
 	BindAddr         string
 	DefaultSrv       string
+	SrvSelector      string
 	Servers          map[string]Server
 	ForceDefaultSrv  bool
 	KickOnNewPool    bool
diff --git a/doc/config.md b/doc/config.md
index 687476e..36a16c9 100644
--- a/doc/config.md
+++ b/doc/config.md
@@ -103,6 +103,14 @@ Default: ""
 Description: The default server to connect new clients to. May be a group.
 ```
 
+> `SrvSelector`
+```
+Type: string
+Default: ""
+Description: The server selection handler to use for new clients.
+Only useful when combined with a plugin providing a handler with that name.
+```
+
 > `Servers`
 ```
 Type: map[string]Server
diff --git a/doc/server_selectors.md b/doc/server_selectors.md
new file mode 100644
index 0000000..ce4a4d5
--- /dev/null
+++ b/doc/server_selectors.md
@@ -0,0 +1,12 @@
+# Server selectors
+
+If needed, plugins can provide a custom function to choose the server to
+connect a new client to. To do this, call the [RegisterSrvSelector](https://pkg.go.dev/github.com/HimbeerserverDE/mt-multiserver-proxy#RegisterSrvSelector)
+function.
+
+Use the `SrvSelector` configuration option to choose one of the registered
+server selectors. This option is reloadable.
+
+If the configured server selector doesn't exist or doesn't return a server, the
+proxy uses the builtin server selection strategy. In the case where the
+configured server selector doesn't exist, a log message is generated.
diff --git a/plugin_srvselect.go b/plugin_srvselect.go
new file mode 100644
index 0000000..327efd5
--- /dev/null
+++ b/plugin_srvselect.go
@@ -0,0 +1,66 @@
+package proxy
+
+import (
+	"errors"
+	"sync"
+)
+
+var (
+	ErrEmptySrvSelectorName = errors.New("server selector name is empty")
+)
+
+var (
+	srvSelectors    map[string]func(*ClientConn) (string, Server)
+	srvSelectorMu   sync.RWMutex
+	srvSelectorOnce sync.Once
+)
+
+// RegisterSrvSelector registers a server selection handler
+// which can be enabled using the `SrvSelector` config option.
+// Empty names are an error.
+// If the handler returns an empty server name,
+// the regular server selection procedure is used.
+// Only one server selector can be active at a time.
+func RegisterSrvSelector(name string, sel func(*ClientConn) (string, Server)) error {
+	initSrvSelectors()
+
+	if name == "" {
+		return ErrEmptySrvSelectorName
+	}
+
+	srvSelectorMu.Lock()
+	defer srvSelectorMu.Unlock()
+
+	srvSelectors[name] = sel
+	return nil
+}
+
+func selectSrv(cc *ClientConn) (string, Server) {
+	sel := Conf().SrvSelector
+
+	if sel == "" {
+		return "", Server{}
+	}
+
+	initSrvSelectors()
+
+	srvSelectorMu.RLock()
+	defer srvSelectorMu.RUnlock()
+
+	handler, ok := srvSelectors[sel]
+	if !ok {
+		cc.Log("<-", "server selector not registered")
+		return "", Server{}
+	}
+
+	return handler(cc)
+}
+
+func initSrvSelectors() {
+	srvSelectorOnce.Do(func() {
+		srvSelectorMu.Lock()
+		defer srvSelectorMu.Unlock()
+
+		srvSelectors = make(map[string]func(*ClientConn) (string, Server))
+	})
+}
diff --git a/run.go b/run.go
index edd0856..21c2664 100644
--- a/run.go
+++ b/run.go
@@ -123,22 +123,26 @@ func runFunc() {
 				cc.Kick("No servers are configured.")
 				return
 			}
-			if _, ok := conf.Servers[""]; !ok && conf.DefaultServerName() == "" {
-				cc.Log("<-", "no default server")
-				cc.Kick("No valid default server is configured.")
-				return
-			}
 
-			srvName, srv := conf.DefaultServerInfo()
-			lastSrv, err := DefaultAuth().LastSrv(cc.Name())
-			if err == nil && !conf.ForceDefaultSrv && lastSrv != srvName {
-				choice, ok := conf.RandomGroupServer(lastSrv)
-				if !ok {
-					cc.Log("<-", "inexistent previous server")
+			srvName, srv := selectSrv(cc)
+			if srvName == "" {
+				if _, ok := conf.Servers[""]; !ok && conf.DefaultServerName() == "" {
+					cc.Log("<-", "no default server")
+					cc.Kick("No valid default server is configured.")
+					return
 				}
 
-				srvName = choice
-				srv, _ = conf.Servers[choice] // Existence already checked.
+				srvName, srv = conf.DefaultServerInfo()
+				lastSrv, err := DefaultAuth().LastSrv(cc.Name())
+				if err == nil && !conf.ForceDefaultSrv && lastSrv != srvName {
+					choice, ok := conf.RandomGroupServer(lastSrv)
+					if !ok {
+						cc.Log("<-", "inexistent previous server")
+					}
+
+					srvName = choice
+					srv, _ = conf.Servers[choice] // Existence already checked.
+				}
 			}
 
 			doConnect := func(srvName string, srv Server) error {