From 731164bec2664ca57d1c167531695e59f9504e88 Mon Sep 17 00:00:00 2001
From: Vinnie Falco <vinnie.falco@gmail.com>
Date: Thu, 6 Jul 2017 20:58:07 -0700
Subject: [PATCH] Add serializer::limit

---
 CHANGELOG.md                              |   1 +
 include/beast/core/buffer_prefix.hpp      |   6 +-
 include/beast/core/impl/buffer_prefix.ipp |   5 +-
 include/beast/http/impl/serializer.ipp    | 107 ++++++++++++----------
 include/beast/http/impl/write.ipp         |  16 ++--
 include/beast/http/serializer.hpp         |  68 +++++++++++---
 test/http/serializer.cpp                  |  54 +++++++++++
 7 files changed, 182 insertions(+), 75 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5b6ba677..4761060e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@ Version 75:
 * Shrink serializer buffers using buffers_ref
 * Tidy up BEAST_NO_BIG_VARIANTS
 * Shrink serializer buffers using buffers_ref
+* Add serializer::limit
 
 --------------------------------------------------------------------------------
 
diff --git a/include/beast/core/buffer_prefix.hpp b/include/beast/core/buffer_prefix.hpp
index 6625b67e..1450504f 100644
--- a/include/beast/core/buffer_prefix.hpp
+++ b/include/beast/core/buffer_prefix.hpp
@@ -28,8 +28,10 @@ namespace beast {
 template<class BufferSequence>
 class buffer_prefix_view
 {
-    using iter_type =
-        typename BufferSequence::const_iterator;
+    using buffers_type = typename
+        std::decay<BufferSequence>::type;
+
+    using iter_type = typename buffers_type::const_iterator;
 
     BufferSequence bs_;
     iter_type back_;
diff --git a/include/beast/core/impl/buffer_prefix.ipp b/include/beast/core/impl/buffer_prefix.ipp
index f6c6470e..b3aa5578 100644
--- a/include/beast/core/impl/buffer_prefix.ipp
+++ b/include/beast/core/impl/buffer_prefix.ipp
@@ -48,11 +48,8 @@ class buffer_prefix_view<BufferSequence>::const_iterator
 {
     friend class buffer_prefix_view<BufferSequence>;
 
-    using iter_type =
-        typename BufferSequence::const_iterator;
-
     buffer_prefix_view const* b_ = nullptr;
-    typename BufferSequence::const_iterator it_;
+    iter_type it_;
 
 public:
     using value_type = typename std::conditional<
diff --git a/include/beast/http/impl/serializer.ipp b/include/beast/http/impl/serializer.ipp
index 17618edc..4173b2c6 100644
--- a/include/beast/http/impl/serializer.ipp
+++ b/include/beast/http/impl/serializer.ipp
@@ -36,6 +36,22 @@ frdinit(std::false_type)
     frd_.emplace(m_, m_.version, m_.result_int());
 }
 
+template<bool isRequest, class Body,
+    class Fields, class ChunkDecorator>
+template<class T1, class T2, class Visit>
+inline
+void
+serializer<isRequest, Body, Fields, ChunkDecorator>::
+do_visit(error_code& ec, Visit& visit)
+{
+    // VFALCO work-around for missing variant::emplace
+    pv_.~variant();
+    new(&pv_) decltype(pv_){
+        T1{limit_, boost::get<T2>(v_)}};
+    visit(ec, beast::detail::make_buffers_ref(
+        boost::get<T1>(pv_)));
+}
+
 template<bool isRequest, class Body,
     class Fields, class ChunkDecorator>
 serializer<isRequest, Body, Fields, ChunkDecorator>::
@@ -83,7 +99,7 @@ next(error_code& ec, Visit&& visit)
         if(! result)
             goto go_header_only;
         more_ = result->second;
-        v_ = cb0_t{
+        v_ = cb2_t{
             boost::in_place_init,
             frd_->get(),
             result->first};
@@ -92,16 +108,14 @@ next(error_code& ec, Visit&& visit)
     }
 
     case do_header:
-        visit(ec, make_buffers_ref(
-            boost::get<cb0_t>(v_)));
+        do_visit<pcb2_t, cb2_t>(ec, visit);
         break;
 
     go_header_only:
-        v_ = ch_t{frd_->get()};
+        v_ = cb1_t{frd_->get()};
         s_ = do_header_only;
     case do_header_only:
-        visit(ec, make_buffers_ref(
-            boost::get<ch_t>(v_)));
+        do_visit<pcb1_t, cb1_t>(ec, visit);
         break;
 
     case do_body:
@@ -122,14 +136,13 @@ next(error_code& ec, Visit&& visit)
         if(! result)
             goto go_complete;
         more_ = result->second;
-        v_ = cb1_t{result->first};
+        v_ = cb3_t{result->first};
         s_ = do_body + 2;
         BEAST_FALLTHROUGH;
     }
 
     case do_body + 2:
-        visit(ec, make_buffers_ref(
-            boost::get<cb1_t>(v_)));
+        do_visit<pcb3_t, cb3_t>(ec, visit);
         break;
 
     //----------------------------------------------------------------------
@@ -155,7 +168,7 @@ next(error_code& ec, Visit&& visit)
         if(! more_)
         {
             // do it all in one buffer
-            v_ = ch3_t{
+            v_ = cb7_t{
                 boost::in_place_init,
                 frd_->get(),
                 detail::chunk_header{
@@ -183,7 +196,7 @@ next(error_code& ec, Visit&& visit)
             goto go_all_c;
         }
     #endif
-        v_ = ch0_t{
+        v_ = cb4_t{
             boost::in_place_init,
             frd_->get(),
             detail::chunk_header{
@@ -203,16 +216,14 @@ next(error_code& ec, Visit&& visit)
     }
 
     case do_header_c:
-        visit(ec, make_buffers_ref(
-            boost::get<ch0_t>(v_)));
+        do_visit<pcb4_t, cb4_t>(ec, visit);
         break;
 
     go_header_only_c:
-        v_ = ch_t{frd_->get()};
+        v_ = cb1_t{frd_->get()};
         s_ = do_header_only_c;
     case do_header_only_c:
-        visit(ec, make_buffers_ref(
-            boost::get<ch_t>(v_)));
+        do_visit<pcb1_t, cb1_t>(ec, visit);
         break;
 
     case do_body_c:
@@ -237,7 +248,7 @@ next(error_code& ec, Visit&& visit)
         if(! more_)
         {
             // do it all in one buffer
-            v_ = ch2_t{
+            v_ = cb6_t{
                 boost::in_place_init,
                 detail::chunk_header{
                     buffer_size(result->first)},
@@ -264,7 +275,7 @@ next(error_code& ec, Visit&& visit)
             goto go_body_final_c;
         }
     #endif
-        v_ = ch1_t{
+        v_ = cb5_t{
             boost::in_place_init,
             detail::chunk_header{
                 buffer_size(result->first)},
@@ -283,29 +294,26 @@ next(error_code& ec, Visit&& visit)
     }
 
     case do_body_c + 2:
-        visit(ec, make_buffers_ref(
-            boost::get<ch1_t>(v_)));
+        do_visit<pcb5_t, cb5_t>(ec, visit);
         break;
 
 #ifndef BEAST_NO_BIG_VARIANTS
     go_body_final_c:
         s_ = do_body_final_c;
     case do_body_final_c:
-        visit(ec, make_buffers_ref(
-            boost::get<ch2_t>(v_)));
+        do_visit<pcb6_t, cb6_t>(ec, visit);
         break;
 
     go_all_c:
         s_ = do_all_c;
     case do_all_c:
-        visit(ec, make_buffers_ref(
-            boost::get<ch3_t>(v_)));
+        do_visit<pcb7_t, cb7_t>(ec, visit);
         break;
 #endif
 
     go_final_c:
     case do_final_c:
-        v_ = ch4_t{
+        v_ = cb8_t{
             boost::in_place_init,
             detail::chunk_final(),
             [&]()
@@ -321,8 +329,7 @@ next(error_code& ec, Visit&& visit)
         BEAST_FALLTHROUGH;
 
     case do_final_c + 1:
-        visit(ec, make_buffers_ref(
-            boost::get<ch4_t>(v_)));
+        do_visit<pcb8_t, cb8_t>(ec, visit);
         break;
 
     //----------------------------------------------------------------------
@@ -349,9 +356,9 @@ consume(std::size_t n)
     {
     case do_header:
         BOOST_ASSERT(n <= buffer_size(
-            boost::get<cb0_t>(v_)));
-        boost::get<cb0_t>(v_).consume(n);
-        if(buffer_size(boost::get<cb0_t>(v_)) > 0)
+            boost::get<cb2_t>(v_)));
+        boost::get<cb2_t>(v_).consume(n);
+        if(buffer_size(boost::get<cb2_t>(v_)) > 0)
             break;
         header_done_ = true;
         v_ = boost::blank{};
@@ -362,9 +369,9 @@ consume(std::size_t n)
 
     case do_header_only:
         BOOST_ASSERT(n <= buffer_size(
-            boost::get<ch_t>(v_)));
-        boost::get<ch_t>(v_).consume(n);
-        if(buffer_size(boost::get<ch_t>(v_)) > 0)
+            boost::get<cb1_t>(v_)));
+        boost::get<cb1_t>(v_).consume(n);
+        if(buffer_size(boost::get<cb1_t>(v_)) > 0)
             break;
         frd_ = boost::none;
         header_done_ = true;
@@ -376,9 +383,9 @@ consume(std::size_t n)
     case do_body + 2:
     {
         BOOST_ASSERT(n <= buffer_size(
-            boost::get<cb1_t>(v_)));
-        boost::get<cb1_t>(v_).consume(n);
-        if(buffer_size(boost::get<cb1_t>(v_)) > 0)
+            boost::get<cb3_t>(v_)));
+        boost::get<cb3_t>(v_).consume(n);
+        if(buffer_size(boost::get<cb3_t>(v_)) > 0)
             break;
         v_ = boost::blank{};
         if(! more_)
@@ -391,9 +398,9 @@ consume(std::size_t n)
 
     case do_header_c:
         BOOST_ASSERT(n <= buffer_size(
-            boost::get<ch0_t>(v_)));
-        boost::get<ch0_t>(v_).consume(n);
-        if(buffer_size(boost::get<ch0_t>(v_)) > 0)
+            boost::get<cb4_t>(v_)));
+        boost::get<cb4_t>(v_).consume(n);
+        if(buffer_size(boost::get<cb4_t>(v_)) > 0)
             break;
         header_done_ = true;
         v_ = boost::blank{};
@@ -406,9 +413,9 @@ consume(std::size_t n)
     case do_header_only_c:
     {
         BOOST_ASSERT(n <= buffer_size(
-            boost::get<ch_t>(v_)));
-        boost::get<ch_t>(v_).consume(n);
-        if(buffer_size(boost::get<ch_t>(v_)) > 0)
+            boost::get<cb1_t>(v_)));
+        boost::get<cb1_t>(v_).consume(n);
+        if(buffer_size(boost::get<cb1_t>(v_)) > 0)
             break;
         frd_ = boost::none;
         header_done_ = true;
@@ -423,9 +430,9 @@ consume(std::size_t n)
 
     case do_body_c + 2:
         BOOST_ASSERT(n <= buffer_size(
-            boost::get<ch1_t>(v_)));
-        boost::get<ch1_t>(v_).consume(n);
-        if(buffer_size(boost::get<ch1_t>(v_)) > 0)
+            boost::get<cb5_t>(v_)));
+        boost::get<cb5_t>(v_).consume(n);
+        if(buffer_size(boost::get<cb5_t>(v_)) > 0)
             break;
         v_ = boost::blank{};
         if(more_)
@@ -437,7 +444,7 @@ consume(std::size_t n)
 #ifndef BEAST_NO_BIG_VARIANTS
     case do_body_final_c:
     {
-        auto& b = boost::get<ch2_t>(v_);
+        auto& b = boost::get<cb6_t>(v_);
         BOOST_ASSERT(n <= buffer_size(b));
         b.consume(n);
         if(buffer_size(b) > 0)
@@ -449,7 +456,7 @@ consume(std::size_t n)
 
     case do_all_c:
     {
-        auto& b = boost::get<ch3_t>(v_);
+        auto& b = boost::get<cb7_t>(v_);
         BOOST_ASSERT(n <= buffer_size(b));
         b.consume(n);
         if(buffer_size(b) > 0)
@@ -463,9 +470,9 @@ consume(std::size_t n)
 
     case do_final_c + 1:
         BOOST_ASSERT(buffer_size(
-            boost::get<ch4_t>(v_)));
-        boost::get<ch4_t>(v_).consume(n);
-        if(buffer_size(boost::get<ch4_t>(v_)) > 0)
+            boost::get<cb8_t>(v_)));
+        boost::get<cb8_t>(v_).consume(n);
+        if(buffer_size(boost::get<cb8_t>(v_)) > 0)
             break;
         v_ = boost::blank{};
         goto go_complete;
diff --git a/include/beast/http/impl/write.ipp b/include/beast/http/impl/write.ipp
index 08a63a2b..8b33e2c5 100644
--- a/include/beast/http/impl/write.ipp
+++ b/include/beast/http/impl/write.ipp
@@ -54,12 +54,12 @@ class write_some_op
         template<class ConstBufferSequence>
         void
         operator()(error_code& ec,
-            ConstBufferSequence const& buffer)
+            ConstBufferSequence const& buffers)
         {
             ec.assign(0, ec.category());
             invoked = true;
             return op_.s_.async_write_some(
-                buffer, std::move(op_));
+                buffers, std::move(op_));
         }
     };
 
