From 4a5f25615a31e35f63af16d245701220035e686a Mon Sep 17 00:00:00 2001
From: Antonio Nesic <anesic@collectivemind.dev>
Date: Mon, 9 Dec 2024 21:32:48 +0100
Subject: [PATCH] Added Partitioned flag for cookies

---
 lib/inc/drogon/Cookie.h           | 24 ++++++++++++++++++++++++
 lib/src/Cookie.cc                 |  6 +++++-
 lib/tests/unittests/CookieTest.cc | 25 +++++++++++++++++++++++++
 3 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/lib/inc/drogon/Cookie.h b/lib/inc/drogon/Cookie.h
index 65c83a40e6..2f02669bd2 100644
--- a/lib/inc/drogon/Cookie.h
+++ b/lib/inc/drogon/Cookie.h
@@ -156,6 +156,18 @@ class DROGON_EXPORT Cookie
         sameSite_ = sameSite;
     }
 
+    /**
+     *  @brief Set the partitioned status of the cookie
+     */
+    void setPartitioned(bool partitioned)
+    {
+        partitioned_ = partitioned;
+        if (partitioned)
+        {
+            setSecure(true);
+        }
+    }
+
     /**
      * @brief Get the string value of the cookie
      */
@@ -282,6 +294,17 @@ class DROGON_EXPORT Cookie
         return secure_;
     }
 
+    /**
+     * @brief Check if the cookie is partitioned.
+     *
+     * @return true means the cookie is partitioned.
+     * @return false means the cookie is not partitioned.
+     */
+    bool isPartitioned() const
+    {
+        return partitioned_;
+    }
+
     /**
      * @brief Get the max-age of the cookie
      */
@@ -394,6 +417,7 @@ class DROGON_EXPORT Cookie
     trantor::Date expiresDate_{(std::numeric_limits<int64_t>::max)()};
     bool httpOnly_{true};
     bool secure_{false};
+    bool partitioned_{false};
     std::string domain_;
     std::string path_;
     std::string key_;
diff --git a/lib/src/Cookie.cc b/lib/src/Cookie.cc
index b46caeb04d..cb89b3f04f 100644
--- a/lib/src/Cookie.cc
+++ b/lib/src/Cookie.cc
@@ -69,7 +69,7 @@ std::string Cookie::cookieString() const
                 ret.append("SameSite=Lax; ");
         }
     }
-    if (secure_ && sameSite_ != SameSite::kNone)
+    if ((secure_ && sameSite_ != SameSite::kNone) || partitioned_)
     {
         ret.append("Secure; ");
     }
@@ -77,6 +77,10 @@ std::string Cookie::cookieString() const
     {
         ret.append("HttpOnly; ");
     }
+    if (partitioned_)
+    {
+        ret.append("Partitioned; ");
+    }
     ret.resize(ret.length() - 2);  // delete last semicolon
     ret.append("\r\n");
     return ret;
diff --git a/lib/tests/unittests/CookieTest.cc b/lib/tests/unittests/CookieTest.cc
index 0a634415ed..bdeb9effe9 100644
--- a/lib/tests/unittests/CookieTest.cc
+++ b/lib/tests/unittests/CookieTest.cc
@@ -35,4 +35,29 @@ DROGON_TEST(CookieTest)
           drogon::Cookie::convertString2SameSite("Strict"));
     CHECK(drogon::Cookie::SameSite::kNone ==
           drogon::Cookie::convertString2SameSite("None"));
+
+    // Test for Partitioned attribute
+    drogon::Cookie cookie6("test", "6");
+    cookie6.setPartitioned(true);
+    CHECK(cookie6.cookieString() ==
+          "Set-Cookie: test=6; Secure; HttpOnly; Partitioned\r\n");
+    // Test that partitioned attribute  automatically sets secure
+    drogon::Cookie cookie7("test", "7");
+    cookie7.setPartitioned(true);
+    CHECK(cookie7.isSecure() == true);
+    // Test other attributes
+    drogon::Cookie cookie8("test", "8");
+    cookie8.setPartitioned(true);
+    cookie8.setDomain("drogon.org");
+    cookie8.setMaxAge(3600);
+    CHECK(cookie8.cookieString() ==
+          "Set-Cookie: test=8; Max-Age=3600; Domain=drogon.org; Secure; "
+          "HttpOnly; Partitioned\r\n");
+    // Teset Partitioned and SameSite can coexist
+    drogon::Cookie cookie9("test", "9");
+    cookie9.setPartitioned(true);
+    cookie9.setSameSite(drogon::Cookie::SameSite::kLax);
+    CHECK(
+        cookie9.cookieString() ==
+        "Set-Cookie: test=9; SameSite=Lax; Secure; HttpOnly; Partitioned\r\n");
 }