Refacto session cookie handling

Signed-off-by: Sacha Al Himdani <sacha@getprobo.com>

pkg/server/api/console/v1/resolver.go

@@ -19,7 +19,6 @@ package console_v1
 import (
 	"context"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"net/http"
 	"strings"
@@ -35,9 +34,9 @@ import (
 	"github.com/getprobo/probo/pkg/gid"
 	"github.com/getprobo/probo/pkg/probo"
 	"github.com/getprobo/probo/pkg/saferedirect"
-	"github.com/getprobo/probo/pkg/securecookie"
 	"github.com/getprobo/probo/pkg/server/api/console/v1/schema"
 	gqlutils "github.com/getprobo/probo/pkg/server/graphql"
+	"github.com/getprobo/probo/pkg/server/session"
 	"github.com/getprobo/probo/pkg/statelesstoken"
 	"github.com/getprobo/probo/pkg/usrmgr"
 	"github.com/go-chi/chi/v5"
@@ -264,65 +263,43 @@ func WithSession(usrmgrSvc *usrmgr.Service, authCfg AuthConfig, next http.Handle
 	return func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 
-		cookieValue, err := securecookie.Get(r, securecookie.DefaultConfig(
+		sessionAuthCfg := session.AuthConfig{
-			authCfg.CookieName,
+			CookieName:   authCfg.CookieName,
-			authCfg.CookieSecret,
+			CookieSecret: authCfg.CookieSecret,
-		))
+		}
-		if err != nil {
+
-			if !errors.Is(err, securecookie.ErrCookieNotFound) {
+		errorHandler := session.ErrorHandler{
+			OnCookieError: func(err error) {
 				panic(fmt.Errorf("failed to get session: %w", err))
-			}
+			},
-
+			OnParseError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
-			next(w, r)
+				session.ClearCookie(w, authCfg)
-			return
+			},
-		}
+			OnSessionError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
-
+				session.ClearCookie(w, authCfg)
-		sessionID, err := gid.ParseGID(cookieValue)
+			},
-		if err != nil {
+			OnUserError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
-			securecookie.Clear(w, securecookie.DefaultConfig(
+				session.ClearCookie(w, authCfg)
-				authCfg.CookieName,
+			},
-				authCfg.CookieSecret,
+			OnTenantError: func(err error) {
-			))
-
-			next(w, r)
-			return
-		}
-
-		session, err := usrmgrSvc.GetSession(ctx, sessionID)
-		if err != nil {
-			securecookie.Clear(w, securecookie.DefaultConfig(
-				authCfg.CookieName,
-				authCfg.CookieSecret,
-			))
-
-			next(w, r)
-			return
-		}
-
-		user, err := usrmgrSvc.GetUserBySession(ctx, sessionID)
-		if err != nil {
-			securecookie.Clear(w, securecookie.DefaultConfig(
-				authCfg.CookieName,
-				authCfg.CookieSecret,
-			))
-
-			next(w, r)
-			return
-		}
-
-		tenantIDs, err := usrmgrSvc.ListTenantsForUserID(ctx, user.ID)
-		if err != nil {
 				panic(fmt.Errorf("failed to list tenants for user: %w", err))
+			},
 		}
 
-		ctx = context.WithValue(ctx, sessionContextKey, session)
+		authResult := session.TryAuth(ctx, w, r, usrmgrSvc, sessionAuthCfg, errorHandler)
-		ctx = context.WithValue(ctx, userContextKey, user)
+		if authResult == nil {
-		ctx = context.WithValue(ctx, userTenantContextKey, &tenantIDs)
+			next(w, r)
+			return
+		}
+
+		ctx = context.WithValue(ctx, sessionContextKey, authResult.Session)
+		ctx = context.WithValue(ctx, userContextKey, authResult.User)
+		ctx = context.WithValue(ctx, userTenantContextKey, &authResult.TenantIDs)
 
 		next(w, r.WithContext(ctx))
 
 		// Update session after the handler completes
-		if err := usrmgrSvc.UpdateSession(ctx, session); err != nil {
+		if err := usrmgrSvc.UpdateSession(ctx, authResult.Session); err != nil {
 			panic(fmt.Errorf("failed to update session: %w", err))
 		}
 	}

pkg/server/api/trust/v1/resolver.go

@@ -28,11 +28,11 @@ import (
 	"github.com/getprobo/probo/pkg/coredata"
 	"github.com/getprobo/probo/pkg/gid"
 	"github.com/getprobo/probo/pkg/probo"
-	"github.com/getprobo/probo/pkg/securecookie"
 	console_v1 "github.com/getprobo/probo/pkg/server/api/console/v1"
 	"github.com/getprobo/probo/pkg/server/api/trust/v1/auth"
 	"github.com/getprobo/probo/pkg/server/api/trust/v1/schema"
 	gqlutils "github.com/getprobo/probo/pkg/server/graphql"
+	"github.com/getprobo/probo/pkg/server/session"
 	"github.com/getprobo/probo/pkg/statelesstoken"
 	"github.com/getprobo/probo/pkg/trust"
 	"github.com/getprobo/probo/pkg/usrmgr"
@@ -167,41 +167,34 @@ func WithSession(usrmgrSvc *usrmgr.Service, trustSvc *trust.Service, authCfg con
 }
 
 func trySessionAuth(ctx context.Context, w http.ResponseWriter, r *http.Request, usrmgrSvc *usrmgr.Service, authCfg console_v1.AuthConfig) context.Context {
-	cookieValue, err := securecookie.Get(r, securecookie.DefaultConfig(
+	sessionAuthCfg := session.AuthConfig{
-		authCfg.CookieName,
+		CookieName:   authCfg.CookieName,
-		authCfg.CookieSecret,
+		CookieSecret: authCfg.CookieSecret,
-	))
+	}
-	if err != nil {
+
+	errorHandler := session.ErrorHandler{
+		OnParseError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
+			session.ClearCookie(w, authCfg)
+		},
+		OnSessionError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
+			session.ClearCookie(w, authCfg)
+		},
+		OnUserError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
+			session.ClearCookie(w, authCfg)
+		},
+		OnTenantError: func(err error) {
+			session.ClearCookie(w, sessionAuthCfg)
+		},
+	}
+
+	authResult := session.TryAuth(ctx, w, r, usrmgrSvc, sessionAuthCfg, errorHandler)
+	if authResult == nil {
 		return nil
 	}
 
-	sessionID, err := gid.ParseGID(cookieValue)
+	ctx = context.WithValue(ctx, sessionContextKey, authResult.Session)
-	if err != nil {
+	ctx = context.WithValue(ctx, userContextKey, authResult.User)
-		clearSessionCookie(w, authCfg)
+	ctx = context.WithValue(ctx, userTenantContextKey, &authResult.TenantIDs)
-		return nil
-	}
-
-	session, err := usrmgrSvc.GetSession(ctx, sessionID)
-	if err != nil {
-		clearSessionCookie(w, authCfg)
-		return nil
-	}
-
-	user, err := usrmgrSvc.GetUserBySession(ctx, sessionID)
-	if err != nil {
-		clearSessionCookie(w, authCfg)
-		return nil
-	}
-
-	tenantIDs, err := usrmgrSvc.ListTenantsForUserID(ctx, user.ID)
-	if err != nil {
-		clearSessionCookie(w, authCfg)
-		return nil
-	}
-
-	ctx = context.WithValue(ctx, sessionContextKey, session)
-	ctx = context.WithValue(ctx, userContextKey, user)
-	ctx = context.WithValue(ctx, userTenantContextKey, &tenantIDs)
 
 	return ctx
 }
