diff --git a/pkg/server/api/console/v1/resolver.go b/pkg/server/api/console/v1/resolver.go index 6502dfe62..e62226de1 100644 --- a/pkg/server/api/console/v1/resolver.go +++ b/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( - 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)) } } diff --git a/pkg/server/api/trust/v1/resolver.go b/pkg/server/api/trust/v1/resolver.go index 59e518d7f..bb27dc47c 100644 --- a/pkg/server/api/trust/v1/resolver.go +++ b/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( - 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, diff --git a/pkg/server/session/session.go b/pkg/server/session/session.go new file mode 100644 index 000000000..68b9c9a04 --- /dev/null +++ b/pkg/server/session/session.go @@ -0,0 +1,110 @@ +// Copyright (c) 2025 Probo Inc . +// +// 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, + )) +}