diff --git a/src/brpc/details/http_message.cpp b/src/brpc/details/http_message.cpp index 14a81e553c..e039719ce0 100644 --- a/src/brpc/details/http_message.cpp +++ b/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, diff --git a/src/brpc/http_header.cpp b/src/brpc/http_header.cpp index 6ef074712e..1404940ec5 100644 --- a/src/brpc/http_header.cpp +++ b/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,22 +84,19 @@ 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 HttpHeader::GetAllSetCookieHeader() const { - return GetMultiLineHeaders(SET_COOKIE); -} - -std::vector -HttpHeader::GetMultiLineHeaders(const std::string& key) const { std::vector 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; @@ -83,6 +104,19 @@ HttpHeader::GetMultiLineHeaders(const std::string& key) const { 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() { diff --git a/src/brpc/http_header.h b/src/brpc/http_header.h index 58b92ecb19..cbaaf64157 100644 --- a/src/brpc/http_header.h +++ b/src/brpc/http_header.h @@ -19,6 +19,8 @@ #ifndef BRPC_HTTP_HEADER_H #define BRPC_HTTP_HEADER_H +#include +#include #include #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 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 _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 SetCookieHeaders; + std::unique_ptr _set_cookie_headers; + // When true, Set-Cookie headers are only stored in _set_cookie_headers. + bool _set_cookie_headers_dirty; }; const HttpHeader& DefaultHttpHeader(); diff --git a/test/brpc_http_message_unittest.cpp b/test/brpc_http_message_unittest.cpp index 078659bef5..916a6e482e 100644 --- a/test/brpc_http_message_unittest.cpp +++ b/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 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);