rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Refacto session cookie handling

Browse Source 
Signed-off-by: Sacha Al Himdani <sacha@getprobo.com>
This commit is contained in:
Sacha Al Himdani
2025-08-05 11:34:59 +02:00
parent 05123185e3
commit 330ba0c18a
3 changed files with 164 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
pkg/server/api/console/v1/resolver.go
View File

@@ -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(
authCfg.CookieName,
authCfg.CookieSecret,
))
if err != nil {
if !errors.Is(err, securecookie.ErrCookieNotFound) {
sessionAuthCfg := session.AuthConfig{
CookieName: authCfg.CookieName,
CookieSecret: authCfg.CookieSecret,
}
errorHandler := session.ErrorHandler{
OnCookieError: func(err error) {
panic(fmt.Errorf("failed to get session: %w", err))
}
},
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) {
panic(fmt.Errorf("failed to list tenants for user: %w", err))
},
}
authResult := session.TryAuth(ctx, w, r, usrmgrSvc, sessionAuthCfg, errorHandler)
if authResult == nil {
next(w, r)
return
}
sessionID, err := gid.ParseGID(cookieValue)
if err != nil {
securecookie.Clear(w, securecookie.DefaultConfig(
authCfg.CookieName,
authCfg.CookieSecret,
))
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)
ctx = context.WithValue(ctx, userContextKey, user)
ctx = context.WithValue(ctx, userTenantContextKey, &tenantIDs)
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))
}
}

66
pkg/server/api/trust/v1/resolver.go
View File

@@ -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(
authCfg.CookieName,
authCfg.CookieSecret,
))
if err != nil {
sessionAuthCfg := session.AuthConfig{
CookieName: authCfg.CookieName,
CookieSecret: authCfg.CookieSecret,
}
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)
if err != nil {
clearSessionCookie(w, authCfg)
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)
ctx = context.WithValue(ctx, sessionContextKey, authResult.Session)
ctx = context.WithValue(ctx, userContextKey, authResult.User)
ctx = context.WithValue(ctx, userTenantContextKey, &authResult.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,

110
pkg/server/session/session.go Normal file
View File

@@ -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,
))
}