Merge pull request #30 from ldsec/memoryLeak

Fixed concurrent map memory leakage

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
 }
 
@@ -274,6 +281,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
 
 	}
@@ -323,6 +336,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
 }
@@ -378,6 +397,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
 }
 

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 {