Add success field, matched rules output

main.go

@@ -31,6 +31,16 @@ var (
 	jobChan = make(chan *Job)
 )
 
+type response struct {
+	Success      bool            `json:"success"`
+	MatchedRules yara.MatchRules `json:"rules,omitempty"`
+}
+
+type multiResponse struct {
+	Success bool                       `json:"success"`
+	Files   map[string]yara.MatchRules `json:"files,omitempty"`
+}
+
 func main() {
 	viper.SetDefault("threads", runtime.NumCPU())
 	viper.SetDefault("rules", "pkg:github/Neo23x0/signature-base#yara")
@@ -106,22 +116,17 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 
 		wg := &sync.WaitGroup{}
 
-		results := make([]*yara.MatchRules, 0)
-
-		jobCallback := func(m *yara.MatchRules, err error) {
-			wg.Done()
-
-			if err != nil {
-				return
-			}
-
-			results = append(results, m)
-		}
-
 		log.WithField("contentType", contentType).Debug("Adding files from multipart form as jobs")
 
 		fileCount := 0
 
+		res := multiResponse{
+			Success: true,
+			Files:   make(map[string]yara.MatchRules),
+		}
+
+		fileLock := &sync.Mutex{}
+
 		for _, files := range r.MultipartForm.File {
 			// Append files
 			for _, file := range files {
@@ -137,7 +142,13 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 
 				job := &Job{
 					Data: f,
-					Callback: jobCallback,
+					Callback: func(m *yara.MatchRules, err error) {
+						fileLock.Lock()
+						res.Files[file.Filename] = *m
+						fileLock.Unlock()
+
+						wg.Done()
+					},
 				}
 
 				jobChan <- job
@@ -148,8 +159,8 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 
 		wg.Wait()
 
-		log.WithField("matches", w).Debug("Matched rules")
-		json.NewEncoder(w).Encode(results)
+		log.WithField("matches", res).Debug("Matched rules")
+		json.NewEncoder(w).Encode(res)
 	default:
 		job := &Job{
 			Data: r.Body,
@@ -160,7 +171,7 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 					w.WriteHeader(http.StatusInternalServerError)
 					w.Write([]byte(err.Error()))
 				} else {
-					json.NewEncoder(w).Encode(m)
+					json.NewEncoder(w).Encode(response{Success: true, MatchedRules: *m})
 				}
 			},
 		}