From da6a569e859f4dfbbebfd084060e15d940de8861 Mon Sep 17 00:00:00 2001
From: Sam Ottenhoff <ottenhoff@longsight.com>
Date: Fri, 23 Feb 2024 14:45:58 -0500
Subject: [PATCH] reverseproxy: cookie should be Secure and SameSite=None when
 TLS (#6115)

* reverseproxy: cookie should be Secure and SameSite=None when TLS

* Update modules/caddyhttp/reverseproxy/selectionpolicies_test.go

Co-authored-by: Mohammed Al Sahaf <mohammed@caffeinatedwonders.com>

---------

Co-authored-by: Mohammed Al Sahaf <mohammed@caffeinatedwonders.com>
---
 .../reverseproxy/selectionpolicies.go         | 14 ++++-
 .../reverseproxy/selectionpolicies_test.go    | 54 +++++++++++++++++++
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/modules/caddyhttp/reverseproxy/selectionpolicies.go b/modules/caddyhttp/reverseproxy/selectionpolicies.go
index b56c8074..b6f807c1 100644
--- a/modules/caddyhttp/reverseproxy/selectionpolicies.go
+++ b/modules/caddyhttp/reverseproxy/selectionpolicies.go
@@ -655,12 +655,22 @@ func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http
 		if err != nil {
 			return upstream
 		}
-		http.SetCookie(w, &http.Cookie{
+		cookie := &http.Cookie{
 			Name:   s.Name,
 			Value:  sha,
 			Path:   "/",
 			Secure: false,
-		})
+		}
+		isProxyHttps := false
+		if trusted, ok := caddyhttp.GetVar(req.Context(), caddyhttp.TrustedProxyVarKey).(bool); ok && trusted {
+			xfp, xfpOk, _ := lastHeaderValue(req.Header, "X-Forwarded-Proto")
+			isProxyHttps = xfpOk && xfp == "https"
+		}
+		if req.TLS != nil || isProxyHttps {
+			cookie.Secure = true
+			cookie.SameSite = http.SameSiteNoneMode
+		}
+		http.SetCookie(w, cookie)
 		return upstream
 	}
 
diff --git a/modules/caddyhttp/reverseproxy/selectionpolicies_test.go b/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
index 9199f619..d7e79626 100644
--- a/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
+++ b/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
@@ -658,6 +658,9 @@ func TestCookieHashPolicy(t *testing.T) {
 	if cookieServer1.Name != "lb" {
 		t.Error("cookieHashPolicy should set a cookie with name lb")
 	}
+	if cookieServer1.Secure {
+		t.Error("cookieHashPolicy should set cookie Secure attribute to false when request is not secure")
+	}
 	if h != pool[0] {
 		t.Error("Expected cookieHashPolicy host to be the first only available host.")
 	}
@@ -687,6 +690,57 @@ func TestCookieHashPolicy(t *testing.T) {
 	}
 }
 
+func TestCookieHashPolicyWithSecureRequest(t *testing.T) {
+    ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
+    defer cancel()
+    cookieHashPolicy := CookieHashSelection{}
+    if err := cookieHashPolicy.Provision(ctx); err != nil {
+        t.Errorf("Provision error: %v", err)
+        t.FailNow()
+    }
+
+    pool := testPool()
+    pool[0].Dial = "localhost:8080"
+    pool[1].Dial = "localhost:8081"
+    pool[2].Dial = "localhost:8082"
+    pool[0].setHealthy(true)
+    pool[1].setHealthy(false)
+    pool[2].setHealthy(false)
+
+    // Create a test server that serves HTTPS requests
+    ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+        h := cookieHashPolicy.Select(pool, r, w)
+        if h != pool[0] {
+            t.Error("Expected cookieHashPolicy host to be the first only available host.")
+        }
+    }))
+    defer ts.Close()
+
+    // Make a new HTTPS request to the test server
+    client := ts.Client()
+    request, err := http.NewRequest(http.MethodGet, ts.URL+"/test", nil)
+    if err != nil {
+        t.Fatal(err)
+    }
+    response, err := client.Do(request)
+    if err != nil {
+        t.Fatal(err)
+    }
+
+    // Check if the cookie set is Secure and has SameSiteNone mode
+    cookies := response.Cookies()
+    if len(cookies) == 0 {
+        t.Fatal("Expected a cookie to be set")
+    }
+    cookie := cookies[0]
+    if !cookie.Secure {
+        t.Error("Expected cookie Secure attribute to be true when request is secure")
+    }
+    if cookie.SameSite != http.SameSiteNoneMode {
+        t.Error("Expected cookie SameSite attribute to be None when request is secure")
+    }
+}
+
 func TestCookieHashPolicyWithFirstFallback(t *testing.T) {
 	ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
 	defer cancel()