Implemented health check route

$HOME/.beast/beast_authorized_keys

@@ -0,0 +1,5 @@
+
+command="/home/kartik/.beast/scripts/60ced46738be4be3655bc289164df893f12fe67390d541621d9dae65849519b0",environment="SSH_USER=2",no-agent-forwarding,no-port-forwarding,no-X11-forwarding hi 
+command="/home/kartik/.beast/scripts/8027ef0d2bdea68f0c2e04e2e5dbc4d649bfaba68d7baa2f7223a04cd6c41edf",environment="SSH_USER=4",no-agent-forwarding,no-port-forwarding,no-X11-forwarding ssh-rsa AAAAB3NzaC1y 
+command="/home/kartik/.beast/scripts/8027ef0d2bdea68f0c2e04e2e5dbc4d649bfaba68d7baa2f7223a04cd6c41edf",environment="SSH_USER=4",no-agent-forwarding,no-port-forwarding,no-X11-forwarding ssh-rsa AAAAB3NzaC1y 
+command="/home/kartik/.beast/scripts/169c1d8053fee22df8553f8eb0957cf8cad6dca68fd77c2adf5ca10fe94fb2af",environment="SSH_USER=6",no-agent-forwarding,no-port-forwarding,no-X11-forwarding ssh-rsa AAAAB3NzaC1y 

_examples/simple/beast.toml

@@ -22,3 +22,6 @@ apt_deps = ["gcc", "socat"]
 setup_scripts = ["setup.sh"]
 run_cmd = "socat tcp-l:10005,fork,reuseaddr exec:./pwn"
 ports = [10005]
+
+[resource]
+memory_limit = 67108864

_examples/web-php/beast.toml

@@ -23,3 +23,6 @@ points = 20
 ports = [10002]
 web_root = "challenge"
 default_port = 10002
+
+[resource]
+memory_limit = 6710886

api/router.go

@@ -74,6 +74,7 @@ func initGinRouter() *gin.Engine {
 			statusGroup.GET("/challenge/:name", challengeStatusHandler)
 			statusGroup.GET("/all", statusHandler)
 			statusGroup.GET("/all/:filter", statusHandler)
+			statusGroup.GET("/health/:name", ChallengeHealthHandler)
 		}
 
 		// Info route group
@@ -92,6 +93,7 @@ func initGinRouter() *gin.Engine {
 			infoGroup.GET("/tags", tagHandler)
 			infoGroup.GET("/hint/:hintID", hintHandler)
 			infoGroup.POST("/hint/:hintID", hintHandler)
+
 		}
 
 		// Notification route group

api/status.go

@@ -4,7 +4,7 @@ import (
 	"fmt"
 	"net/http"
 	"time"
-
+	"github.com/sdslabs/beastv4/pkg/cr"
 	"github.com/gin-gonic/gin"
 	"github.com/sdslabs/beastv4/core"
 	"github.com/sdslabs/beastv4/core/database"
@@ -109,3 +109,46 @@ func statusHandler(c *gin.Context) {
 		c.JSON(http.StatusOK, resp)
 	}
 }
+
+func ChallengeHealthHandler(c *gin.Context) {
+    name := c.Param("name")
+	if name == "" {
+		c.JSON(http.StatusBadRequest, HTTPPlainResp{
+			Message: "Name of the challenge is a required parameter to process request.",
+		})
+		return
+	}
+
+	challenge, err := database.QueryChallengeEntries("name", name)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, HTTPPlainResp{
+			Message: "DATABASE ERROR while processing the request.",
+		})
+		return
+	}
+	chall := challenge[0]
+
+	memory_usage, cpu_perc, err := cr.GetContainerStats(chall.ContainerId)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, HTTPPlainResp{
+			Message: "error while fetching container stats",
+		})
+		return
+	}
+
+	memory_limit, _, err := cr.GetContainerLimits(chall.ContainerId)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, HTTPPlainResp{
+			Message: "error while fetching container limits",
+		})
+		return
+	}
+	memory_perc := (float64(memory_usage) / float64(memory_limit)) * 100
+
+    c.JSON(http.StatusOK, gin.H{
+        "status": "healthy",
+        "memory(bytes)": memory_usage,
+		"memory_perc": memory_perc,
+		"cpu_perc": cpu_perc,
+    })
+}

cmd/beast/healthprobe.go

