diff --git a/_examples/service/beast.toml b/_examples/service/beast.toml index e95a4a2d..b5fafcda 100644 --- a/_examples/service/beast.toml +++ b/_examples/service/beast.toml @@ -24,3 +24,6 @@ apt_deps = ["gcc", "socat"] setup_scripts = ["setup.sh"] service_path = "pwn" ports = [10004] + +[resource] +memory_limit = 6710886 diff --git a/_examples/simple/beast.toml b/_examples/simple/beast.toml index 661245f3..a6979552 100644 --- a/_examples/simple/beast.toml +++ b/_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 diff --git a/_examples/web-php/beast.toml b/_examples/web-php/beast.toml index 3caf7425..d4d37443 100644 --- a/_examples/web-php/beast.toml +++ b/_examples/web-php/beast.toml @@ -23,3 +23,6 @@ points = 20 ports = [10002] web_root = "challenge" default_port = 10002 + +[resource] +memory_limit = 6710886 diff --git a/api/router.go b/api/router.go index cd98d4e8..95aabb0d 100644 --- a/api/router.go +++ b/api/router.go @@ -49,8 +49,12 @@ func initGinRouter() *gin.Engine { router.GET("/api/info/competition-info", competitionInfoHandler) router.GET("/api/info/download", serveAssets) + router.GET("/api/health/:name", adminAuthorize, ChallengeHealthHandler) + + // API routes group - apiGroup := router.Group("/api", authorize) + apiGroup := router.Group("/api", +) { // Deploy route group manageGroup := apiGroup.Group("/manage", managerAuthorize) @@ -92,6 +96,7 @@ func initGinRouter() *gin.Engine { infoGroup.GET("/tags", tagHandler) infoGroup.GET("/hint/:hintID", hintHandler) infoGroup.POST("/hint/:hintID", hintHandler) + } // Notification route group diff --git a/api/status.go b/api/status.go index 54ae4943..6e9acaab 100644 --- a/api/status.go +++ b/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, + }) +} \ No newline at end of file diff --git a/cmd/beast/healthprobe.go b/cmd/beast/healthprobe.go index 65bc5691..b027f934 100644 --- a/cmd/beast/healthprobe.go +++ b/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) }, } diff --git a/core/manager/health_check.go b/core/manager/health_check.go index 767dcd10..35382b26 100644 --- a/core/manager/health_check.go +++ b/core/manager/health_check.go @@ -57,6 +57,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 +130,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) diff --git a/pkg/cr/containers.go b/pkg/cr/containers.go index 777e8f4c..577d9674 100644 --- a/pkg/cr/containers.go +++ b/pkg/cr/containers.go @@ -4,7 +4,7 @@ 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" @@ -12,7 +12,7 @@ import ( "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,73 @@ 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 +} + + + +