From fc8130f064aff1e049f7cf7799a0b42247bf46d9 Mon Sep 17 00:00:00 2001
From: Joao Sa Sousa <jagsousa1992@gmail.com>
Date: Wed, 5 Feb 2020 12:57:04 +0100
Subject: [PATCH] Fixed concurrent map memory leakage

---
 services/service.go | 26 ++++++++++++++++++++++++++
 services/structs.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/services/service.go b/services/service.go
index edfda7b..ef305f9 100644
--- a/services/service.go
+++ b/services/service.go
@@ -164,6 +164,13 @@ func (s *Service) HandleSurveyKSRequest(skr *SurveyKSRequest) (network.Message,
 	}
 	surveyKS.TR.MapTR[KSTimeExec] = execTime
 	surveyKS.TR.MapTR[KSTimeCommunication] = communicationTime
+
+	// remove query from map
+	_, err = s.deleteSurveyKS(skr.SurveyID)
+	if err != nil {
+		return nil, xerrors.Errorf("%+v", err)
+	}
+
 	return &Result{Result: keySwitchingResult, TR: surveyKS.TR}, nil
 }
 
@@ -278,6 +285,12 @@ func (s *Service) HandleSurveyShuffleRequest(ssr *SurveyShuffleRequest) (network
 		surveyShuffle.TR.MapTR[KSTimeExec] = execTime
 		surveyShuffle.TR.MapTR[KSTimeCommunication] = communicationTime
 
+		// remove query from map
+		_, err = s.deleteSurveyShuffle(ssr.SurveyID)
+		if err != nil {
+			return nil, xerrors.Errorf("%+v", err)
+		}
+
 		return &Result{Result: libunlynx.CipherVector{keySwitchingResult[index]}, TR: surveyShuffle.TR}, nil
 
 	}
@@ -327,6 +340,12 @@ func (s *Service) HandleSurveyShuffleRequest(ssr *SurveyShuffleRequest) (network
 		return nil, xerrors.New("couldn't find this node in the roster")
 	}
 
+	// remove query from map
+	_, err = s.deleteSurveyShuffle(ssr.SurveyID)
+	if err != nil {
+		return nil, xerrors.Errorf("%+v", err)
+	}
+
 	return &Result{Result: libunlynx.CipherVector{keySwitchingResult[index]},
 		TR: surveyShuffle.TR}, nil
 }
@@ -382,6 +401,13 @@ func (s *Service) HandleSurveyAggRequest(sar *SurveyAggRequest) (network.Message
 	}
 	surveyAgg.TR.MapTR[KSTimeExec] = execTime
 	surveyAgg.TR.MapTR[KSTimeCommunication] = communicationTime
+
+	// remove query from map
+	_, err = s.deleteSurveyAgg(sar.SurveyID)
+	if err != nil {
+		return nil, xerrors.Errorf("%+v", err)
+	}
+
 	return &Result{Result: keySwitchingResult, TR: surveyAgg.TR}, nil
 }
 
diff --git a/services/structs.go b/services/structs.go
index 348d82f..51efe6b 100644
--- a/services/structs.go
+++ b/services/structs.go
@@ -159,6 +159,39 @@ type SurveyAggGenerated struct {
 	SurveyID SurveyID
 }
 
+func (s *Service) deleteSurveyKS(sid SurveyID) (SurveyKS, error) {
+	surv, err := s.MapSurveyKS.Remove(string(sid))
+	if err != nil {
+		return SurveyKS{}, errors.New("Error" + err.Error() + "while deleting surveyID: " + string(sid))
+	}
+	if surv == nil {
+		return SurveyKS{}, errors.New("No entry in map with surveyID: " + string(sid))
+	}
+	return surv.(SurveyKS), nil
+}
+
+func (s *Service) deleteSurveyShuffle(sid SurveyID) (SurveyShuffle, error) {
+	surv, err := s.MapSurveyShuffle.Remove(string(sid))
+	if err != nil {
+		return SurveyShuffle{}, errors.New("Error" + err.Error() + "while deleting surveyID: " + string(sid))
+	}
+	if surv == nil {
+		return SurveyShuffle{}, errors.New("No entry in map with surveyID: " + string(sid))
+	}
+	return surv.(SurveyShuffle), nil
+}
+
+func (s *Service) deleteSurveyAgg(sid SurveyID) (SurveyAgg, error) {
+	surv, err := s.MapSurveyAgg.Remove(string(sid))
+	if err != nil {
+		return SurveyAgg{}, errors.New("Error" + err.Error() + "while deleting surveyID: " + string(sid))
+	}
+	if surv == nil {
+		return SurveyAgg{}, errors.New("No entry in map with surveyID: " + string(sid))
+	}
+	return surv.(SurveyAgg), nil
+}
+
 func (s *Service) getSurveyKS(sid SurveyID) (SurveyKS, error) {
 	surv, err := s.MapSurveyKS.Get(string(sid))
 	if err != nil {