@@ -4,6 +4,7 @@ import (
 	"github.com/sdslabs/beastv4/core/config"
 	"github.com/sdslabs/beastv4/core/manager"
 	"github.com/spf13/cobra"
+	"github.com/sdslabs/beastv4/core/database"
 )
 
 var healthProbeCmd = &cobra.Command{
@@ -13,7 +14,8 @@ var healthProbeCmd = &cobra.Command{
 
 	Run: func(cmd *cobra.Command, args []string) {
 		config.InitConfig()
+		database.Init()
 
-		go manager.BeastHeathCheckProber(config.Cfg.TickerFrequency)
+		manager.BeastHeathCheckProber(config.Cfg.TickerFrequency)
 	},
 }

core/manager/health_check.go

@@ -19,6 +19,45 @@ import (
 
 var HEALTH_CHECKER = false
 
+func RestartChallenge(chall *database.Challenge) error {
+	const maxRetries = 3;
+	success := false
+	var newContainerId string
+	var restartErr error
+
+	for i:=0; i<maxRetries; i++ {
+		log.Infof("Restart attempt %d/%d for challenge '%s'", i+1, maxRetries, chall.Name)
+		newContainerId ,restartErr = cr.RestartContainer(chall.ContainerId, chall.Name)
+		if restartErr==nil {
+			success = true
+			break
+		}
+	}
+	if (success) {
+		msg := fmt.Sprintf("Container restarted successfully! : %s",chall.Name)
+		log.WithFields(log.Fields{
+		"ChallName": chall.Name,
+		}).Info(msg)
+
+		chall.ContainerId = newContainerId
+
+		err := database.UpdateChallenge(chall, map[string]interface{}{"container_id": newContainerId})
+		if err != nil {
+			return fmt.Errorf("CRITICAL: created new container but failed to update database: %w", err)
+		}
+
+	} else
+	 {
+		msg := fmt.Sprintf("Container restart failed! : %s",chall.Name)
+		log.WithFields(log.Fields{
+		"ChallName": chall.Name,
+		}).Error(msg)
+	}
+
+    return nil
+
+}
+
 // Check for static challenegs' assets to be present on staging server.
 // At the time of writing, Beast deploys assets to localhost only.
 // So it will check only on localhost
@@ -57,6 +96,35 @@ func containerProber(chall database.Challenge) error {
 			return err
 		}
 	}
+
+	memory_usage, _, err := cr.GetContainerStats(chall.ContainerId)
+	if err != nil {
+		err = fmt.Errorf("error while fetching container stats with id %s ", chall.ContainerId)
+		return err
+	}
+
+	memory_limit, _, err := cr.GetContainerLimits(chall.ContainerId)
+	if err != nil {
+		err = fmt.Errorf("error while fetching container limits with id %s ", chall.ContainerId)
+		return err
+	}
+	memory_perc := (float64(memory_usage) / float64(memory_limit)) * 100
+
+	log.WithFields(log.Fields{
+		"ChallName": chall.Name,
+	}).Infof("Container ID: %s", chall.ContainerId)
+	log.WithFields(log.Fields{
+		"ChallName": chall.Name,
+	}).Infof("Memory Usage: %.2f MB", float64(memory_usage)/1024/1024)
+	log.WithFields(log.Fields{
+		"ChallName": chall.Name,
+	}).Infof("Memory Usage Percentage: %.2f%%", memory_perc)
+
+	if memory_perc > 100 {
+		err := fmt.Errorf("Memory limit exceeded for container with id %s ", chall.ContainerId)
+		return err
+	}
+
 	return nil
 }
 
