mc-router/server/routes.go

package server

import (
	"context"
	"regexp"
	"strings"
	"sync"

	"github.com/sirupsen/logrus"
)

// WakerFunc is a function that wakes up a server and returns its address.
type WakerFunc func(ctx context.Context) (string, error)

// SleeperFunc is a function that puts a server to sleep.
type SleeperFunc func(ctx context.Context) error

func buildWakerFromSleeper(endpoint string, sleeper SleeperFunc) WakerFunc {
	if sleeper == nil {
		return nil
	}
	return func(ctx context.Context) (string, error) {
		if err := sleeper(ctx); err != nil {
			return "", err
		}
		return endpoint, nil
	}
}

var tcpShieldPattern = regexp.MustCompile("///.*")

// RouteFinder implementations find new routes in the system that can be tracked by a RoutesHandler
type RouteFinder interface {
	Start(ctx context.Context, handler RoutesHandler) error
	String() string
}

type RoutesHandler interface {
	CreateMapping(serverAddress string, backend string, scalingTarget string, waker WakerFunc, sleeper SleeperFunc, asleepMOTD string, loadingMOTD string)
	SetDefaultRoute(backend string, scalingTarget string, waker WakerFunc, sleeper SleeperFunc, asleepMOTD string, loadingMOTD string)
	// DeleteMapping requests that the serverAddress be removed from routes.
	// Returns true if the route existed.
	DeleteMapping(serverAddress string) bool
}

type IRoutes interface {
	RoutesHandler

	Reset()
	RegisterAll(mappings map[string]string)
	// FindBackendForServerAddress returns the host:port for the external server address, if registered.
	// Otherwise, an empty string is returned. Also returns the normalized version of the given serverAddress.
	// The 3rd value returned is the scalingTarget which indicates what endpoint to scale (may differ from backend when using proxy).
	// The 4th value returned is an (optional) "waker" function which a caller must invoke to wake up serverAddress.
	// The 5th value returned is an (optional) "sleeper" function which a caller must invoke to shut down serverAddress.
	HasRoute(serverAddress string) bool
	FindBackendForServerAddress(ctx context.Context, serverAddress string) (string, string, string, WakerFunc, SleeperFunc)
	GetSleepers(scalingTarget string) []SleeperFunc
	GetMappings() map[string]string
	GetDefaultRoute() (string, string, WakerFunc, SleeperFunc)
	GetAsleepMOTD(serverAddress string) string
	GetLoadingMOTD(serverAddress string) string
	SimplifySRV(srvEnabled bool)
}

var Routes = NewRoutes()

func NewRoutes() IRoutes {
	r := &routesImpl{
		mappings: make(map[string]mapping),
	}

	return r
}

func (r *routesImpl) RegisterAll(mappings map[string]string) {
	for k, v := range mappings {
		r.CreateMapping(k, v, "", nil, nil, "", "")
	}
}

type mapping struct {
	backend       string
	waker         WakerFunc
	sleeper       SleeperFunc
	asleepMOTD    string
	loadingMOTD   string
	scalingTarget string // The endpoint to scale (may differ from backend when using proxy)
}

type routesImpl struct {
	sync.RWMutex
	mappings     map[string]mapping
	defaultRoute mapping
	simplifySRV  bool
}

func (r *routesImpl) Reset() {
	r.mappings = make(map[string]mapping)
	DownScaler.Reset()
}

func (r *routesImpl) SetDefaultRoute(backend string, scalingTarget string, waker WakerFunc, sleeper SleeperFunc, asleepMOTD string, loadingMOTD string) {
	if scalingTarget == "" {
		scalingTarget = backend
	}
	r.defaultRoute = mapping{backend: backend, scalingTarget: scalingTarget, waker: waker, sleeper: sleeper, asleepMOTD: asleepMOTD, loadingMOTD: loadingMOTD}

	logrus.WithFields(logrus.Fields{
		"backend": backend,
	}).Info("Using default route")
}

func (r *routesImpl) GetDefaultRoute() (string, string, WakerFunc, SleeperFunc) {
	return r.defaultRoute.backend, r.defaultRoute.scalingTarget, r.defaultRoute.waker, r.defaultRoute.sleeper
}

func (r *routesImpl) GetAsleepMOTD(serverAddress string) string {
	r.RLock()
	defer r.RUnlock()

	if serverAddress == "" {
		return r.defaultRoute.asleepMOTD
	}

	if m, ok := r.mappings[serverAddress]; ok {
		return m.asleepMOTD
	}
	return ""
}