@@ -211,12 +211,12 @@ class write_op
         template<class ConstBufferSequence>
         void
         operator()(error_code& ec,
-            ConstBufferSequence const& buffer)
+            ConstBufferSequence const& buffers)
         {
             ec.assign(0, ec.category());
             invoked = true;
             return op_.s_.async_write_some(
-                buffer, std::move(op_));
+                buffers, std::move(op_));
         }
     };
 
@@ -464,11 +464,11 @@ public:
     template<class ConstBufferSequence>
     void
     operator()(error_code& ec,
-        ConstBufferSequence const& buffer)
+        ConstBufferSequence const& buffers)
     {
         invoked = true;
         bytes_transferred =
-            stream_.write_some(buffer, ec);
+            stream_.write_some(buffers, ec);
     }
 };
 
@@ -490,11 +490,11 @@ public:
     template<class ConstBufferSequence>
     void
     operator()(error_code& ec,
-        ConstBufferSequence const& buffer)
+        ConstBufferSequence const& buffers)
     {
         invoked = true;
         bytes_transferred = boost::asio::write(
-            stream_, buffer, ec);
+            stream_, buffers, ec);
     }
 };
 
diff --git a/include/beast/http/serializer.hpp b/include/beast/http/serializer.hpp
index 0ba8f431..b37a829d 100644
--- a/include/beast/http/serializer.hpp
+++ b/include/beast/http/serializer.hpp
@@ -10,6 +10,7 @@
 
 #include <beast/config.hpp>
 #include <beast/core/buffer_cat.hpp>
