From 01ee7c4dc848efb9ac7712a3035d109057567de3 Mon Sep 17 00:00:00 2001
From: Lukasz Raczylo <lukasz@raczylo.com>
Date: Fri, 6 Dec 2024 23:43:47 +0000
Subject: [PATCH] Improve cookie setting.

---
 helpers.go   | 14 ++++++++++++--
 main.go      | 18 +++++++++---------
 main_test.go |  2 ++
 settings.go  | 10 ----------
 4 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/helpers.go b/helpers.go
index 216a593..2a3765b 100644
--- a/helpers.go
+++ b/helpers.go
@@ -17,6 +17,16 @@ import (
 	"github.com/gorilla/sessions"
 )
 
+func newSessionOptions(isSecure bool) *sessions.Options {
+	return &sessions.Options{
+			HttpOnly: true,
+			Secure:   isSecure,
+			SameSite: http.SameSiteLaxMode,
+			MaxAge:   ConstSessionTimeout,
+			Path:     "/",
+	}
+}
+
 // generateNonce generates a random nonce
 func generateNonce() (string, error) {
 	nonceBytes := make([]byte, 32)
@@ -101,7 +111,7 @@ func (t *TraefikOidc) handleExpiredToken(rw http.ResponseWriter, req *http.Reque
 	session.Values["csrf"] = uuid.New().String()
 	session.Values["incoming_path"] = req.URL.Path
 	session.Values["nonce"], _ = generateNonce()
-	session.Options = defaultSessionOptions
+	session.Options = newSessionOptions(t.determineScheme(req) == "https")
 
 	// Save the session before initiating authentication
 	if err := session.Save(req, rw); err != nil {
@@ -222,7 +232,7 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 	session.Values["email"] = email
 	session.Values["id_token"] = idToken
 	session.Values["refresh_token"] = tokenResponse.RefreshToken
-	session.Options = defaultSessionOptions
+	session.Options = newSessionOptions(t.determineScheme(req) == "https")
 
 	// Remove CSRF and nonce from session
 	delete(session.Values, "csrf")
diff --git a/main.go b/main.go
index 7be689a..3d5e9a9 100644
--- a/main.go
+++ b/main.go
@@ -186,7 +186,9 @@ func (t *TraefikOidc) VerifyJWTSignatureAndClaims(jwt *JWT, token string) error
 // New creates a new instance of the OIDC middleware
 func New(ctx context.Context, next http.Handler, config *Config, name string) (http.Handler, error) {
 	store := sessions.NewCookieStore([]byte(config.SessionEncryptionKey))
-	store.Options = defaultSessionOptions
+	store.Options = newSessionOptions(func() bool {
+		return config.ForceHTTPS
+	}())
 
 	// Setup HTTP client
 	transport := &http.Transport{
@@ -200,7 +202,7 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 		},
 		ForceAttemptHTTP2:     true,
 		TLSHandshakeTimeout:   10 * time.Second,
-		ExpectContinueTimeout: 1 * time.Second,
+		ExpectContinueTimeout: 0,
 		MaxIdleConns:          100,
 		MaxIdleConnsPerHost:   100,
 		IdleConnTimeout:       90 * time.Second,
@@ -377,12 +379,9 @@ func (t *TraefikOidc) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	}
 
 	// Determine the scheme (http/https) and host
-	t.scheme = t.determineScheme(req)
-	defaultSessionOptions.Secure = t.scheme == "https"
+	scheme := t.determineScheme(req)
 	host := t.determineHost(req)
-
-	redirectURL := buildFullURL(t.scheme, host, t.redirURLPath)
-
+	redirectURL := buildFullURL(scheme, host, t.redirURLPath)
 	// Build the redirect URL if not already set
 	if redirectURL == "" {
 		redirectURL = buildFullURL(t.scheme, host, t.redirURLPath)
@@ -397,6 +396,7 @@ func (t *TraefikOidc) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
+	session.Options = newSessionOptions(scheme == "https")
 	t.logger.Debugf("Session contents at start: %+v", session.Values)
 
 	// Handle logout URL
@@ -584,7 +584,7 @@ func (t *TraefikOidc) defaultInitiateAuthentication(rw http.ResponseWriter, req
 	csrfToken := uuid.New().String()
 	session.Values["csrf"] = csrfToken
 	session.Values["incoming_path"] = req.URL.Path
-	session.Options = defaultSessionOptions
+	session.Options = newSessionOptions(t.determineScheme(req) == "https")
 	t.logger.Debugf("Setting CSRF token: %s", csrfToken)
 
 	// Generate nonce
@@ -710,7 +710,7 @@ func (t *TraefikOidc) refreshToken(rw http.ResponseWriter, req *http.Request, se
 	// Update session with new tokens
 	session.Values["id_token"] = newToken.IDToken
 	session.Values["refresh_token"] = newToken.RefreshToken
-	session.Options = defaultSessionOptions
+	session.Options = newSessionOptions(t.determineScheme(req) == "https")
 	if err := session.Save(req, rw); err != nil {
 		t.logger.Errorf("Failed to save refreshed session: %v", err)
 		return false
diff --git a/main_test.go b/main_test.go
index 530d28f..7f73da4 100644
--- a/main_test.go
+++ b/main_test.go
@@ -1223,6 +1223,8 @@ func TestHandleExpiredToken(t *testing.T) {
 				t.Error("Nonce not set")
 			}
 
+			defaultSessionOptions := newSessionOptions(tOidc.determineScheme(req) == "https")
+
 			// Verify session options
 			if session.Options.MaxAge != defaultSessionOptions.MaxAge {
 				t.Error("Session MaxAge not set correctly")
diff --git a/settings.go b/settings.go
index c1f11c1..e3d37de 100644
--- a/settings.go
+++ b/settings.go
@@ -6,8 +6,6 @@ import (
 	"log"
 	"net/http"
 	"os"
-
-	"github.com/gorilla/sessions"
 )
 
 const (
@@ -35,14 +33,6 @@ type Config struct {
 	HTTPClient            *http.Client
 }
 
-var defaultSessionOptions = &sessions.Options{
-	HttpOnly: true,
-	Secure:   false,
-	SameSite: http.SameSiteLaxMode,
-	MaxAge:   ConstSessionTimeout,
-	Path:     "/",
-}
-
 // CreateConfig creates a new Config with default values
 func CreateConfig() *Config {
 	c := &Config{}