@@ -234,13 +227,6 @@ func tryTokenAuth(ctx context.Context, w http.ResponseWriter, r *http.Request, t
 	return context.WithValue(ctx, tokenAccessContextKey, tokenAccess)
 }
 
-func clearSessionCookie(w http.ResponseWriter, authCfg console_v1.AuthConfig) {
-	securecookie.Clear(w, securecookie.DefaultConfig(
-		authCfg.CookieName,
-		authCfg.CookieSecret,
-	))
-}
-
 func clearTokenCookie(w http.ResponseWriter, trustAuthCfg TrustAuthConfig) {
 	http.SetCookie(w, &http.Cookie{
 		Name:     trustAuthCfg.CookieName,

pkg/server/session/session.go

@@ -0,0 +1,110 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package session
+
+import (
+	"context"
+	"errors"
+	"net/http"
+
+	"github.com/getprobo/probo/pkg/coredata"
+	"github.com/getprobo/probo/pkg/gid"
+	"github.com/getprobo/probo/pkg/securecookie"
+	"github.com/getprobo/probo/pkg/usrmgr"
+)
+
+type AuthConfig struct {
+	CookieName   string
+	CookieSecret string
+}
+
+type AuthResult struct {
+	Session   *coredata.Session
+	User      *coredata.User
+	TenantIDs []gid.TenantID
+}
+
+type ErrorHandler struct {
+	OnCookieError  func(err error)
+	OnParseError   func(w http.ResponseWriter, authCfg AuthConfig)
+	OnSessionError func(w http.ResponseWriter, authCfg AuthConfig)
+	OnUserError    func(w http.ResponseWriter, authCfg AuthConfig)
+	OnTenantError  func(err error)
+}
+
+func TryAuth(
+	ctx context.Context,
+	w http.ResponseWriter,
+	r *http.Request,
+	usrmgrSvc *usrmgr.Service,
+	authCfg AuthConfig,
+	errorHandler ErrorHandler,
+) *AuthResult {
+	cookieValue, err := securecookie.Get(r, securecookie.DefaultConfig(
+		authCfg.CookieName,
+		authCfg.CookieSecret,
+	))
+	if err != nil {
+		if !errors.Is(err, securecookie.ErrCookieNotFound) && errorHandler.OnCookieError != nil {
+			errorHandler.OnCookieError(err)
+		}
+		return nil
+	}
+
+	sessionID, err := gid.ParseGID(cookieValue)
+	if err != nil {
+		if errorHandler.OnParseError != nil {
+			errorHandler.OnParseError(w, authCfg)
+		}
+		return nil
+	}
+
+	session, err := usrmgrSvc.GetSession(ctx, sessionID)
+	if err != nil {
+		if errorHandler.OnSessionError != nil {
+			errorHandler.OnSessionError(w, authCfg)
+		}
+		return nil
+	}
+
+	user, err := usrmgrSvc.GetUserBySession(ctx, sessionID)
+	if err != nil {
+		if errorHandler.OnUserError != nil {
+			errorHandler.OnUserError(w, authCfg)
+		}
+		return nil
+	}
+
+	tenantIDs, err := usrmgrSvc.ListTenantsForUserID(ctx, user.ID)
+	if err != nil {
+		if errorHandler.OnTenantError != nil {
+			errorHandler.OnTenantError(err)
+		}
+		return nil
+	}
+
+	return &AuthResult{
+		Session:   session,
+		User:      user,
+		TenantIDs: tenantIDs,
+	}
+}
+
+func ClearCookie(w http.ResponseWriter, authCfg AuthConfig) {
+	securecookie.Clear(w, securecookie.DefaultConfig(
+		authCfg.CookieName,
+		authCfg.CookieSecret,
+	))
+}