@@ -101,7 +169,7 @@ func ChallengesHealthProber(waitTime int) {
 				}
 				err = containerProber(chall)
 				if err != nil {
-					msg := fmt.Sprintf("CONTAINER HEALTH CHECK %s: %s : %s", result, chall.Name, err)
+ 					msg := fmt.Sprintf("CONTAINER HEALTH CHECK %s: %s : %s", result, chall.Name, err)
 					log.WithFields(log.Fields{
 						"ChallName": chall.Name,
 					}).Error(msg)

pkg/cr/containers.go

@@ -4,15 +4,15 @@ import (
 	"fmt"
 	"io/ioutil"
 	"strconv"
-
+	"encoding/json"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/filters"
 	"github.com/docker/docker/api/types/mount"
 	"github.com/docker/docker/client"
 	"github.com/docker/go-connections/nat"
 	"github.com/sdslabs/beastv4/pkg/defaults"
-
+	"time"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/net/context"
 )
@@ -279,3 +279,137 @@ func CommitContainer(containerId string) (string, error) {
 
 	return commitResp.ID, nil
 }
+
+func GetContainerStats(containerId string) (int64, float64, error) {
+	ctx := context.Background()
+	var data1 types.StatsJSON
+	var data2 types.StatsJSON
+	cli, err := client.NewClientWithOpts(client.FromEnv)
+	if err != nil {
+		log.Error("Failed to connect to docker sdk")
+		return 0, 0, err
+	}
+	defer cli.Close()
+
+	stats, err := cli.ContainerStats(ctx, containerId, false)
+	if err != nil {
+		log.Error("Failed to fetch container stats : ", containerId)
+		return 0, 0, err
+	}
+	time.Sleep(1 * time.Second)
+	stats2, err := cli.ContainerStats(ctx, containerId, false)
+	if err != nil {
+		log.Error("Failed to fetch container stats : ", containerId)
+		return 0, 0, err
+	}
+	defer stats.Body.Close()
+
+	if err := json.NewDecoder(stats.Body).Decode(&data1); err != nil {
+		return 0, 0, err
+	}
+	if err := json.NewDecoder(stats2.Body).Decode(&data2); err != nil {
+		return 0, 0, err
+	}
+
+	memoryUsage := data2.MemoryStats.Usage
+
+	cpuDelta := data2.CPUStats.CPUUsage.TotalUsage - data1.CPUStats.CPUUsage.TotalUsage
+	systemDelta := data2.CPUStats.SystemUsage - data1.CPUStats.SystemUsage
+	numCPUs := float64(len(data2.CPUStats.CPUUsage.PercpuUsage))
+	cpuPercent := 0.0
+	if systemDelta > 0 && cpuDelta > 0 {
+		cpuPercent = (float64(cpuDelta) / float64(systemDelta)) * numCPUs * 100
+	}
+
+	return int64(memoryUsage), cpuPercent, nil
+
+}
+
+func GetContainerLimits(containerId string) (int64, int64, error) {
+	ctx := context.Background()
+	cli, err := client.NewClientWithOpts(client.FromEnv)
+	if err != nil {
+		log.Error("Failed to connect to docker sdk")
+		return 0, 0, err
+	}
+	defer cli.Close()
+
+	stats, err := cli.ContainerInspect(ctx, containerId)
+	if err != nil {
+		log.Error("Failed to fetch container stats : ", containerId)
+		return 0, 0, err
+	}
+
+	memory_limit := stats.HostConfig.Resources.Memory
+	cpu_shares := stats.HostConfig.Resources.CPUShares
+
+	return memory_limit, cpu_shares, nil
+}
+
+func RestartContainer(containerId string, name string) (string,error) {
+	log.Infof("Attempting to restart container for challenge: %s (ID: %s)", name, containerId)
+
+	cli, err := client.NewClientWithOpts()
+    if err != nil {
+        return "", fmt.Errorf("could not create docker client: %w", err)
+    }
+
+	oldContainer, err := cli.ContainerInspect(context.Background(), containerId)
+    if err != nil {
+        return "", fmt.Errorf("could not inspect old container %s: %w", containerId, err)
+    }
+
+	err = StopAndRemoveContainer(containerId)
+	if err!=nil {
+		return "", err
+	}
+
+	mountsMap := make(map[string]string)
+    for _, m := range oldContainer.HostConfig.Mounts {
+        if m.Type == mount.TypeBind {
+            mountsMap[m.Source] = m.Target
+        }
+    }
+
+    var portMappings []PortMapping
+    for containerPort, hostBindings := range oldContainer.HostConfig.PortBindings {
+        if len(hostBindings) > 0 {
+            hostPortStr := hostBindings[0].HostPort
+            hostPort, _ := strconv.ParseUint(hostPortStr, 10, 32)
+            containerPortInt, _ := strconv.ParseUint(containerPort.Port(), 10, 32)
+
+            portMappings = append(portMappings, PortMapping{
+                HostPort:      uint32(hostPort),
+                ContainerPort: uint32(containerPortInt),
+            })
+        }
+    }
+
+    newName := fmt.Sprintf("%s-restarted-%d", name, time.Now().Unix())
+
+    config := &CreateContainerConfig{
+        ImageId:          oldContainer.Config.Image,
+        ContainerName:    newName,
+        ContainerEnv:     oldContainer.Config.Env,
+        ContainerNetwork: string(oldContainer.HostConfig.NetworkMode),
+        MountsMap:        mountsMap,
+        PortMapping:      portMappings,
+        Memory:           oldContainer.HostConfig.Resources.Memory,
+        CPUShares:        oldContainer.HostConfig.Resources.CPUShares,
+    }
+    if oldContainer.HostConfig.Resources.PidsLimit != nil {
+        config.PidsLimit = *oldContainer.HostConfig.Resources.PidsLimit
+    }
+
+    newContainerId, err := CreateContainerFromImage(config)
+    if err != nil {
+        return "", fmt.Errorf("failed to create new container: %w", err)
+    }
+    log.Infof("Successfully created new container %s for challenge %s", newContainerId, name)
+
+    return newContainerId, nil
+}
+
+
+
+