Merge pull request #18 from aau-network-security/develop

Close lab feature and cancel queue feature

internal/daemon/adminEvents.go

@@ -113,6 +113,12 @@ func (d *daemon) newEvent(c *gin.Context) {
 			return
 		}
 
+		if req.Type == int32(TypeBeginner) {
+			log.Warn().Msg("admin user tried to create an event with beginner type")
+			c.JSON(http.StatusBadRequest, APIResponse{Status: "Beginner events are no longer supported"})
+			return
+		}
+
 		uniqueExercisesList := removeDuplicates(req.ExerciseTags)
 
 		exClientResp, err := d.exClient.GetExerciseByTags(ctx, &eproto.GetExerciseByTagsRequest{Tag: uniqueExercisesList})

internal/daemon/eventLabs.go

@@ -22,6 +22,11 @@ func (d *daemon) eventLabsSubrouter(r *gin.RouterGroup) {
 	labs.GET("/hosts", d.getHostsInLab)
 	labs.GET("/vpnconf/:id", d.getVpnConf)
 	labs.GET("/resetlab", d.resetLab)
+	labs.DELETE("/close", d.closeLab)
+
+	queue := labs.Group("/queue")
+	queue.Use(d.eventAuthMiddleware())
+	queue.DELETE("/cancel", d.cancelLabConfigurationRequest)
 }
 
 type LabRequest struct {
@@ -101,10 +106,12 @@ func (d *daemon) configureLab(c *gin.Context) {
 		return
 	}
 
-	if err := d.agentPool.createLabForEvent(ctx, req.IsVpn, event, d.eventpool); err != nil {
-		log.Error().Err(err).Msg("Error creating lab")
-		c.JSON(http.StatusInternalServerError, APIResponse{Status: "error when creating lab, please try again..."})
-		return
+	if (req.IsVpn && event.TeamsWaitingForVpnLabs.Len() == 0 && len(event.UnassignedVpnLabs) == 0) || (!req.IsVpn && event.TeamsWaitingForBrowserLabs.Len() == 0 && len(event.UnassignedBrowserLabs) == 0) {
+		if err := d.agentPool.createLabForEvent(ctx, req.IsVpn, event, d.eventpool); err != nil {
+			log.Error().Err(err).Msg("Error creating lab")
+			c.JSON(http.StatusInternalServerError, APIResponse{Status: "error when creating lab, please try again..."})
+			return
+		}
 	}
 
 	//Add team to queue
@@ -160,6 +167,89 @@ func (d *daemon) getLabInfo(c *gin.Context) {
 	c.JSON(http.StatusOK, APIResponse{Status: "OK", TeamLab: labResponse})
 }
 
+// Closes the lab for the requesting team
+func (d *daemon) closeLab(c *gin.Context) {
+	teamClaims := unpackTeamClaims(c)
+
+	event, err := d.eventpool.GetEvent(teamClaims.EventTag)
+	if err != nil {
+		log.Error().Err(err).Msg("could not find event in event pool")
+		c.JSON(http.StatusBadRequest, APIResponse{Status: "event for team is not currently running"})
+		return
+	}
+
+	team, err := event.GetTeam(teamClaims.Username)
+	if err != nil {
+		log.Error().Err(err).Msg("could not find team for event")
+		c.JSON(http.StatusBadRequest, APIResponse{Status: "could not find team for event"})
+		return
+	}
+
+	team.M.Lock()
+	defer team.M.Unlock()
+
+	if team.Lab == nil {
+		log.Debug().Str("team", team.Username).Msg("lab not found for team")
+		c.JSON(http.StatusNotFound, APIResponse{Status: "lab not found"})
+		return
+	}
+
+	defer saveState(d.eventpool, d.conf.StatePath)
+
+	event.M.Lock()
+	delete(event.Labs, team.Lab.LabInfo.Tag)
+	event.M.Unlock()
+	if team.Lab.Conn != nil {
+		if err := team.Lab.close(); err != nil {
+			log.Error().Err(err).Str("team", team.Username).Msg("Error closing lab for team")
+		}
+	}
+	team.Lab = nil
+	sendCommandToTeam(team, updateTeam)
+
+	c.JSON(http.StatusOK, APIResponse{Status: "OK"})
+}
+
+func (d *daemon) cancelLabConfigurationRequest(c *gin.Context) {
+	teamClaims := unpackTeamClaims(c)
+
+	event, err := d.eventpool.GetEvent(teamClaims.EventTag)
+	if err != nil {
+		log.Error().Err(err).Msg("could not find event in event pool")
+		c.JSON(http.StatusBadRequest, APIResponse{Status: "event for team is not currently running"})
+		return
+	}
+
+	team, err := event.GetTeam(teamClaims.Username)
+	if err != nil {
+		log.Error().Err(err).Msg("could not find team for event")
+		c.JSON(http.StatusBadRequest, APIResponse{Status: "could not find team for event"})
+		return
+	}
+
+	team.M.Lock()
+	defer team.M.Unlock()
+
+	if team.Status == InQueue {
+		team.Status = Idle
+		if team.QueueElement != nil {
+			event.TeamsWaitingForBrowserLabs.Remove(team.QueueElement)
+		}
+		sendCommandToTeam(team, updateTeam)
+		c.JSON(http.StatusOK, APIResponse{Status: "OK"})
+		return
+	}
+
+	if team.Status == WaitingForLab {
+		team.Status = Idle
+		sendCommandToTeam(team, updateTeam)
+		c.JSON(http.StatusOK, APIResponse{Status: "OK"})
+		return
+	}
+
+	c.JSON(http.StatusBadRequest, APIResponse{Status: "team is not in queue"})
+}
+
 // Returns current hosts running in their lab
 // It is returned as a list of string in format "<IP> \t <DNS>"
 func (d *daemon) getHostsInLab(c *gin.Context) {

internal/daemon/eventpool.go

@@ -145,6 +145,7 @@ Short minded fix is currently inserting a 1 milisecond delay...
 func (event *Event) startQueueHandlers(eventPool *EventPool, statePath string, labExpiry time.Duration) {
 	browserQueueHandler := func() {
 		log.Debug().Msg("Waiting for teams to enter browser lab queue")
+	Outer:
 		for {
 			time.Sleep(1 * time.Millisecond)
 			e := event.TeamsWaitingForBrowserLabs.Front()
@@ -169,6 +170,32 @@ func (event *Event) startQueueHandlers(eventPool *EventPool, statePath string, l
 				return
 			}
 
+			team.M.RLock()
+			TeamStatus := team.Status
+			team.M.RUnlock()
+			if TeamStatus != WaitingForLab {
+				for {
+					e := event.TeamsWaitingForBrowserLabs.Front()
+					// If no more teams are waiting for labs, close the lab, remove it from the event and continue
+					if e == nil {
+						log.Info().Msg("No more teams waiting for labs closing abunadant lab")
+						event.M.Lock()
+						delete(event.Labs, lab.LabInfo.Tag)
+						event.M.Unlock()
+						saveState(eventPool, statePath)
+						if err := lab.close(); err != nil {
+							log.Error().Err(err).Msg("error closing lab no longer needed")
+						}
+						continue Outer
+					}
+					log.Debug().Msg("New team pulled from browser queue")
+					event.TeamsWaitingForBrowserLabs.Remove(e)
+
+					team = e.Value.(*Team)
+					break
+				}
+			}
+
 			team.M.Lock()
 			lab.IsAssigned = true
 			lab.ExpiresAtTime = time.Now().Add(labExpiry * time.Minute)
@@ -184,6 +211,7 @@ func (event *Event) startQueueHandlers(eventPool *EventPool, statePath string, l
 
 	vpnQueueHandler := func() {
 		log.Debug().Msg("Waiting for team to enter vpn lab queue")
+	Outer:
 		for {
 			time.Sleep(1 * time.Millisecond)
 			e := event.TeamsWaitingForVpnLabs.Front()
@@ -206,6 +234,32 @@ func (event *Event) startQueueHandlers(eventPool *EventPool, statePath string, l
 				return
 			}
 
+			team.M.RLock()
+			TeamStatus := team.Status
+			team.M.RUnlock()
+			if TeamStatus != WaitingForLab {
+				for {
+					e := event.TeamsWaitingForVpnLabs.Front()
+					// If no more teams are waiting for labs, close the lab, remove it from the event and continue
+					if e == nil {
+						log.Info().Msg("No more teams waiting for labs closing abunadant lab")
+						event.M.Lock()
+						delete(event.Labs, lab.LabInfo.Tag)
+						event.M.Unlock()
+						saveState(eventPool, statePath)
+						if err := lab.close(); err != nil {
+							log.Error().Err(err).Msg("error closing lab no longer needed")
+						}
+						continue Outer
+					}
+					log.Debug().Msg("New team pulled from browser queue")
+					event.TeamsWaitingForVpnLabs.Remove(e)
+
+					team = e.Value.(*Team)
+					break
+				}
+			}
+
 			team.M.Lock()
 			lab.IsAssigned = true
 			lab.ExpiresAtTime = time.Now().Add(labExpiry * time.Minute)