+#include <beast/core/buffer_prefix.hpp>
 #include <beast/core/consuming_buffers.hpp>
 #include <beast/core/string.hpp>
 #include <beast/core/type_traits.hpp>
@@ -157,35 +158,44 @@ class serializer
     void frdinit(std::true_type);
     void frdinit(std::false_type);
 
+    template<class T1, class T2, class Visit>
+    void
+    do_visit(error_code& ec, Visit& visit);
+
     using reader = typename Body::reader;
 
-    using ch_t = consuming_buffers<typename
+    using cb1_t = consuming_buffers<typename
         Fields::reader::const_buffers_type>;        // header
+    using pcb1_t  = buffer_prefix_view<cb1_t const&>;
 
-    using cb0_t = consuming_buffers<buffer_cat_view<
+    using cb2_t = consuming_buffers<buffer_cat_view<
         typename Fields::reader::const_buffers_type,// header
         typename reader::const_buffers_type>>;      // body
+    using pcb2_t = buffer_prefix_view<cb2_t const&>;
 
-    using cb1_t = consuming_buffers<
+    using cb3_t = consuming_buffers<
         typename reader::const_buffers_type>;       // body
+    using pcb3_t = buffer_prefix_view<cb3_t const&>;
 
-    using ch0_t = consuming_buffers<buffer_cat_view<
+    using cb4_t = consuming_buffers<buffer_cat_view<
         typename Fields::reader::const_buffers_type,// header
         detail::chunk_header,                       // chunk-header
         boost::asio::const_buffers_1,               // chunk-ext
         boost::asio::const_buffers_1,               // crlf
         typename reader::const_buffers_type,        // body
         boost::asio::const_buffers_1>>;             // crlf
