Merge pull request #13 from aau-network-security/feature/agent-syncro…

…nization

Feature/agent syncronization

internal/agent/agent.go

@@ -126,14 +126,18 @@ func New(conf *Config) (*Agent, error) {
 	if err != nil {
 		log.Error().Err(err).Msg("error resuming state")
 		envPool = &env.EnvPool{
-			M:    &sync.RWMutex{},
-			Envs: make(map[string]*env.Environment),
+			M:            &sync.RWMutex{},
+			Envs:         make(map[string]*env.Environment),
+			StartingEnvs: make(map[string]bool),
+			ClosingEnvs:  make(map[string]bool),
 		}
 	}
 	if envPool == nil {
 		envPool = &env.EnvPool{
-			M:    &sync.RWMutex{},
-			Envs: make(map[string]*env.Environment),
+			M:            &sync.RWMutex{},
+			Envs:         make(map[string]*env.Environment),
+			StartingEnvs: make(map[string]bool),
+			ClosingEnvs:  make(map[string]bool),
 		}
 	}
 	// Creating agent struct

internal/agent/environment.go

@@ -29,6 +29,8 @@ var (
 // beginner events where the user would just need to press the connect button and a lab would be ready with all challenges running.
 func (a *Agent) CreateEnvironment(ctx context.Context, req *proto.CreatEnvRequest) (*proto.StatusResponse, error) {
 	// Env for event already exists, Do not start a new guac container
+	a.EnvPool.AddStartingEnv(req.EventTag)
+	defer a.EnvPool.RemoveStartingEnv(req.EventTag)
 	log.Debug().Msgf("got createEnv request: %v", req)
 
 	if a.EnvPool.DoesEnvExist(req.EventTag) {
@@ -176,6 +178,8 @@ func (a *Agent) CreateEnvironment(ctx context.Context, req *proto.CreatEnvReques
 
 // Closes environment and attached containers/vms, and removes the environment from the event pool
 func (a *Agent) CloseEnvironment(ctx context.Context, req *proto.CloseEnvRequest) (*proto.StatusResponse, error) {
+	a.EnvPool.AddClosingEnv(req.EventTag)
+	defer a.EnvPool.RemoveClosingEnv(req.EventTag)
 	env, err := a.EnvPool.GetEnv(req.EventTag)
 	if err != nil {
 		log.Error().Str("envTag", req.EventTag).Msg("error finding finding environment with tag")
@@ -266,6 +270,15 @@ func (a *Agent) AddExercisesToEnv(ctx context.Context, req *proto.ExerciseReques
 	return &proto.StatusResponse{Message: "OK"}, nil
 }
 
+// Lists currently running, starting and closing environments.
+func (a *Agent) ListEnvironments(ctx context.Context, req *proto.Empty) (*proto.ListEnvResponse, error) {
+	return &proto.ListEnvResponse{
+		EventTags:         a.EnvPool.GetEnvList(),
+		StartingEventTags: a.EnvPool.GetStartingEnvs(),
+		ClosingEventTags:  a.EnvPool.GetClosingEnvs(),
+	}, nil
+}
+
 func getVPNIP() (string, error) {
 	// Get VPN IP address from ip pool
 	ip, err := vpnIPPool.Get()

internal/agent/lab.go

@@ -17,12 +17,10 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-// TODO: Rethink func name as this should be the function that configures a lab for a user
+
 func (a *Agent) CreateLabForEnv(ctx context.Context, req *proto.CreateLabRequest) (*proto.StatusResponse, error) {
-	a.EnvPool.M.RLock()
-	env, ok := a.EnvPool.Envs[req.EventTag]
-	a.EnvPool.M.RUnlock()
-	if !ok {
+	env, err := a.EnvPool.GetEnv(req.EventTag)
+	if err != nil {
 		return nil, errors.New("environment for event does not exist")
 	}
 

internal/environment/envpool.go

@@ -80,3 +80,62 @@ func (ep *EnvPool) RemoveEnv(tag string) error {
 	delete(ep.Envs, tag)
 	return nil
 }
+
+func (ep *EnvPool) LockForFunc(function func()) {
+	ep.M.Lock()
+	defer ep.M.Unlock()
+
+	function()
+}
+
+func (ep *EnvPool) GetEnvList() (envList map[string]bool) {
+	ep.M.RLock()
+	defer ep.M.RUnlock()
+	envList = make(map[string]bool)
+	for eventTag := range ep.Envs {
+		envList[eventTag] = true
+	}
+	return
+}
+
+func (ep *EnvPool) GetStartingEnvs() map[string]bool {
+	ep.M.RLock()
+	defer ep.M.RUnlock()
+
+	return ep.StartingEnvs
+}
+
+func (ep *EnvPool) AddStartingEnv(eventTag string) {
+	ep.M.Lock()
+	defer ep.M.Unlock()
+
+	ep.StartingEnvs[eventTag] = true
+}
+
+func (ep *EnvPool) RemoveStartingEnv(eventTag string) {
+	ep.M.Lock()
+	defer ep.M.Unlock()
+
+	delete(ep.StartingEnvs, eventTag)
+}
+
+func (ep *EnvPool) GetClosingEnvs() map[string]bool {
+	ep.M.RLock()
+	defer ep.M.RUnlock()
+
+	return ep.ClosingEnvs
+}
+
+func (ep *EnvPool) AddClosingEnv(eventTag string) {
+	ep.M.Lock()
+	defer ep.M.Unlock()
+
+	ep.ClosingEnvs[eventTag] = true
+}
+
+func (ep *EnvPool) RemoveClosingEnv(eventTag string) {
+	ep.M.Lock()
+	defer ep.M.Unlock()
+
+	delete(ep.ClosingEnvs, eventTag)
+}

internal/environment/models.go

@@ -13,8 +13,11 @@ import (
 
 // General environment types
 type EnvPool struct {
-	M    *sync.RWMutex
-	Envs map[string]*Environment
+	M *sync.RWMutex
+	// Map of environments with eventTag as key
+	Envs         map[string]*Environment
+	StartingEnvs map[string]bool
+	ClosingEnvs  map[string]bool
 }
 
 type Environment struct {

internal/state/state.go

@@ -64,8 +64,10 @@ func ResumeState(vlib *virtual.VboxLibrary, workerPool worker.WorkerPool, stateP
 	}
 
 	envPool := &environment.EnvPool{
-		M:    &sync.RWMutex{},
-		Envs: make(map[string]*environment.Environment),
+		M:            &sync.RWMutex{},
+		Envs:         make(map[string]*environment.Environment),
+		StartingEnvs: make(map[string]bool),
+		ClosingEnvs:  make(map[string]bool),
 	}
 	for k, envState := range state.Environments {
 		env, err := convertEnvState(envState, vlib, workerPool)