Merge pull request #19 from aau-network-security/hotfix/labs-not-bein…

…g-queued Fixed labs not being set in queue
aau-network-security · Nov 16, 2024 · 6eb042a · 6eb042a
2 parents fc4fa61 + ed80d30
commit 6eb042a
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 90 deletions.
diff --git a/internal/daemon/eventLabs.go b/internal/daemon/eventLabs.go
@@ -106,12 +106,10 @@ func (d *daemon) configureLab(c *gin.Context) {
 		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
-		}
+	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
@@ -132,7 +130,7 @@ func (d *daemon) configureLab(c *gin.Context) {
 	if event.IsMaxLabsReached() {
 		broadCastCommandToEventTeams(event, updateEventInfo)
 	}
-	log.Info().Str("username", team.Username).Msg("putting team into queue for vpn lab")
+	log.Info().Str("username", team.Username).Msg("putting team into queue for browser lab")
 	queueElement := event.TeamsWaitingForBrowserLabs.PushBack(team)
 	team.QueueElement = queueElement
 	sendCommandToTeam(team, updateTeam)
@@ -234,19 +232,13 @@ func (d *daemon) cancelLabConfigurationRequest(c *gin.Context) {
 		team.Status = Idle
 		if team.QueueElement != nil {
 			event.TeamsWaitingForBrowserLabs.Remove(team.QueueElement)
+			event.TeamsWaitingForVpnLabs.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"})
 }
 

diff --git a/internal/daemon/eventpool.go b/internal/daemon/eventpool.go
@@ -145,23 +145,8 @@ 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()
-			if e == nil {
-				continue
-			}
-			log.Debug().Msg("team pulled from browser queue")
-			event.TeamsWaitingForBrowserLabs.Remove(e)
 
-			team := e.Value.(*Team)
-			team.M.Lock()
-			team.Status = WaitingForLab
-			team.M.Unlock()
-
-			// TODO Make agent send return even if it fails to create the lab
-			// TODO Make unblocking and implement cancel
+		for {
 			lab, ok := <-event.UnassignedBrowserLabs
 			if ok {
 				log.Debug().Msgf("pulled lab from browser queue: %v", lab)
@@ -170,31 +155,22 @@ 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
+			e := event.TeamsWaitingForBrowserLabs.Front()
+			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
 			}
+			log.Debug().Msg("team pulled from browser queue")
+			event.TeamsWaitingForBrowserLabs.Remove(e)
+
+			team := e.Value.(*Team)
 
 			team.M.Lock()
 			lab.IsAssigned = true
@@ -211,21 +187,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()
-			if e == nil {
-				continue
-			}
-			event.TeamsWaitingForVpnLabs.Remove(e)
-			log.Debug().Msg("team pulled from vpn queue")
-
-			team := e.Value.(*Team)
-			team.M.Lock()
-			team.Status = WaitingForLab
-			team.M.Unlock()
-
 			lab, ok := <-event.UnassignedVpnLabs
 			if ok {
 				log.Debug().Msgf("pulled lab from vpn queue: %v", lab)
@@ -234,31 +196,23 @@ 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
+			e := event.TeamsWaitingForVpnLabs.Front()
+			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
 			}
+			event.TeamsWaitingForVpnLabs.Remove(e)
+			log.Debug().Msg("team pulled from vpn queue")
+
+			team := e.Value.(*Team)
 
 			team.M.Lock()
 			lab.IsAssigned = true