From 7753d6eda2d08a3e8fd05f9773a7c6f38199e1e1 Mon Sep 17 00:00:00 2001
From: Vinnie Falco <vinnie.falco@gmail.com>
Date: Wed, 28 Jun 2017 04:01:33 -0700
Subject: [PATCH] Add parser::on_header to set a callback

---
 CHANGELOG.md                  |  1 +
 include/beast/http/parser.hpp | 29 ++++++++++++++++++++++-
 test/http/basic_parser.cpp    |  1 +
 test/http/parser.cpp          | 44 +++++++++++++++++++++++++++++++++++
 4 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c94c8a00..40fa2b43 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@ Version 70:
 
 * Serialize in one step when possible
 * Add basic_parser header and body limits
+* Add parser::on_header to set a callback
 
 API Changes:
 
diff --git a/include/beast/http/parser.hpp b/include/beast/http/parser.hpp
index 4c35e746..d7801171 100644
--- a/include/beast/http/parser.hpp
+++ b/include/beast/http/parser.hpp
@@ -14,6 +14,7 @@
 #include <beast/http/type_traits.hpp>
 #include <boost/optional.hpp>
 #include <boost/throw_exception.hpp>
+#include <functional>
 #include <memory>
 #include <type_traits>
 #include <utility>
@@ -60,6 +61,7 @@ class parser
 
     message<isRequest, Body, basic_fields<Allocator>> m_;
     boost::optional<typename Body::writer> wr_;
+    std::function<void(parser&, error_code&)> cb_;
 
 public:
     /// The type of message returned by the parser
@@ -191,6 +193,28 @@ public:
         return std::move(m_);
     }
 
+    /** Set the on_header callback.
+
+        This optional callback is invoked after the parser receives
+        a complete header. The function must be invokable with
+        this signature:
+        @code
+        void callback(
+            parser<isRequest, Body, Fields>& p,     // `*this`
+            error_code& ec)                         // Set to the error, if any
+        @endcode
+        The callback will ensure that `!ec` is `true` if there was
+        no error or set to the appropriate error code if there was one. 
+
+        The callback may not call @ref put or @ref put_eof, or 
+        else the behavior is undefined.
+    */
+    void
+    on_header(std::function<void(parser&, error_code&)> cb)
+    {
+        cb_ = cb;
+    }
+
 private:
     friend class basic_parser<isRequest, parser>;
 
@@ -250,7 +274,10 @@ private:
     void
     on_header(error_code& ec)
     {
-        ec.assign(0, ec.category());
+        if(cb_)
+            cb_(*this, ec);
+        else
+            ec.assign(0, ec.category());
     }
 
     void
diff --git a/test/http/basic_parser.cpp b/test/http/basic_parser.cpp
index b3ed989f..af51212d 100644
--- a/test/http/basic_parser.cpp
+++ b/test/http/basic_parser.cpp
@@ -984,6 +984,7 @@ public:
         }
     }
 
+
     //--------------------------------------------------------------------------
 
     template<bool isRequest, class Derived>
diff --git a/test/http/parser.cpp b/test/http/parser.cpp
index f63e2834..cae9d904 100644
--- a/test/http/parser.cpp
+++ b/test/http/parser.cpp
@@ -328,6 +328,49 @@ public:
         BEAST_EXPECT(used == 0);
     }
 
+    void
+    testCallback()
+    {
+        {
+            multi_buffer b;
+            ostream(b) << 
+                "POST / HTTP/1.1\r\n"
+                "Content-Length: 2\r\n"
+                "\r\n"
+                "**";
+            error_code ec;
+            parser<true, string_body> p;
+            p.eager(true);
+            p.put(b.data(), ec);
+            p.on_header(
+                [this](parser<true, string_body>& p, error_code& ec)
+                {
+                    BEAST_EXPECT(p.is_header_done());
+                    ec.assign(0, ec.category());
+                });
+            BEAST_EXPECTS(! ec, ec.message());
+        }
+        {
+            multi_buffer b;
+            ostream(b) << 
+                "POST / HTTP/1.1\r\n"
+                "Content-Length: 2\r\n"
+                "\r\n"
+                "**";
+            error_code ec;
+            parser<true, string_body> p;
+            p.eager(true);
+            p.put(b.data(), ec);
+            p.on_header(
+                [this](parser<true, string_body>&, error_code& ec)
+                {
+                    ec.assign(errc::bad_message,
+                        generic_category());
+                });
+            BEAST_EXPECTS(! ec, ec.message());
+        }
+    }
+
     void
     run() override
     {
@@ -335,6 +378,7 @@ public:
         testNeedMore<flat_buffer>();
         testNeedMore<multi_buffer>();
         testGotSome();
+        testCallback();
     }
 };