check finish within myqueue (refs #448 #436)

framework/fuzzer/Fuzzer.py

@@ -159,46 +159,23 @@ def __init__(self, options):
     def __iter__(self):
 	return self
 
-    def process(self):
+    def next(self):
 	# http://bugs.python.org/issue1360
-	prio, item = self.results_queue.get(True, 365 * 24 * 60 * 60)
-
+	prio, res = self.results_queue.get(True, 365 * 24 * 60 * 60)
 	self.results_queue.task_done()
 
-        if item is None:
-            return None
-
-        if item.type == FuzzResult.result:
-	    if item.is_processable: self.genReq.stats.processed.inc()
-	    self.genReq.stats.pending_fuzz.dec()
-	    if not item.is_visible: self.genReq.stats.filtered.inc()
-        elif item.type == FuzzResult.endseed:
-	    self.genReq.stats.pending_seeds.dec()
-        elif item.type == FuzzResult.error:
-	    raise item.exception
-
-	# check if we are done. If so, send None to everyone so they can stop nicely
-	if item and self.genReq.stats.pending_fuzz() == 0 and self.genReq.stats.pending_seeds() == 0:
-	    self.qmanager.stop()
-
-	return item
-
-    def next(self):
-	# ignore end seed marks and not processable items
-	res = self.process()
-	while res and (not res.is_processable or res.type == FuzzResult.cancel or res.type == FuzzResult.endseed):
-
-	    res = self.process()
-
 	# done! (None sent has gone through all queues).
 	if not res:
+	    self.qmanager.stop()
 	    self.genReq.stats.mark_end()
 
             if self.printer:
                 self.printer.footer(self.genReq.stats)
 
 	    if self.output_fn: self.output_fn.close()
 	    raise StopIteration
+        elif res.type == FuzzResult.error:
+	    raise res.exception
 
 	# Save results?
 	if res and self.output_fn: 

framework/utils/myqueue.py

@@ -85,7 +85,11 @@ def _throw(self, e):
 
     def get_stats(self):
         return {self.get_name(): self.qsize()}
-
+
+    def _check_finish(self):
+	if self.stats.pending_fuzz() == 0 and self.stats.pending_seeds() == 0:
+	    self.send_last(None)
+
     def run(self):
 	cancelling = False
 
@@ -94,28 +98,45 @@ def run(self):
 
 	    try:
                 if item == None:
-                    self.send_last(None)
-                    if not cancelling: self.qout_join()
+                    if self.type != FuzzQueue.last:
+                        self.send_last(None)
+                        if not cancelling: self.qout_join()
                     self.task_done()
                     break
                 elif cancelling:
                     self.task_done()
                     continue
                 elif item.type == FuzzResult.endseed:
-                    self.send_last(item)
+                    if self.type == FuzzQueue.last:
+                        self.stats.pending_seeds.dec()
+                        self._check_finish()
+                    else:
+                        self.send_last(item)
                     self.task_done()
                     continue
-                elif item.type in [FuzzResult.error, FuzzResult.cancel]:
-                    cancelling = True if item.type == FuzzResult.cancel else False
+                elif item.type == FuzzResult.error:
                     self.send_first(item)
                     self.task_done()
                     continue
+                elif item.type == FuzzResult.cancel:
+                    cancelling = True
+                    if self.type != FuzzQueue.last: self.send_first(item)
+                    self.task_done()
+                    continue
                 elif not item.is_processable:
                     self.send(item)
                     self.task_done()
                     continue
 
 		self.process(prio, item)
+
+                if self.type == FuzzQueue.last:
+                    if item.type == FuzzResult.result:
+                        if item.is_processable: self.stats.processed.inc()
+                        self.stats.pending_fuzz.dec()
+                        if not item.is_visible: self.stats.filtered.inc()
+                    self._check_finish()
+
 		self.task_done()
 	    except Exception, e:
 		self.task_done()