fix iterating headers with `set-cookie` (#4048)

* fix iterating headers with `set-cookie`

* a test

* move work to `HTTPHeaderMap::set`

* append set-cookie after sort

* remove compare function

5 files changed, 51 insertions, 53 deletions

diff --git a/src/bun.js/bindings/webcore/FetchHeaders.cpp b/src/bun.js/bindings/webcore/FetchHeaders.cpp
index 02a3c8b3a..f54214733 100644
--- a/src/bun.js/bindings/webcore/FetchHeaders.cpp
+++ b/src/bun.js/bindings/webcore/FetchHeaders.cpp
@@ -28,6 +28,7 @@
 
 #include "config.h"
 #include "FetchHeaders.h"
+#include "HTTPHeaderNames.h"
 
 #include "HTTPParsers.h"
 
@@ -261,30 +262,48 @@ void FetchHeaders::filterAndFill(const HTTPHeaderMap& headers, Guard guard)
     }
 }
 
-static NeverDestroyed<const String> setCookieLowercaseString(MAKE_STATIC_STRING_IMPL("set-cookie"));
-
 std::optional<KeyValuePair<String, String>> FetchHeaders::Iterator::next()
 {
     if (m_keys.isEmpty() || m_updateCounter != m_headers->m_updateCounter) {
+        bool hasSetCookie = !m_headers->getSetCookieHeaders().isEmpty();
         m_keys.resize(0);
-        m_keys.reserveCapacity(m_headers->m_headers.size());
+        m_keys.reserveCapacity(m_headers->m_headers.size() + (hasSetCookie ? 1 : 0));
         for (auto& header : m_headers->m_headers)
             m_keys.uncheckedAppend(header.asciiLowerCaseName());
         std::sort(m_keys.begin(), m_keys.end(), WTF::codePointCompareLessThan);
+        if (hasSetCookie)
+            m_keys.uncheckedAppend(String());
+
+        m_currentIndex += m_cookieIndex;
+        if (hasSetCookie) {
+            size_t setCookieKeyIndex = m_keys.size() - 1;
+            if (m_currentIndex < setCookieKeyIndex)
+                m_cookieIndex = 0;
+            else {
+                m_cookieIndex = std::min(m_currentIndex - setCookieKeyIndex, m_headers->getSetCookieHeaders().size());
+                m_currentIndex -= m_cookieIndex;
+            }
+        } else
+            m_cookieIndex = 0;
+
         m_updateCounter = m_headers->m_updateCounter;
-        m_cookieIndex = 0;
     }
 
     auto& setCookieHeaders = m_headers->m_headers.getSetCookieHeaders();
 
     while (m_currentIndex < m_keys.size()) {
-        auto key = m_keys[m_currentIndex++];
+        auto key = m_keys[m_currentIndex];
 
-        if (!setCookieHeaders.isEmpty() && key == setCookieLowercaseString) {
-            auto cookie = setCookieHeaders[m_cookieIndex++];
-            return KeyValuePair<String, String> { WTFMove(key), WTFMove(cookie) };
+        if (key.isNull()) {
+            if (m_cookieIndex < setCookieHeaders.size()) {
+                String value = setCookieHeaders[m_cookieIndex++];
+                return KeyValuePair<String, String> { WTF::staticHeaderNames[static_cast<uint8_t>(HTTPHeaderName::SetCookie)], WTFMove(value) };
+            }
+            m_currentIndex++;
+            continue;
         }
 
+        m_currentIndex++;
         auto value = m_headers->m_headers.get(key);
         if (!value.isNull())
             return KeyValuePair<String, String> { WTFMove(key), WTFMove(value) };
diff --git a/src/bun.js/bindings/webcore/HTTPHeaderMap.cpp b/src/bun.js/bindings/webcore/HTTPHeaderMap.cpp
index 013aad750..99fc9cf13 100644
--- a/src/bun.js/bindings/webcore/HTTPHeaderMap.cpp
+++ b/src/bun.js/bindings/webcore/HTTPHeaderMap.cpp
@@ -236,15 +236,7 @@ String HTTPHeaderMap::get(HTTPHeaderName name) const
 void HTTPHeaderMap::set(HTTPHeaderName name, const String& value)
 {
     if (name == HTTPHeaderName::SetCookie) {
-        auto cookieName = extractCookieName(value);
-        size_t length = m_setCookieHeaders.size();
-        const auto& cookies = m_setCookieHeaders.data();
-        for (size_t i = 0; i < length; ++i) {
-            if (extractCookieName(cookies[i]) == cookieName) {
-                m_setCookieHeaders[i] = value;
-                return;
-            }
-        }
+        m_setCookieHeaders.clear();
         m_setCookieHeaders.append(value);
         return;
     }
diff --git a/src/bun.js/bindings/webcore/HTTPHeaderMap.h b/src/bun.js/bindings/webcore/HTTPHeaderMap.h
index 13fc806cb..42fdbf41a 100644
--- a/src/bun.js/bindings/webcore/HTTPHeaderMap.h
+++ b/src/bun.js/bindings/webcore/HTTPHeaderMap.h
@@ -69,12 +69,9 @@ public:
             : m_table(table)
             , m_commonHeadersIt(commonHeadersIt)
             , m_uncommonHeadersIt(uncommonHeadersIt)
-            , m_setCookiesIter(setCookiesIter)
         {
             if (!updateKeyValue(m_commonHeadersIt)) {
-                if (!updateSetCookieHeaderPosition(setCookiesIter)) {
-                    updateKeyValue(m_uncommonHeadersIt);
-                }
+                updateKeyValue(m_uncommonHeadersIt);
             }
         }
 
@@ -107,9 +104,6 @@ public:
             if (m_commonHeadersIt != m_table.m_commonHeaders.end()) {
                 if (updateKeyValue(++m_commonHeadersIt))
                     return *this;
-            } else if (m_setCookiesIter != m_table.m_setCookieHeaders.end()) {
-                if (updateSetCookieHeaderPosition(++m_setCookiesIter))
-                    return *this;
             } else {
                 ++m_uncommonHeadersIt;
             }
@@ -122,7 +116,7 @@ public:
         bool operator!=(const HTTPHeaderMapConstIterator &other) const { return !(*this == other); }
         bool operator==(const HTTPHeaderMapConstIterator &other) const
         {
-            return m_commonHeadersIt == other.m_commonHeadersIt && m_uncommonHeadersIt == other.m_uncommonHeadersIt && m_setCookiesIter == other.m_setCookiesIter;
+            return m_commonHeadersIt == other.m_commonHeadersIt && m_uncommonHeadersIt == other.m_uncommonHeadersIt;
         }
 
     private:
@@ -145,22 +139,9 @@ public:
             return true;
         }
 
-        bool updateSetCookieHeaderPosition(Vector<String, 0>::const_iterator cookieI)
-        {
-            if (cookieI == m_table.m_setCookieHeaders.end()) {
-                return false;
-            }
-
-            m_keyValue.key = httpHeaderNameString(HTTPHeaderName::SetCookie).toStringWithoutCopying();
-            m_keyValue.keyAsHTTPHeaderName = HTTPHeaderName::SetCookie;
-            m_keyValue.value = *cookieI;
-            return true;
-        }
-
         const HTTPHeaderMap &m_table;
         CommonHeadersVector::const_iterator m_commonHeadersIt;
         UncommonHeadersVector::const_iterator m_uncommonHeadersIt;
-        Vector<String, 0>::const_iterator m_setCookiesIter;
         KeyValue m_keyValue;
     };
     typedef HTTPHeaderMapConstIterator const_iterator;
@@ -178,14 +159,12 @@ public:
     {
         m_commonHeaders.clear();
         m_uncommonHeaders.clear();
-        m_setCookieHeaders.clear();
     }
 
     void shrinkToFit()
     {
         m_commonHeaders.shrinkToFit();
         m_uncommonHeaders.shrinkToFit();
-        m_setCookieHeaders.shrinkToFit();
     }
 
     WEBCORE_EXPORT String get(const String &name) const;
diff --git a/test/js/deno/fetch/headers.test.ts b/test/js/deno/fetch/headers.test.ts
index ee2466e32..efc4e6a6e 100644
--- a/test/js/deno/fetch/headers.test.ts
+++ b/test/js/deno/fetch/headers.test.ts
@@ -312,16 +312,16 @@ test(function headersInitMultiple() {
     ];
     assertEquals(actual, [
         [
+            "x-deno",
+            "foo, bar"
+        ],
+        [
             "set-cookie",
             "foo=bar"
         ],
         [
             "set-cookie",
             "bar=baz"
-        ],
-        [
-            "x-deno",
-            "foo, bar"
         ]
     ]);
 });
@@ -363,16 +363,16 @@ test(function headersAppendMultiple() {
     ];
     assertEquals(actual, [
         [
+            "x-deno",
+            "foo, bar"
+        ],
+        [
             "set-cookie",
             "foo=bar"
         ],
         [
             "set-cookie",
             "bar=baz"
-        ],
-        [
-            "x-deno",
-            "foo, bar"
         ]
     ]);
 });
