From 3111c563f21463fa27df2fed0254b8d61df993be Mon Sep 17 00:00:00 2001
From: Tyler <tystuyfzand@gmail.com>
Date: Sun, 7 Mar 2021 21:48:22 -0500
Subject: [PATCH] Properly hold handler until jobs done

---
 main.go | 37 ++++++++++++++++++++++++-------------
 1 file changed, 24 insertions(+), 13 deletions(-)

diff --git a/main.go b/main.go
index 6adf454..19d55ca 100644
--- a/main.go
+++ b/main.go
@@ -109,6 +109,10 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 		contentType = contentType[0:idx]
 	}
 
+	var res interface{}
+
+	wg := &sync.WaitGroup{}
+
 	switch contentType {
 	case "multipart/form-data":
 		if r.MultipartForm == nil {
@@ -119,15 +123,13 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 
 		fileCount := 0
 
-		res := multiResponse{
+		response := multiResponse{
 			Success: true,
 			Files:   make(map[string]yara.MatchRules),
 		}
 
 		fileLock := &sync.Mutex{}
 
-		wg := &sync.WaitGroup{}
-
 		for _, files := range r.MultipartForm.File {
 			// Append files
 			for _, file := range files {
@@ -145,7 +147,7 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 					Data: f,
 					Callback: func(m yara.MatchRules, err error) {
 						fileLock.Lock()
-						res.Files[file.Filename] = m
+						response.Files[file.Filename] = m
 						fileLock.Unlock()
 
 						wg.Done()
@@ -156,31 +158,32 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 			}
 		}
 
-		wg.Wait()
-
 		log.WithField("count", fileCount).Debug("Waiting for jobs to finish")
 
-		log.WithField("matches", res).Debug("Matched rules")
-		err := json.NewEncoder(w).Encode(res)
+		wg.Wait()
 
-		if err != nil {
-			log.WithError(err).Error("Unable to send response")
-		}
+		log.WithField("matches", res).Debug("Matched rules")
+
+		res = response
 	default:
+		wg.Add(1)
+
 		job := &Job{
 			Data: r.Body,
 			Callback: func(m yara.MatchRules, err error) {
 				log.WithField("match", m).Debug("Matched rules")
 
 				if err != nil {
-					err = json.NewEncoder(w).Encode(response{Success: false, Error: err})
+					res = response{Success: false, Error: err}
 				} else {
-					err = json.NewEncoder(w).Encode(response{Success: true, MatchedRules: m})
+					res = response{Success: true, MatchedRules: m}
 				}
 
 				if err != nil {
 					log.WithError(err).Error("Unable to send response")
 				}
+
+				wg.Done()
 			},
 		}
 
@@ -188,6 +191,14 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 
 		jobChan <- job
 	}
+
+	wg.Wait()
+
+	err := json.NewEncoder(w).Encode(res)
+
+	if err != nil {
+		log.WithError(err).Error("Unable to send response")
+	}
 }
 
 // Load rules from the rules configuration option