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