+    using pcb4_t = buffer_prefix_view<cb4_t const&>;
 
-    using ch1_t = consuming_buffers<buffer_cat_view<
+    using cb5_t = consuming_buffers<buffer_cat_view<
         detail::chunk_header,                       // chunk-header
         boost::asio::const_buffers_1,               // chunk-ext
         boost::asio::const_buffers_1,               // crlf
         typename reader::const_buffers_type,        // body
         boost::asio::const_buffers_1>>;             // crlf
+    using pcb5_t = buffer_prefix_view<cb5_t const&>;
 
 #ifndef BEAST_NO_BIG_VARIANTS
-    using ch2_t = consuming_buffers<buffer_cat_view<
+    using cb6_t = consuming_buffers<buffer_cat_view<
         detail::chunk_header,                       // chunk-header
         boost::asio::const_buffers_1,               // chunk-ext
         boost::asio::const_buffers_1,               // crlf
@@ -194,8 +204,9 @@ class serializer
         boost::asio::const_buffers_1,               // chunk-final
         boost::asio::const_buffers_1,               // trailers 
         boost::asio::const_buffers_1>>;             // crlf
+    using pcb6_t = buffer_prefix_view<cb6_t const&>;
 
-    using ch3_t = consuming_buffers<buffer_cat_view<
+    using cb7_t = consuming_buffers<buffer_cat_view<
         typename Fields::reader::const_buffers_type,// header
         detail::chunk_header,                       // chunk-header
         boost::asio::const_buffers_1,               // chunk-ext