func (r *routesImpl) GetLoadingMOTD(serverAddress string) string {
	r.RLock()
	defer r.RUnlock()

	if serverAddress == "" {
		return r.defaultRoute.loadingMOTD
	}

	if m, ok := r.mappings[serverAddress]; ok {
		return m.loadingMOTD
	}
	return ""
}

func (r *routesImpl) SimplifySRV(srvEnabled bool) {
	r.simplifySRV = srvEnabled
}

func (r *routesImpl) HasRoute(serverAddress string) bool {
	r.RLock()
	defer r.RUnlock()

	_, exists := r.mappings[serverAddress]
	return exists
}

func (r *routesImpl) FindBackendForServerAddress(_ context.Context, serverAddress string) (string, string, string, WakerFunc, SleeperFunc) {
	r.RLock()
	defer r.RUnlock()

	// Trim off Forge null-delimited address parts like \x00FML3\x00
	serverAddress = strings.Split(serverAddress, "\x00")[0]

	// Trim off infinity-filter backslash address parts like \\GUID\\CLIENT_IP...
	serverAddress = strings.Split(serverAddress, "\\")[0]

	serverAddress = strings.ToLower(
		// trim the root zone indicator, see https://en.wikipedia.org/wiki/Fully_qualified_domain_name
		strings.TrimSuffix(serverAddress, "."))

	logrus.WithFields(logrus.Fields{
		"serverAddress": serverAddress,
	}).Debug("Finding backend for server address")

	if r.simplifySRV {
		parts := strings.Split(serverAddress, ".")
		tcpIndex := -1
		for i, part := range parts {
			if part == "_tcp" {
				tcpIndex = i
				break
			}
		}
		if tcpIndex != -1 {
			parts = parts[tcpIndex+1:]
		}

		serverAddress = strings.Join(parts, ".")
	}

	// Strip suffix of TCP Shield
	serverAddress = tcpShieldPattern.ReplaceAllString(serverAddress, "")

	if r.mappings != nil {
		if mapping, exists := r.mappings[serverAddress]; exists {
			return mapping.backend, serverAddress, mapping.scalingTarget, mapping.waker, mapping.sleeper
		}
	}
	return r.defaultRoute.backend, serverAddress, r.defaultRoute.scalingTarget, r.defaultRoute.waker, r.defaultRoute.sleeper
}

func (r *routesImpl) GetSleepers(scalingTarget string) []SleeperFunc {
	r.RLock()
	defer r.RUnlock()

	var sleepers []SleeperFunc
	for _, m := range r.mappings {
		if m.scalingTarget == scalingTarget && m.sleeper != nil {
			sleepers = append(sleepers, m.sleeper)
		}
	}
	if r.defaultRoute.scalingTarget == scalingTarget && r.defaultRoute.sleeper != nil {
		sleepers = append(sleepers, r.defaultRoute.sleeper)
	}
	return sleepers
}

func (r *routesImpl) GetMappings() map[string]string {
	r.RLock()
	defer r.RUnlock()

	result := make(map[string]string, len(r.mappings))
	for k, v := range r.mappings {
		result[k] = v.backend
	}
	return result
}

func (r *routesImpl) DeleteMapping(serverAddress string) bool {
	r.Lock()
	defer r.Unlock()
	logrus.WithField("serverAddress", serverAddress).Info("Deleting route")

	if m, ok := r.mappings[serverAddress]; ok {
		DownScaler.Cancel(m.scalingTarget)
		delete(r.mappings, serverAddress)
		return true
	} else {
		return false
	}
}

func (r *routesImpl) CreateMapping(serverAddress string, backend string, scalingTarget string, waker WakerFunc, sleeper SleeperFunc, asleepMOTD string, loadingMOTD string) {
	r.Lock()
	defer r.Unlock()

	serverAddress = strings.ToLower(serverAddress)

	if scalingTarget == "" {
		scalingTarget = backend
	}

	logrus.WithFields(logrus.Fields{
		"serverAddress": serverAddress,
		"backend":       backend,
	}).Info("Created route mapping")
	r.mappings[serverAddress] = mapping{backend: backend, scalingTarget: scalingTarget, waker: waker, sleeper: sleeper, asleepMOTD: asleepMOTD, loadingMOTD: loadingMOTD}

	// Trigger auto scale down when mapping is created to ensure servers are shut down if router restarts
	if DownScaler != nil && scalingTarget != "" {
		DownScaler.Begin(scalingTarget)
	}
}