diff --git a/main.go b/main.go index 4263dba..38f2353 100644 --- a/main.go +++ b/main.go @@ -109,14 +109,14 @@ func scanHandler(w http.ResponseWriter, r *http.Request) { contentType = contentType[0:idx] } + wg := &sync.WaitGroup{} + switch contentType { case "multipart/form-data": if r.MultipartForm == nil { r.ParseMultipartForm(32 << 20) } - wg := &sync.WaitGroup{} - log.WithField("contentType", contentType).Debug("Adding files from multipart form as jobs") fileCount := 0 @@ -156,6 +156,8 @@ 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") @@ -164,6 +166,8 @@ func scanHandler(w http.ResponseWriter, r *http.Request) { job := &Job{ Data: r.Body, Callback: func(m yara.MatchRules, err error) { + wg.Add(1) + log.WithField("match", m).Debug("Matched rules") if err != nil { @@ -177,6 +181,8 @@ func scanHandler(w http.ResponseWriter, r *http.Request) { log.WithField("contentType", contentType).Debug("Scanning contents of body") jobChan <- job + + wg.Wait() } }