diff --git a/test/js/web/fetch/fetch.test.ts b/test/js/web/fetch/fetch.test.ts
index ec73069de..a1585ba14 100644
--- a/test/js/web/fetch/fetch.test.ts
+++ b/test/js/web/fetch/fetch.test.ts
@@ -234,11 +234,11 @@ describe("Headers", () => {
     ]);
     const actual = [...headers];
     expect(actual).toEqual([
+      ["x-bun", "abc, def"],
       ["set-cookie", "foo=bar"],
       ["set-cookie", "bar=baz"],
-      ["x-bun", "abc, def"],
     ]);
-    expect([...headers.values()]).toEqual(["foo=bar", "bar=baz", "abc, def"]);
+    expect([...headers.values()]).toEqual(["abc, def", "foo=bar", "bar=baz"]);
   });
 
   it("Headers append multiple", () => {
@@ -253,9 +253,9 @@ describe("Headers", () => {
     // we do not preserve the order
     // which is kind of bad
     expect(actual).toEqual([
+      ["x-bun", "foo, bar"],
       ["set-cookie", "foo=bar"],
       ["set-cookie", "bar=baz"],
-      ["x-bun", "foo, bar"],
     ]);
   });
 
@@ -276,9 +276,17 @@ describe("Headers", () => {
     headers.set("set-Cookie", "foo=baz");
     headers.set("set-cookie", "bar=qat");
     const actual = [...headers];
+    expect(actual).toEqual([["set-cookie", "bar=qat"]]);
+  });
+
+  it("should include set-cookie headers in array", () => {
+    const headers = new Headers();
+    headers.append("Set-Cookie", "foo=bar");
+    headers.append("Content-Type", "text/plain");
+    const actual = [...headers];
     expect(actual).toEqual([
-      ["set-cookie", "foo=baz"],
-      ["set-cookie", "bar=qat"],
+      ["content-type", "text/plain"],
+      ["set-cookie", "foo=bar"],
     ]);
   });
 });