add max-age, samesite options to Cookie (#1159)

Co-authored-by: an-tao <antao2002@gmail.com>
drogonframework · Jan 16, 2022 · bbc3161
1 parent bfb25a3
commit bbc3161
Show file tree

Hide file tree

Showing 5 changed files with 117 additions and 3 deletions.
diff --git a/lib/inc/drogon/Cookie.h b/lib/inc/drogon/Cookie.h
@@ -14,6 +14,7 @@
 #pragma once
 
 #include <drogon/exports.h>
+#include <drogon/utils/optional.h>
 #include <trantor/utils/Date.h>
 #include <string>
 #include <limits>
@@ -40,7 +41,13 @@ class DROGON_EXPORT Cookie
     {
     }
     Cookie() = default;
-
+    enum class SameSite
+    {
+        kNull,
+        kLax,
+        kStrict,
+        kNone
+    };
     /**
      * @brief Set the Expires Date
      *
@@ -113,6 +120,20 @@ class DROGON_EXPORT Cookie
     {
         value_ = std::move(value);
     }
+    /**
+     * @brief Set the max-age of the cookie.
+     */
+    void setMaxAge(int value)
+    {
+        maxAge_ = value;
+    }
+    /**
+     * @brief Set the same site of the cookie.
+     */
+    void setSameSite(SameSite sameSite)
+    {
+        sameSite_ = sameSite;
+    }
 
     /**
      * @brief Get the string value of the cookie
@@ -240,6 +261,38 @@ class DROGON_EXPORT Cookie
         return secure_;
     }
 
+    /**
+     * @brief Get the max-age of the cookie
+     */
+    optional<int> maxAge() const
+    {
+        return maxAge_;
+    }
+
+    /**
+     * @brief Get the max-age of the cookie
+     */
+    optional<int> getMaxAge() const
+    {
+        return maxAge_;
+    }
+
+    /**
+     * @brief Get the same site of the cookie
+     */
+    SameSite sameSite() const
+    {
+        return sameSite_;
+    }
+
+    /**
+     * @brief Get the same site of the cookie
+     */
+    SameSite getSameSite() const
+    {
+        return sameSite_;
+    }
+
   private:
     trantor::Date expiresDate_{(std::numeric_limits<int64_t>::max)()};
     bool httpOnly_{true};
@@ -248,6 +301,8 @@ class DROGON_EXPORT Cookie
     std::string path_;
     std::string key_;
     std::string value_;
+    optional<int> maxAge_;
+    SameSite sameSite_{SameSite::kNull};
 };
 
 }  // namespace drogon
diff --git a/lib/src/Cookie.cc b/lib/src/Cookie.cc
@@ -29,6 +29,12 @@ std::string Cookie::cookieString() const
             .append(utils::getHttpFullDate(expiresDate_))
             .append("; ");
     }
+    if (maxAge_.has_value())
+    {
+        ret.append("Max-Age=")
+            .append(std::to_string(maxAge_.value()))
+            .append("; ");
+    }
     if (!domain_.empty())
     {
         ret.append("Domain=").append(domain_).append("; ");
@@ -37,7 +43,29 @@ std::string Cookie::cookieString() const
     {
         ret.append("Path=").append(path_).append("; ");
     }
-    if (secure_)
+    if (sameSite_ != SameSite::kNull)
+    {
+        switch (sameSite_)
+        {
+            case SameSite::kLax:
+                ret.append("SameSite=Lax; ");
+                break;
+            case SameSite::kStrict:
+                ret.append("SameSite=Strict; ");
+                break;
+            case SameSite::kNone:
+                ret.append("SameSite=None; ");
+                // Cookies with SameSite=None must now also specify the Secure
+                // attribute (they require a secure context/HTTPS).
+                ret.append("Secure; ");
+                break;
+            default:
+                // Lax replaced None as the default value to ensure that users
+                // have reasonably robust defense against some CSRF attacks
+                ret.append("SameSite=Lax; ");
+        }
+    }
+    if (secure_ && sameSite_ != SameSite::kNone)
     {
         ret.append("Secure; ");
     }

diff --git a/lib/src/HttpResponseImpl.cc b/lib/src/HttpResponseImpl.cc
@@ -742,6 +742,25 @@ void HttpResponseImpl::addHeader(const char *start,
                 {
                     cookie.setHttpOnly(true);
                 }
+                else if (cookie_name == "samesite")
+                {
+                    if (cookie_value == "Lax")
+                    {
+                        cookie.setSameSite(Cookie::SameSite::kLax);
+                    }
+                    else if (cookie_value == "Strict")
+                    {
+                        cookie.setSameSite(Cookie::SameSite::kStrict);
+                    }
+                    else if (cookie_value == "None")
+                    {
+                        cookie.setSameSite(Cookie::SameSite::kNone);
+                    }
+                }
+                else if (cookie_name == "max-age")
+                {
+                    cookie.setMaxAge(std::stoi(cookie_value));
+                }
             }
         }
         if (!cookie.key().empty())

diff --git a/lib/tests/unittests/CookieTest.cc b/lib/tests/unittests/CookieTest.cc
@@ -16,4 +16,16 @@ DROGON_TEST(CookieTest)
     CHECK(cookie3.cookieString() ==
           "Set-Cookie: test=3; Expires=Fri, 21 May 2021 01:45:57 GMT; "
           "Domain=drogon.org; HttpOnly\r\n");
+
+    drogon::Cookie cookie4("test", "4");
+    cookie4.setMaxAge(3600);
+    cookie4.setSameSite(drogon::Cookie::SameSite::kStrict);
+    CHECK(cookie4.cookieString() ==
+          "Set-Cookie: test=4; Max-Age=3600; "