Resolve server host issue

Author: kunrex

api/info.go

@@ -256,6 +256,12 @@ func challengeInfoHandler(c *gin.Context) {
 			PreRequisite:   strings.Split(challenge.PreReqs, core.DELIMITER),
 			DeployedStatus: challenge.Status,
 		}
+		deployedHost := challenge.ServerDeployed
+		if deployedHost != core.LOCALHOST && deployedHost != "" {
+			if s, ok := cfg.Cfg.AvailableServers[deployedHost]; ok {
+				deployedHost = s.Host
+			}
+		}
 		challengeInfo := Challenge{
 			ChallengeMetadata: challMetadata,
 			Description:       challenge.Description,
@@ -265,7 +271,7 @@ func challengeInfoHandler(c *gin.Context) {
 			AdditionalLinks:   strings.Split(challenge.AdditionalLinks, core.DELIMITER),
 			PreviousTries:     previousTries,
 			MaxAttemptLimit:   challenge.MaxAttemptLimit,
-			DeployedLink:      challenge.ServerDeployed,
+			DeployedLink:      deployedHost,
 		}
 		if user.Role == core.USER_ROLES["contestant"] {
 			c.JSON(http.StatusOK, challengeInfo)

cmd/beast/config.go

@@ -115,11 +115,14 @@ func promptServerDetails(configuration *config.BeastConfig) {
 		var server config.AvailableServer
 
 		server.Host = utils.PromptString("Enter Host Name, leave empty for localhost")
+		if server.Host == "" {
+			server.Host = core.LOCALHOST
+		}
 		server.Username = utils.PromptString("Enter Username")
 		server.SSHKeyPath = utils.PromptString("Enter SSH Key Path")
 		server.Active = utils.PromptBinary("Enable this server?")
 
-		configuration.AvailableServers[server.Username] = server
+		configuration.AvailableServers[server.Host] = server
 	}
 }
 

core/config/config.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"strings"
 	"time"
 
 	"github.com/sdslabs/beastv4/core"
@@ -126,7 +127,6 @@ type BeastConfig struct {
 	HealthProber         bool                       `toml:"health_prober"`
 	RemoteSyncPeriod     time.Duration              `toml:"-"`
 	Rsp                  string                     `toml:"remote_sync_period"`
-	LocalHostPortRange   string                     `toml:"local_host_port_range"`
 	InstanceConfig       InstanceConfig             `toml:"instance_config"`
 
 	CPUShares int64 `toml:"default_cpu_shares"`
@@ -225,15 +225,23 @@ func (config *BeastConfig) ValidateConfig() error {
 		log.Warn("No available servers provided for challenges. Using default localhost")
 		config.AvailableServers = map[string]AvailableServer{
 			core.LOCALHOST: {
+				Name:       core.LOCALHOST,
 				Host:       core.LOCALHOST,
 				Username:   os.Getenv("USER"),
 				SSHKeyPath: "",
 				Active:     true,
+				PortRange:  fmt.Sprintf("%v%s%v", core.ALLOWED_MIN_PORT_VALUE, core.MappingDelimeter, core.ALLOWED_MAX_PORT_VALUE),
 			},
 		}
 	}
 
-	for _, server := range config.AvailableServers {
+	for name, server := range config.AvailableServers {
+		if strings.Contains(name, ":") {
+			return fmt.Errorf("server key %q contains invalid character ':'", name)
+		}
+
+		server.Name = name
+		config.AvailableServers[name] = server
 		if server.Active {
 			err := server.ValidateServerConfig()
 			if err != nil {
@@ -284,11 +292,6 @@ func (config *BeastConfig) ValidateConfig() error {
 		}
 	}
 
-	err = ValidatePortRange(config.LocalHostPortRange)
-	if err != nil {
-		return fmt.Errorf("error while validating port range in global beast config: %s", err)
-	}
-
 	if config.CPUShares <= 0 {
 		log.Debug("Per container CPU shares not provided using default value")
 		config.CPUShares = core.DEFAULT_CPU_SHARE
@@ -314,6 +317,7 @@ func (config *BeastConfig) ValidateConfig() error {
 }
 
 type AvailableServer struct {
+	Name       string `toml:"-"`
 	Host       string `toml:"host"`
 	Username   string `toml:"username"`
 	SSHKeyPath string `toml:"ssh_key_path"`

core/constants.go

@@ -222,3 +222,4 @@ var NOTIFICATION_SERVICES = []string{
 	"discord",
 }
 
+const MappingDelimeter = ":"

core/manager/challenge.go

@@ -698,11 +698,9 @@ func undeployChallenge(challengeName string, purge bool) error {
 			}
 		}
 
-		var host string
-		if challenge.ServerDeployed == core.LOCALHOST || challenge.ServerDeployed == "" {
+		host := challenge.ServerDeployed
+		if host == "" {
 			host = core.LOCALHOST
-		} else {
-			host = config.Cfg.AvailableServers[challenge.ServerDeployed].Host
 		}
 
 		err = cache.FreeContainerPortsOnHost(host, challenge.ContainerId)

core/manager/health_check.go

@@ -87,8 +87,14 @@ func ChallengesHealthProber(waitTime int) {
 			// Do a better job at health probing mechanism.
 			if len(allocatedPorts) > 0 {
 				port := int(allocatedPorts[0].PortNo)
+				probeHost := chall.ServerDeployed
+				if probeHost != core.LOCALHOST && probeHost != "" {
+					if s, ok := config.Cfg.AvailableServers[probeHost]; ok {
+						probeHost = s.Host
+					}
+				}
 				prober := probes.NewTcpProber()
-				result, err := prober.Probe(chall.ServerDeployed, port, time.Duration(core.DEFAULT_PROBE_TIMEOUT)*time.Second)
+				result, err := prober.Probe(probeHost, port, time.Duration(core.DEFAULT_PROBE_TIMEOUT)*time.Second)
 				if err != nil {
 					msg := fmt.Sprintf("NETWORK HEALTH CHECK %s: %s : %s", result, chall.Name, err)
 					log.WithFields(log.Fields{

core/manager/instance.go

@@ -269,12 +269,8 @@ func allocateInstancePort(host string) (uint32, error) {
 	var firstPort, lastPort uint32
 	var err error
 
-	if host == core.LOCALHOST || host == "" {
-		firstPort, lastPort, err = utils.ParsePortMapping(cfg.Cfg.LocalHostPortRange)
-	} else {
-		server := cfg.Cfg.AvailableServers[host]
-		firstPort, lastPort, err = utils.ParsePortMapping(server.PortRange)
-	}
+	server := cfg.Cfg.AvailableServers[host]
+	firstPort, lastPort, err = utils.ParsePortMapping(server.PortRange)
 
 	if err != nil {
 		return 0, fmt.Errorf("failed to parse port range: %w", err)
@@ -291,8 +287,8 @@ func allocateInstancePort(host string) (uint32, error) {
 
 func selectServerForInstance() string {
 	availableServer, err := remoteManager.ServerQueue.GetNextAvailableInstance()
-	if err == nil && availableServer.Host != "" {
-		return availableServer.Host
+	if err == nil && availableServer.Name != "" {
+		return availableServer.Name
 	}
 	return core.LOCALHOST
 }

core/manager/pipeline.go

@@ -273,11 +273,8 @@ func deployChallenge(challenge *database.Challenge, config cfg.BeastChallengeCon
 			return fmt.Errorf("failed to extract port variables: %w", err)
 		}
 
-		var serverDeployed string
-		if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" {
-			server := cfg.Cfg.AvailableServers[challenge.ServerDeployed]
-			serverDeployed = server.Host
-		} else {
+		serverDeployed := challenge.ServerDeployed
+		if serverDeployed == "" {
 			serverDeployed = core.LOCALHOST
 		}
 
@@ -356,18 +353,15 @@ func deployChallenge(challenge *database.Challenge, config cfg.BeastChallengeCon
 	)
 
 	var err error
-	var host string
-	var firstPort, lastPort uint32
-	if challenge.ServerDeployed == core.LOCALHOST || challenge.ServerDeployed == "" {
+	host := challenge.ServerDeployed
+	if host == "" {
 		host = core.LOCALHOST
-		firstPort, lastPort, err = utils.ParsePortMapping(cfg.Cfg.LocalHostPortRange)
-	} else {
-		server := cfg.Cfg.AvailableServers[challenge.ServerDeployed]
-
-		host = server.Host
-		firstPort, lastPort, err = utils.ParsePortMapping(server.PortRange)
 	}
 
+	var firstPort, lastPort uint32
+	server := cfg.Cfg.AvailableServers[host]
+	firstPort, lastPort, err = utils.ParsePortMapping(server.PortRange)
+
 	if err != nil {
 		return fmt.Errorf("error while allocating ports on server %s for challenge %s: %s", host, challenge.Name, err.Error())
 	}

core/manager/utils.go

@@ -492,7 +492,7 @@ func UpdateOrCreateChallengeDbEntry(challEntry *database.Challenge, config cfg.B
 		availableServerHostname := core.LOCALHOST
 		if config.Challenge.Metadata.Type != core.STATIC_CHALLENGE_TYPE_NAME {
 			availableServer, _ := remoteManager.ServerQueue.GetNextAvailableInstance()
-			availableServerHostname = availableServer.Host
+			availableServerHostname = availableServer.Name
 		}
 		if config.Challenge.Metadata.Difficulty == "" {
 			log.Debug("Setting difficulty to default(medium)")
@@ -570,14 +570,7 @@ func UpdateOrCreateChallengeDbEntry(challEntry *database.Challenge, config cfg.B
 	}
 
 	if challEntry.ContainerId != "" {
-		var host string
-		if challEntry.ServerDeployed == core.LOCALHOST || challEntry.ServerDeployed == "" {
-			host = core.LOCALHOST
-		} else {
-			host = cfg.Cfg.AvailableServers[challEntry.ServerDeployed].Host
-		}
-
-		hostPorts, err := cache.GetContainerPortsOnHost(host, challEntry.ContainerId)
+		hostPorts, err := cache.GetContainerPortsOnHost(challEntry.ServerDeployed, challEntry.ContainerId)
 		if err != nil {
 			return fmt.Errorf("error while parsing host port for challenge %s : %s", challEntry.Name, err)
 		}

utils/datatypes.go

@@ -3,12 +3,11 @@ package utils
 import (
 	"errors"
 	"fmt"
+	"github.com/sdslabs/beastv4/core"
 	"strconv"
 	"strings"
 )
 
-const mappingDelimeter = ":"
-
 // From a list of strings generate a list containing only unique strings
 // from the list.
 func GetUniqueStrings(list []string) []string {
@@ -48,7 +47,7 @@ func UInt32InList(a uint32, list []uint32) bool {
 // If the portMapping string is not valid, this returns an error.
 // The format of the port mapping is `PORT_FIRST:PORT_LAST`
 func ParsePortMapping(portMap string) (uint32, uint32, error) {
-	ports := strings.Split(portMap, mappingDelimeter)
+	ports := strings.Split(portMap, core.MappingDelimeter)
 
 	if len(ports) != 2 {
 		return 0, 0, errors.New("port mapping string is not valid")