Fix stale Set-Cookie pointer after HttpHeader mutations

src/brpc/details/http_message.cpp

@@ -186,6 +186,9 @@ int HttpMessage::on_headers_complete(http_parser *parser) {
         parser->http_major = 1;
     }
     HttpHeader& headers = http_message->header();
+    if (headers.HasSetCookieHeaders()) {
+        headers.SyncSetCookieHeaders();
+    }
     headers.set_version(parser->http_major, parser->http_minor);
     // Only for response
     // http_parser may set status_code to 0 when the field is not needed,

src/brpc/http_header.cpp

@@ -30,8 +30,28 @@ HttpHeader::HttpHeader()
     : _status_code(HTTP_STATUS_OK)
     , _method(HTTP_METHOD_GET)
     , _version(1, 1)
-    , _first_set_cookie(nullptr) {
-    // NOTE: don't forget to clear the field in Clear() as well.
+    , _set_cookie_headers_dirty(false) {}
+
+HttpHeader::HttpHeader(const HttpHeader& rhs)
+    : _headers(rhs._headers)
+    , _uri(rhs._uri)
+    , _status_code(rhs._status_code)
+    , _method(rhs._method)
+    , _content_type(rhs._content_type)
+    , _unresolved_path(rhs._unresolved_path)
+    , _version(rhs._version)
+    , _set_cookie_headers_dirty(rhs._set_cookie_headers_dirty) {
+    if (rhs._set_cookie_headers) {
+        _set_cookie_headers.reset(new SetCookieHeaders(*rhs._set_cookie_headers));
+    }
+}
+
+HttpHeader& HttpHeader::operator=(const HttpHeader& rhs) {
+    if (this != &rhs) {
+        HttpHeader tmp(rhs);
+        Swap(tmp);
+    }
+    return *this;
 }
 
 void HttpHeader::Swap(HttpHeader &rhs) {
@@ -42,6 +62,8 @@ void HttpHeader::Swap(HttpHeader &rhs) {
     _content_type.swap(rhs._content_type);
     _unresolved_path.swap(rhs._unresolved_path);
     std::swap(_version, rhs._version);
+    _set_cookie_headers.swap(rhs._set_cookie_headers);
+    std::swap(_set_cookie_headers_dirty, rhs._set_cookie_headers_dirty);
 }
 
 void HttpHeader::Clear() {
@@ -52,6 +74,8 @@ void HttpHeader::Clear() {
     _content_type.clear();
     _unresolved_path.clear();
     _version = std::make_pair(1, 1);
+    _set_cookie_headers.reset();
+    _set_cookie_headers_dirty = false;
 }
 
 const std::string* HttpHeader::GetHeader(const char* key) const {
@@ -60,29 +84,39 @@ const std::string* HttpHeader::GetHeader(const char* key) const {
 
 const std::string* HttpHeader::GetHeader(const std::string& key) const {
     if (IsSetCookie(key)) {
-        return _first_set_cookie;
+        return HasSetCookieHeaders()
+            ? &_set_cookie_headers->front() : nullptr;
     }
     std::string* val = _headers.seek(key);
     return val;
 }
 
 std::vector<const std::string*> HttpHeader::GetAllSetCookieHeader() const {
-    return GetMultiLineHeaders(SET_COOKIE);
-}
-
-std::vector<const std::string*>
-HttpHeader::GetMultiLineHeaders(const std::string& key) const {
     std::vector<const std::string*> headers;
-    for (const auto& iter : _headers) {
-        if (_header_key_equal(iter.first, key)) {
-            headers.push_back(&iter.second);
+    if (_set_cookie_headers) {
+        headers.reserve(_set_cookie_headers->size());
+        for (const auto& header : *_set_cookie_headers) {
+            headers.push_back(&header);
         }
     }
     return headers;
 }
 
 void HttpHeader::SetHeader(const std::string& key,
                            const std::string& value) {
+    if (IsSetCookie(key)) {
+        if (!HasSetCookieHeaders()) {
+            if (!_set_cookie_headers) {
+                _set_cookie_headers.reset(new SetCookieHeaders);
+            }
+            _set_cookie_headers->push_back(value);
+        } else {
+            _set_cookie_headers->front() = value;
+        }
+        _set_cookie_headers_dirty = true;
+        SyncSetCookieHeaders();
+        return;
+    }
     GetOrAddHeader(key) = value;
 }
 
@@ -92,7 +126,8 @@ void HttpHeader::RemoveHeader(const char* key) {
     } else {
         _headers.erase(key);
         if (IsSetCookie(key)) {
-            _first_set_cookie = nullptr;
+            _set_cookie_headers.reset();
+            _set_cookie_headers_dirty = false;
         }
     }
 }
@@ -101,8 +136,14 @@ void HttpHeader::AppendHeader(const std::string& key,
     const butil::StringPiece& value) {
     if (!CanFoldedInLine(key)) {
         // Add a new Set-Cookie header field.
-        std::string& slot = AddHeader(key);
-        slot.assign(value.data(), value.size());
+        if (!_set_cookie_headers) {
+            _set_cookie_headers.reset(new SetCookieHeaders);
+        }
+        if (_set_cookie_headers_dirty) {
+            SyncSetCookieHeaders();
+        }
+        _set_cookie_headers->push_back(std::string(value.data(), value.size()));
+        _headers.insert({ SET_COOKIE, _set_cookie_headers->back() });
     } else {
         std::string& slot = GetOrAddHeader(key);
         if (slot.empty()) {
@@ -128,28 +169,39 @@ std::string& HttpHeader::GetOrAddHeader(const std::string& key) {
         return _content_type;
     }
 
-    bool is_set_cookie = IsSetCookie(key);
-    // Only returns the first Set-Cookie header field for compatibility.
-    if (is_set_cookie && nullptr != _first_set_cookie) {
-        return *_first_set_cookie;
-    }
-
     std::string* val = _headers.seek(key);
     if (nullptr == val) {
         val = _headers.insert({ key, "" });
-        if (is_set_cookie) {
-            _first_set_cookie = val;
-        }
     }
     return *val;
 }
 
 std::string& HttpHeader::AddHeader(const std::string& key) {
-    std::string* val = _headers.insert({ key, "" });
-    if (IsSetCookie(key) && nullptr == _first_set_cookie) {
-        _first_set_cookie = val;
+    if (!IsSetCookie(key)) {
+        return *_headers.insert({ key, "" });
     }
-    return *val;
+    if (!_set_cookie_headers) {
+        _set_cookie_headers.reset(new SetCookieHeaders);
+    }
+    if (!_set_cookie_headers_dirty) {
+        _headers.erase(SET_COOKIE);
+        _set_cookie_headers_dirty = true;
+    }
+    _set_cookie_headers->push_back(std::string());
+    return _set_cookie_headers->back();
+}
+
+void HttpHeader::SyncSetCookieHeaders() {
+    if (!_set_cookie_headers_dirty) {
+        return;
+    }
+    _headers.erase(SET_COOKIE);
+    if (_set_cookie_headers) {
+        for (const auto& header : *_set_cookie_headers) {
+            _headers.insert({ SET_COOKIE, header });
+        }
+    }
+    _set_cookie_headers_dirty = false;
 }
 
 const HttpHeader& DefaultHttpHeader() {

src/brpc/http_header.h

@@ -19,6 +19,8 @@
 #ifndef  BRPC_HTTP_HEADER_H
 #define  BRPC_HTTP_HEADER_H
 
+#include <deque>
+#include <memory>
 #include <vector>
 #include "butil/strings/string_piece.h"  // StringPiece
 #include "butil/containers/case_ignored_flat_map.h"
@@ -45,6 +47,8 @@ class HttpHeader {
     typedef HeaderMap::key_equal HeaderKeyEqual;
 
     HttpHeader();
+    HttpHeader(const HttpHeader& rhs);
+    HttpHeader& operator=(const HttpHeader& rhs);
 
     // Exchange internal fields with another HttpHeader.
     void Swap(HttpHeader &rhs);
@@ -120,7 +124,12 @@ class HttpHeader {
     HeaderIterator HeaderBegin() const { return _headers.begin(); }
     HeaderIterator HeaderEnd() const { return _headers.end(); }
     // #headers
-    size_t HeaderCount() const { return _headers.size(); }
+    size_t HeaderCount() const {
+        if (_set_cookie_headers_dirty) {
+            return _headers.size() + _set_cookie_headers->size();
+        }
+        return _headers.size();
+    }
 
     // Get the URI object, check src/brpc/uri.h for details.
     const URI& uri() const { return _uri; }
@@ -165,16 +174,20 @@ friend void policy::ProcessHttpRequest(InputMessageBase *msg);
     static const char* COOKIE;
     static const char* CONTENT_TYPE;
 
-    std::vector<const std::string*> GetMultiLineHeaders(const std::string& key) const;
-
     std::string& GetOrAddHeader(const std::string& key);
 
     std::string& AddHeader(const std::string& key);
 
+    void SyncSetCookieHeaders();
+
     bool IsSetCookie(const std::string& key) const {
         return _header_key_equal(key, SET_COOKIE);
     }
 
+    bool HasSetCookieHeaders() const {
+        return _set_cookie_headers && !_set_cookie_headers->empty();
+    }
+
     bool IsCookie(const std::string& key) const {
         return _header_key_equal(key, COOKIE);
     }
@@ -202,7 +215,13 @@ friend void policy::ProcessHttpRequest(InputMessageBase *msg);
     std::string _content_type;
     std::string _unresolved_path;
     std::pair<int, int> _version;
-    std::string* _first_set_cookie;
+    // HeaderMap rehashing invalidates value pointers and does not preserve
+    // duplicate-key order. deque keeps existing value pointers stable when
+    // another Set-Cookie header is appended.
+    typedef std::deque<std::string> SetCookieHeaders;
+    std::unique_ptr<SetCookieHeaders> _set_cookie_headers;
+    // When true, Set-Cookie headers are only stored in _set_cookie_headers.
+    bool _set_cookie_headers_dirty;
 };
 
 const HttpHeader& DefaultHttpHeader();

test/brpc_http_message_unittest.cpp

@@ -341,6 +341,25 @@ TEST(HttpMessageTest, parse_http_set_cookie) {
     }
     ASSERT_EQ(1, set_cookie_value1_count);
     ASSERT_EQ(1, set_cookie_value2_count);
+
+    const char first_part[] = "HTTP/1.1 200 OK\r\nSet-Cookie: a";
+    const char second_part[] = "=1\r\nSet-Cookie: b=2\r\n\r\n";
+    brpc::HttpMessage fragmented_message(false, brpc::HTTP_METHOD_HEAD);
+    ASSERT_EQ((ssize_t)sizeof(first_part) - 1,
+              fragmented_message.ParseFromArray(first_part,
+                                                 sizeof(first_part) - 1));
+    ASSERT_FALSE(fragmented_message.Completed());
+    ASSERT_EQ((ssize_t)sizeof(second_part) - 1,
+              fragmented_message.ParseFromArray(second_part,
+                                                 sizeof(second_part) - 1));
+    ASSERT_TRUE(fragmented_message.Completed());
+    set_cookie = fragmented_message.header().GetHeader("set-cookie");
+    ASSERT_NE(nullptr, set_cookie);
+    ASSERT_EQ("a=1", *set_cookie);
+    all_set_cookie = fragmented_message.header().GetAllSetCookieHeader();
+    ASSERT_EQ(2u, all_set_cookie.size());
+    ASSERT_EQ("a=1", *all_set_cookie[0]);
+    ASSERT_EQ("b=2", *all_set_cookie[1]);
 }
 
 TEST(HttpMessageTest, cl_and_te) {
@@ -852,12 +871,28 @@ TEST(HttpMessageTest, http_header) {
     header.SetHeader(brpc::HttpHeader::SET_COOKIE, set_cookie_value1);
     value = header.GetHeader(brpc::HttpHeader::SET_COOKIE);
     ASSERT_TRUE(value && *value == set_cookie_value1);
+    const std::string* first_set_cookie = value;
     header.AppendHeader(brpc::HttpHeader::SET_COOKIE, set_cookie_value2);
+    ASSERT_EQ(set_cookie_value1, *first_set_cookie);
     value = header.GetHeader(brpc::HttpHeader::SET_COOKIE);
     ASSERT_TRUE(value && *value == set_cookie_value1);
     header.SetHeader(brpc::HttpHeader::SET_COOKIE, set_cookie_value3);
     value = header.GetHeader(brpc::HttpHeader::SET_COOKIE);
     ASSERT_TRUE(value && *value == set_cookie_value3);
+
+    // set-cookie and user-agent collide in the default HeaderMap bucket
+    // layout. Updating Set-Cookie must not corrupt the surviving header.
+    brpc::HttpHeader collision_header;
+    collision_header.SetHeader(brpc::HttpHeader::SET_COOKIE,
+                               set_cookie_value1);
+    collision_header.SetHeader("user-agent", "brpc");
+    collision_header.SetHeader(brpc::HttpHeader::SET_COOKIE,
+                               set_cookie_value2);
+    value = collision_header.GetHeader(brpc::HttpHeader::SET_COOKIE);
+    ASSERT_TRUE(value && *value == set_cookie_value2);
+    value = collision_header.GetHeader("user-agent");
+    ASSERT_TRUE(value && *value == "brpc");
+
     std::vector<const std::string*> all_set_cookie
         = header.GetAllSetCookieHeader();
     ASSERT_EQ(2u, all_set_cookie.size());
@@ -868,7 +903,7 @@ TEST(HttpMessageTest, http_header) {
     int set_cookie_value2_count = 0;
     int set_cookie_value3_count = 0;
     for (auto iter = header.HeaderBegin(); iter != header.HeaderEnd(); ++iter) {
-        if (!header.IsSetCookie(brpc::HttpHeader::SET_COOKIE)) {
+        if (!header.IsSetCookie(iter->first)) {
             continue;
         }
         if (iter->second == set_cookie_value2) {
@@ -879,9 +914,62 @@ TEST(HttpMessageTest, http_header) {
     }
     ASSERT_EQ(1, set_cookie_value2_count);
     ASSERT_EQ(1, set_cookie_value3_count);
+    brpc::HttpHeader copied_header(header);
+    copied_header.SetHeader(brpc::HttpHeader::SET_COOKIE, set_cookie_value1);
+    value = header.GetHeader(brpc::HttpHeader::SET_COOKIE);
+    ASSERT_TRUE(value && *value == set_cookie_value3);
+    value = copied_header.GetHeader(brpc::HttpHeader::SET_COOKIE);
+    ASSERT_TRUE(value && *value == set_cookie_value1);
+    brpc::HttpHeader assigned_header;
+    assigned_header = header;
+    assigned_header.SetHeader(brpc::HttpHeader::SET_COOKIE, set_cookie_value2);
+    value = header.GetHeader(brpc::HttpHeader::SET_COOKIE);
+    ASSERT_TRUE(value && *value == set_cookie_value3);
+    value = assigned_header.GetHeader(brpc::HttpHeader::SET_COOKIE);
+    ASSERT_TRUE(value && *value == set_cookie_value2);
     header.RemoveHeader(brpc::HttpHeader::SET_COOKIE);
     ASSERT_FALSE(header.GetHeader(brpc::HttpHeader::SET_COOKIE));
-    ASSERT_EQ(header._first_set_cookie, nullptr);
+    ASSERT_TRUE(header.GetAllSetCookieHeader().empty());
+
+    header.SetHeader(brpc::HttpHeader::SET_COOKIE, set_cookie_value1);
+    const std::string* first_set_cookie_after_reset =
+        header.GetHeader(brpc::HttpHeader::SET_COOKIE);
+    ASSERT_NE(nullptr, first_set_cookie_after_reset);
+    header.AppendHeader(brpc::HttpHeader::SET_COOKIE, set_cookie_value2);
+    const size_t bucket_count = header._headers.bucket_count();
+    ASSERT_TRUE(header._headers.resize(bucket_count * 2));
+    ASSERT_NE(bucket_count, header._headers.bucket_count());
+    ASSERT_EQ(set_cookie_value1, *first_set_cookie_after_reset);
+    value = header.GetHeader(brpc::HttpHeader::SET_COOKIE);
+    ASSERT_TRUE(value && *value == set_cookie_value1);
+    all_set_cookie = header.GetAllSetCookieHeader();
+    ASSERT_EQ(2u, all_set_cookie.size());
+    ASSERT_EQ(set_cookie_value1, *all_set_cookie[0]);
+    ASSERT_EQ(set_cookie_value2, *all_set_cookie[1]);
+    int set_cookie_value1_after_rehash_count = 0;
+    int set_cookie_value2_after_rehash_count = 0;
+    for (auto iter = header.HeaderBegin(); iter != header.HeaderEnd(); ++iter) {
+        if (!header.IsSetCookie(iter->first)) {
+            continue;
+        }
+        if (iter->second == set_cookie_value1) {
+            ++set_cookie_value1_after_rehash_count;
+        } else if (iter->second == set_cookie_value2) {
+            ++set_cookie_value2_after_rehash_count;
+        }
+    }
+    ASSERT_EQ(1, set_cookie_value1_after_rehash_count);
+    ASSERT_EQ(1, set_cookie_value2_after_rehash_count);
+
+    brpc::HttpHeader other_header;
+    other_header.SetHeader(brpc::HttpHeader::SET_COOKIE, set_cookie_value2);
+    header.Swap(other_header);
+    value = header.GetHeader(brpc::HttpHeader::SET_COOKIE);
+    ASSERT_TRUE(value && *value == set_cookie_value2);
+    value = other_header.GetHeader(brpc::HttpHeader::SET_COOKIE);
+    ASSERT_TRUE(value && *value == set_cookie_value1);
+    header.Clear();
+    ASSERT_FALSE(header.GetHeader(brpc::HttpHeader::SET_COOKIE));
 
     ASSERT_EQ(brpc::HTTP_METHOD_GET, header.method());
     header.set_method(brpc::HTTP_METHOD_POST);