From c8395edfbe13493ecf9b6110175d504903fa23ed Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Wed, 10 Feb 2016 22:35:21 +0900
Subject: [PATCH] asio: client: Limit incoming response header field buffer
 size

---
 src/asio_client_request_impl.cc  |  8 +++++++-
 src/asio_client_request_impl.h   |  4 ++++
 src/asio_client_response_impl.cc |  9 ++++++++-
 src/asio_client_response_impl.h  |  4 ++++
 src/asio_client_session_impl.cc  | 13 +++++++++++++
 5 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/asio_client_request_impl.cc b/src/asio_client_request_impl.cc
index 5512c967..1e4e2dd1 100644
--- a/src/asio_client_request_impl.cc
+++ b/src/asio_client_request_impl.cc
@@ -32,7 +32,7 @@ namespace nghttp2 {
 namespace asio_http2 {
 namespace client {
 
-request_impl::request_impl() : strm_(nullptr) {}
+request_impl::request_impl() : strm_(nullptr), header_buffer_size_(0) {}
 
 void request_impl::write_trailer(header_map h) {
   auto sess = strm_->session();
@@ -105,6 +105,12 @@ void request_impl::method(std::string s) { method_ = std::move(s); }
 
 const std::string &request_impl::method() const { return method_; }
 
+size_t request_impl::header_buffer_size() const { return header_buffer_size_; }
+
+void request_impl::update_header_buffer_size(size_t len) {
+  header_buffer_size_ += len;
+}
+
 } // namespace client
 } // namespace asio_http2
 } // namespace nghttp2
diff --git a/src/asio_client_request_impl.h b/src/asio_client_request_impl.h
index 802a83d1..e0d43d2c 100644
--- a/src/asio_client_request_impl.h
+++ b/src/asio_client_request_impl.h
@@ -75,6 +75,9 @@ public:
   void method(std::string s);
   const std::string &method() const;
 
+  size_t header_buffer_size() const;
+  void update_header_buffer_size(size_t len);
+
 private:
   header_map header_;
   response_cb response_cb_;
@@ -84,6 +87,7 @@ private:
   class stream *strm_;
   uri_ref uri_;
   std::string method_;
+  size_t header_buffer_size_;
 };
 
 } // namespace client
diff --git a/src/asio_client_response_impl.cc b/src/asio_client_response_impl.cc
index fce25a4f..bd2cdf5f 100644
--- a/src/asio_client_response_impl.cc
+++ b/src/asio_client_response_impl.cc
@@ -30,7 +30,8 @@ namespace nghttp2 {
 namespace asio_http2 {
 namespace client {
 
-response_impl::response_impl() : content_length_(-1), status_code_(0) {}
+response_impl::response_impl()
+    : content_length_(-1), header_buffer_size_(0), status_code_(0) {}
 
 void response_impl::on_data(data_cb cb) { data_cb_ = std::move(cb); }
 
@@ -52,6 +53,12 @@ header_map &response_impl::header() { return header_; }
 
 const header_map &response_impl::header() const { return header_; }
 
+size_t response_impl::header_buffer_size() const { return header_buffer_size_; }
+
+void response_impl::update_header_buffer_size(size_t len) {
+  header_buffer_size_ += len;
+}
+
 } // namespace client
 } // namespace asio_http2
 } // namespace nghttp2
diff --git a/src/asio_client_response_impl.h b/src/asio_client_response_impl.h
index e2c22862..524d7285 100644
--- a/src/asio_client_response_impl.h
+++ b/src/asio_client_response_impl.h
@@ -53,12 +53,16 @@ public:
   header_map &header();
   const header_map &header() const;
 
+  size_t header_buffer_size() const;
+  void update_header_buffer_size(size_t len);
+
 private:
   data_cb data_cb_;
 
   header_map header_;
 
   int64_t content_length_;
+  size_t header_buffer_size_;
   int status_code_;
 };
 
diff --git a/src/asio_client_session_impl.cc b/src/asio_client_session_impl.cc
index a4cce25a..d7630c3f 100644
--- a/src/asio_client_session_impl.cc
+++ b/src/asio_client_session_impl.cc
@@ -174,6 +174,12 @@ int on_header_callback(nghttp2_session *session, const nghttp2_frame *frame,
     if (token == http2::HD__STATUS) {
       res.status_code(util::parse_uint(value, valuelen));
     } else {
+      if (res.header_buffer_size() + namelen + valuelen > 64_k) {
+        nghttp2_submit_rst_stream(session, NGHTTP2_FLAG_NONE,
+                                  frame->hd.stream_id, NGHTTP2_INTERNAL_ERROR);
+        break;
+      }
+      res.update_header_buffer_size(namelen + valuelen);
 
       if (token == http2::HD_CONTENT_LENGTH) {
         res.content_length(util::parse_uint(value, valuelen));
@@ -214,6 +220,13 @@ int on_header_callback(nghttp2_session *session, const nghttp2_frame *frame,
       }
     // fall through
     default:
+      if (req.header_buffer_size() + namelen + valuelen > 64_k) {
+        nghttp2_submit_rst_stream(session, NGHTTP2_FLAG_NONE,
+                                  frame->hd.stream_id, NGHTTP2_INTERNAL_ERROR);
+        break;
+      }
+      req.update_header_buffer_size(namelen + valuelen);
+
       req.header().emplace(
           std::string(name, name + namelen),
           header_value{std::string(value, value + valuelen),