From d39e1e30d7cc166b6616a49c345c0fbb57150d74 Mon Sep 17 00:00:00 2001
From: dreamhead <dreamhead.cn@gmail.com>
Date: Sat, 26 Aug 2023 22:20:36 +0800
Subject: [PATCH] introduced sameSite in standalone cookie

---
 .../deserializer/CookieContainerDeserializer.java |  3 ++-
 .../moco/parser/model/CookieContainer.java        |  8 ++++++++
 .../dreamhead/moco/MocoCookieStandaloneTest.java  | 15 +++++++++++++--
 moco-runner/src/test/resources/cookie.json        | 13 +++++++++++++
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/moco-runner/src/main/java/com/github/dreamhead/moco/parser/deserializer/CookieContainerDeserializer.java b/moco-runner/src/main/java/com/github/dreamhead/moco/parser/deserializer/CookieContainerDeserializer.java
index 674cd9153..ac803b037 100644
--- a/moco-runner/src/main/java/com/github/dreamhead/moco/parser/deserializer/CookieContainerDeserializer.java
+++ b/moco-runner/src/main/java/com/github/dreamhead/moco/parser/deserializer/CookieContainerDeserializer.java
@@ -37,10 +37,11 @@ public static final class InternalCookieContainer {
         private LatencyContainer maxAge;
         private boolean secure;
         private boolean httpOnly;
+        private String sameSite;
         private String template;
 
         public CookieContainer toContainer() {
-            return CookieContainer.newContainer(value, path, domain, maxAge, secure, httpOnly, template);
+            return CookieContainer.newContainer(value, path, domain, maxAge, secure, httpOnly, sameSite, template);
         }
     }
 }
diff --git a/moco-runner/src/main/java/com/github/dreamhead/moco/parser/model/CookieContainer.java b/moco-runner/src/main/java/com/github/dreamhead/moco/parser/model/CookieContainer.java
index 77c3508ef..725805e9c 100644
--- a/moco-runner/src/main/java/com/github/dreamhead/moco/parser/model/CookieContainer.java
+++ b/moco-runner/src/main/java/com/github/dreamhead/moco/parser/model/CookieContainer.java
@@ -17,6 +17,7 @@ public final class CookieContainer implements Container {
     private LatencyContainer maxAge;
     private boolean secure;
     private boolean httpOnly;
+    private String sameSite;
     private String template;
 
     public static CookieContainer newContainer(final String text) {
@@ -28,6 +29,7 @@ public static CookieContainer newContainer(final String text) {
     public static CookieContainer newContainer(final String text, final String path,
                                                final String domain, final LatencyContainer maxAge,
                                                final boolean secure, final boolean httpOnly,
+                                               final String sameSite,
                                                final String template) {
         CookieContainer container = new CookieContainer();
         container.value = text;
@@ -36,6 +38,7 @@ public static CookieContainer newContainer(final String text, final String path,
         container.maxAge = maxAge;
         container.secure = secure;
         container.httpOnly = httpOnly;
+        container.sameSite = sameSite;
         container.template = template;
         return container;
     }
@@ -74,6 +77,10 @@ public CookieAttribute[] getOptions() {
             options.add(CookieAttribute.maxAge(this.maxAge.getLatency(), this.maxAge.getUnit()));
         }
 
+        if (this.sameSite != null) {
+            options.add(CookieAttribute.sameSite(this.sameSite));
+        }
+
         return options.toArray(new CookieAttribute[options.size()]);
     }
 
@@ -87,6 +94,7 @@ public String toString() {
                 .add("max age", maxAge)
                 .add("secure", secure)
                 .add("HTTP only", httpOnly)
+                .add("same site", sameSite)
                 .add("template", template)
                 .toString();
 
diff --git a/moco-runner/src/test/java/com/github/dreamhead/moco/MocoCookieStandaloneTest.java b/moco-runner/src/test/java/com/github/dreamhead/moco/MocoCookieStandaloneTest.java
index b9b971b37..cfcd261ad 100644
--- a/moco-runner/src/test/java/com/github/dreamhead/moco/MocoCookieStandaloneTest.java
+++ b/moco-runner/src/test/java/com/github/dreamhead/moco/MocoCookieStandaloneTest.java
@@ -3,6 +3,8 @@
 import com.google.common.net.HttpHeaders;
 import io.netty.handler.codec.http.cookie.ClientCookieDecoder;
 import io.netty.handler.codec.http.cookie.Cookie;
+import io.netty.handler.codec.http.cookie.CookieHeaderNames;
+import io.netty.handler.codec.http.cookie.DefaultCookie;
 import org.junit.Test;
 
 import java.io.IOException;
@@ -65,10 +67,19 @@ public void should_set_and_recognize_cookie_with_max_age() throws IOException {
         assertThat(decodeCookie.maxAge(), is(3600L));
     }
 
-    private Cookie getCookie(String uri) throws IOException {
+    @Test
+    public void should_set_and_recognize_cookie_with_same_site() throws IOException {
+        runWithConfiguration("cookie.json");
+        DefaultCookie decodeCookie = getCookie("/cookie-with-same-site");
+        assertThat(decodeCookie.name(), is("login"));
+        assertThat(decodeCookie.value(), is("true"));
+        assertThat(decodeCookie.sameSite(), is(CookieHeaderNames.SameSite.Lax));
+    }
+
+    private DefaultCookie getCookie(String uri) throws IOException {
         org.apache.hc.core5.http.HttpResponse response = helper.getResponse(remoteUrl(uri));
 
         String value = response.getFirstHeader(HttpHeaders.SET_COOKIE).getValue();
-        return ClientCookieDecoder.STRICT.decode(value);
+        return (DefaultCookie)ClientCookieDecoder.STRICT.decode(value);
     }
 }
diff --git a/moco-runner/src/test/resources/cookie.json b/moco-runner/src/test/resources/cookie.json
index 2678d337f..ea57e55a5 100644
--- a/moco-runner/src/test/resources/cookie.json
+++ b/moco-runner/src/test/resources/cookie.json
@@ -76,5 +76,18 @@
         }
       }
     }
+  },
+  {
+    "request": {
+      "uri": "/cookie-with-same-site"
+    },
+    "response": {
+      "cookies": {
+        "login": {
+          "value": "true",
+          "sameSite": "Lax"
+        }
+      }
+    }
   }
 ]
\ No newline at end of file