Properly hold handler until jobs done
continuous-integration/drone/push Build is passing Details

This commit is contained in:
Tyler 2021-03-07 21:48:22 -05:00
parent 37851c8ee2
commit 3111c563f2
1 changed files with 24 additions and 13 deletions

37
main.go
View File

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