@@ -205,22 +216,32 @@ class serializer
         boost::asio::const_buffers_1,               // chunk-final
         boost::asio::const_buffers_1,               // trailers 
         boost::asio::const_buffers_1>>;             // crlf
+    using pcb7_t = buffer_prefix_view<cb7_t const&>;
 #endif
 
-    using ch4_t = consuming_buffers<buffer_cat_view<
+    using cb8_t = consuming_buffers<buffer_cat_view<
         boost::asio::const_buffers_1,               // chunk-final
         boost::asio::const_buffers_1,               // trailers 
         boost::asio::const_buffers_1>>;             // crlf
+    using pcb8_t = buffer_prefix_view<cb8_t const&>;
 
     message<isRequest, Body, Fields> const& m_;
     boost::optional<typename Fields::reader> frd_;
     boost::optional<reader> rd_;
     boost::variant<boost::blank,
-        ch_t, cb0_t, cb1_t, ch0_t, ch1_t
+        cb1_t, cb2_t, cb3_t, cb4_t, cb5_t
     #ifndef BEAST_NO_BIG_VARIANTS
-        ,ch2_t, ch3_t
+        ,cb6_t, cb7_t
     #endif
-        , ch4_t> v_;
+        , cb8_t> v_;
+    boost::variant<boost::blank,
+        pcb1_t, pcb2_t, pcb3_t, pcb4_t, pcb5_t
+    #ifndef BEAST_NO_BIG_VARIANTS
+        ,pcb6_t, pcb7_t
+    #endif
+        , pcb8_t> pv_;
+    std::size_t limit_ =
+        (std::numeric_limits<std::size_t>::max)();
     int s_ = do_construct;
     bool split_ = false;
     bool header_done_ = false;
@@ -246,6 +267,31 @@ public:
     serializer(message<isRequest, Body, Fields> const& msg,
         ChunkDecorator const& decorator = ChunkDecorator{});
 
+    /// Returns the serialized buffer size limit
+    std::size_t
+    limit() const
+    {
+        return limit_;
+    }
+
+    /** Set the serialized buffer size limit
+
+        This function adjusts the limit on the maximum size of the
+        buffers passed to the visitor. The new size limit takes effect
+        in the following call to @ref next.
+
+        The default is no buffer size limit.
+
+        @param limit The new buffer size limit. If this number
+        is zero, the size limit is removed.
+    */
+    void
+    limit(std::size_t limit)
+    {
+        limit_ = limit > 0 ? limit:
+            (std::numeric_limits<std::size_t>::max)();
+    }
+
     /** Returns `true` if we will pause after writing the complete header.
     */
     bool
diff --git a/test/http/serializer.cpp b/test/http/serializer.cpp
index 3fb95423..022a09cd 100644
--- a/test/http/serializer.cpp
+++ b/test/http/serializer.cpp
@@ -7,3 +7,57 @@
 
 // Test that header file is self-contained.
 #include <beast/http/serializer.hpp>
+
+#include <beast/http/string_body.hpp>
+#include <beast/unit_test/suite.hpp>
+
+namespace beast {
+namespace http {
+
+class serializer_test : public beast::unit_test::suite
+{
+public:
+    struct lambda
+    {
+        std::size_t size;
+
+        template<class ConstBufferSequence>
+        void
+        operator()(error_code&,
+            ConstBufferSequence const& buffers)
+        {
+            size = boost::asio::buffer_size(buffers);
+        }
+    };
+
+    void
+    testWriteLimit()
+    {
+        auto const limit = 30;
+        lambda visit;
+        error_code ec;
+        response<string_body> res;
+        res.body.append(1000, '*');
+        serializer<false, string_body> sr{res};
+        sr.limit(limit);
+        for(;;)
+        {
+            sr.next(ec, visit);
+            BEAST_EXPECT(visit.size <= limit);
+            sr.consume(visit.size);
+            if(sr.is_done())
+                break;
+        }
+    }
+
+    void
+    run() override
+    {
+        testWriteLimit();
+    }
+};
+
+BEAST_DEFINE_TESTSUITE(serializer,http,beast);
+
+} // http
+} // beast