From 0cef7df735a46892871c5bf842b3f269c2fd85a3 Mon Sep 17 00:00:00 2001
From: filp <dobreira@gmail.com>
Date: Wed, 24 Apr 2013 16:49:10 +0100
Subject: [PATCH] Move output handling back up to Whoops\Run

---
 src/Whoops/Handler/JsonResponseHandler.php |  7 +--
 src/Whoops/Handler/PrettyPageHandler.php   |  3 --
 src/Whoops/Run.php                         | 63 ++++++++++++++++------
 3 files changed, 49 insertions(+), 24 deletions(-)

diff --git a/src/Whoops/Handler/JsonResponseHandler.php b/src/Whoops/Handler/JsonResponseHandler.php
index 2229de1..f34ea6b 100644
--- a/src/Whoops/Handler/JsonResponseHandler.php
+++ b/src/Whoops/Handler/JsonResponseHandler.php
@@ -8,7 +8,7 @@ namespace Whoops\Handler;
 use Whoops\Handler\Handler;
 
 /**
- * Catches an exception and converts it to a JSON 
+ * Catches an exception and converts it to a JSON
  * response. Additionally can also return exception
  * frames for consumption by an API.
  */
@@ -71,9 +71,6 @@ class JsonResponseHandler extends Handler
             return Handler::DONE;
         }
 
-        // Remove any previous output
-        ob_get_level() and ob_end_clean();
-
         $exception = $this->getException();
 
         $response = array(
@@ -84,7 +81,7 @@ class JsonResponseHandler extends Handler
                 'line'    => $exception->getLine()
             )
         );
-        
+
         if($this->addTraceToOutput()) {
             $inspector = $this->getInspector();
             $frames    = $inspector->getFrames();
diff --git a/src/Whoops/Handler/PrettyPageHandler.php b/src/Whoops/Handler/PrettyPageHandler.php
index a2ed1df..0d0d632 100644
--- a/src/Whoops/Handler/PrettyPageHandler.php
+++ b/src/Whoops/Handler/PrettyPageHandler.php
@@ -58,9 +58,6 @@ class PrettyPageHandler extends Handler
         if(php_sapi_name() === 'cli' && !isset($_ENV['whoops-test'])) {
             return Handler::DONE;
         }
-        
-        // Remove any previous output
-        ob_get_level() and ob_end_clean();
 
         // Get the 'pretty-template.php' template file
         // @todo: this can be made more dynamic &&|| cleaned-up
diff --git a/src/Whoops/Run.php b/src/Whoops/Run.php
index 9bd29cc..7ada7d4 100644
--- a/src/Whoops/Run.php
+++ b/src/Whoops/Run.php
@@ -20,7 +20,8 @@ class Run
     const SHUTDOWN_HANDLER  = 'handleShutdown';
 
     protected $isRegistered;
-    protected $allowQuit = true;
+    protected $allowQuit  = true;
+    protected $sendOutput = true;
 
     /**
      * @var DarnIt\Handler\HandlerInterface[]
@@ -125,7 +126,7 @@ class Run
     /**
      * Should Whoops allow Handlers to force the script to quit?
      * @param bool|num $exit
-     * @return bool|null
+     * @return bool
      */
     public function allowQuit($exit = null)
     {
@@ -136,6 +137,21 @@ class Run
         return $this->allowQuit = (bool) $exit;
     }
 
+    /**
+     * Should Whoops push output directly to the client?
+     * If this is false, output will be returned by handleException
+     * @param bool|num $send
+     * @return bool
+     */
+    public function writeToOutput($send = null)
+    {
+        if(func_num_args() == 0) {
+            return $this->sendOutput;
+        }
+
+        return $this->sendOutput = (bool) $exit;
+    }
+
     /**
      * Handles an exception, ultimately generating a Whoops error
      * page.
@@ -148,6 +164,11 @@ class Run
         // they were registered, and pass off the exception
         $inspector = $this->getInspector($exception);
 
+        // Capture output produced while handling the exception,
+        // we might want to send it straight away to the client,
+        // or return it silently.
+        ob_start();
+
         for($i = count($this->handlerStack) - 1; $i >= 0; $i--) {
             $handler = $this->handlerStack[$i];
 
@@ -157,21 +178,31 @@ class Run
 
             $handlerResponse = $handler->handle($exception);
 
-            if($handlerResponse === Handler::LAST_HANDLER) {
-                // The Handler has handled the exception in some way,
-                // or signals that no further handlers should be queried,
-                // but the script execution will continue
+            if(in_array($handlerResponse, array(Handler::LAST_HANDLER, Handler::QUIT))) {
+                // The Handler has handled the exception in some way, and
+                // wishes to quit execution (Handler::QUIT), or skip any
+                // other handlers (Handler::LAST_HANDLER). If $this->allowQuit
+                // is false, Handler::QUIT behaves like Handler::LAST_HANDLER
                 break;
-            } elseif($handlerResponse === Handler::QUIT) {
-                // The Handler has handled the exception in some way,
-                // and script execution should terminate, unless specifically
-                // disallowed, in which case the behavior is the same as
-                // Handler::LAST_HANDLER
-                if($this->allowQuit()) {
-                    exit;
-                } else {
-                    break;
-                }
+            }
+        }
+
+        $output = ob_get_clean();
+
+        // Handlers are done! Check if we got here because of Handler::QUIT
+        // ($handlerResponse will be the response from the last queried handler)
+        // and if so, try to quit execution.
+        if($this->allowQuit()) {
+            echo $output;
+            exit;
+        } else {
+            // If we're allowed to, send output generated by handlers directly
+            // to the output, otherwise, return it so that it may be used by
+            // the caller.
+            if($this->writeToOutput()) {
+                echo $output;
+            } else {
+                return $output;
             }
         }
     }