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

Open OIDC child session when assuming organization

Browse Source 
OIDC login dropped organization_id before the provider redirect, so
callbacks with an existing matching root session never created an org
child session. Persist organization_id in OIDC state, open the child
session on callback, and forward the parameter from the sign-in UI.

Signed-off-by: Bryan Frimin <bryan@probo.com>
This commit is contained in:
Bryan Frimin
2026-06-17 18:16:47 +02:00
parent ee1439656b
commit 2c8ae26ea1
7 changed files with 202 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
pkg/coredata/migrations/20260617T120000Z.sql Normal file
View File

@@ -0,0 +1,15 @@
-- Copyright (c) 2026 Probo Inc <hello@probo.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.
ALTER TABLE iam_oidc_states ADD COLUMN organization_id TEXT;

37
pkg/coredata/oidc_state.go
View File

@@ -22,32 +22,35 @@ import (
"github.com/jackc/pgx/v5"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/gid"
)
type OIDCState struct {
ID string `db:"id"`
Provider OIDCProvider `db:"provider"`
Nonce string `db:"nonce"`
CodeVerifier string `db:"code_verifier"`
ContinueURL string `db:"continue_url"`
CreatedAt time.Time `db:"created_at"`
ExpiresAt time.Time `db:"expires_at"`
ID string `db:"id"`
Provider OIDCProvider `db:"provider"`
Nonce string `db:"nonce"`
CodeVerifier string `db:"code_verifier"`
ContinueURL string `db:"continue_url"`
OrganizationID *gid.GID `db:"organization_id"`
CreatedAt time.Time `db:"created_at"`
ExpiresAt time.Time `db:"expires_at"`
}
func (s *OIDCState) Insert(ctx context.Context, conn pg.Tx) error {
query := `
INSERT INTO iam_oidc_states (id, provider, nonce, code_verifier, continue_url, created_at, expires_at)
VALUES (@id, @provider, @nonce, @code_verifier, @continue_url, @created_at, @expires_at)
INSERT INTO iam_oidc_states (id, provider, nonce, code_verifier, continue_url, organization_id, created_at, expires_at)
VALUES (@id, @provider, @nonce, @code_verifier, @continue_url, @organization_id, @created_at, @expires_at)
`
args := pgx.StrictNamedArgs{
"id": s.ID,
"provider": s.Provider,
"nonce": s.Nonce,
"code_verifier": s.CodeVerifier,
"continue_url": s.ContinueURL,
"created_at": s.CreatedAt,
"expires_at": s.ExpiresAt,
"id": s.ID,
"provider": s.Provider,
"nonce": s.Nonce,
"code_verifier": s.CodeVerifier,
"continue_url": s.ContinueURL,
"organization_id": s.OrganizationID,
"created_at": s.CreatedAt,
"expires_at": s.ExpiresAt,
}
_, err := conn.Exec(ctx, query, args)
@@ -60,7 +63,7 @@ VALUES (@id, @provider, @nonce, @code_verifier, @continue_url, @created_at, @exp
func (s *OIDCState) LoadByIDForUpdate(ctx context.Context, conn pg.Tx, id string) error {
query := `
SELECT id, provider, nonce, code_verifier, continue_url, created_at, expires_at
SELECT id, provider, nonce, code_verifier, continue_url, organization_id, created_at, expires_at
FROM iam_oidc_states
WHERE id = @id
FOR UPDATE

44
pkg/iam/oidc/service.go
View File

@@ -274,6 +274,7 @@ func (s *Service) InitiateLogin(
ctx context.Context,
provider coredata.OIDCProvider,
continueURL string,
organizationID *gid.GID,
) (string, error) {
info, ok := s.providers[provider]
if !ok {
@@ -297,13 +298,14 @@ func (s *Service) InitiateLogin(
now := time.Now()
oidcState := &coredata.OIDCState{
ID: state,
Provider: provider,
Nonce: nonce,
CodeVerifier: codeVerifier,
ContinueURL: continueURL,
CreatedAt: now,
ExpiresAt: now.Add(10 * time.Minute),
ID: state,
Provider: provider,
Nonce: nonce,
CodeVerifier: codeVerifier,
ContinueURL: continueURL,
OrganizationID: organizationID,
CreatedAt: now,
ExpiresAt: now.Add(10 * time.Minute),
}
err = s.pg.WithTx(
@@ -337,10 +339,10 @@ func (s *Service) HandleCallback(
provider coredata.OIDCProvider,
stateParam string,
code string,
) (*coredata.Identity, string, error) {
) (*coredata.Identity, string, *gid.GID, error) {
info, ok := s.providers[provider]
if !ok {
return nil, "", NewProviderNotEnabledError(provider)
return nil, "", nil, NewProviderNotEnabledError(provider)
}
var oidcState coredata.OIDCState
@@ -364,15 +366,15 @@ func (s *Service) HandleCallback(
},
)
if err != nil {
return nil, "", err
return nil, "", nil, err
}
if time.Now().After(oidcState.ExpiresAt) {
return nil, "", NewInvalidStateError()
return nil, "", nil, NewInvalidStateError()
}
if oidcState.Provider != provider {
return nil, "", NewInvalidStateError()
return nil, "", nil, NewInvalidStateError()
}
token, err := info.oauth2Config.Exchange(
@@ -381,34 +383,34 @@ func (s *Service) HandleCallback(
oauth2.SetAuthURLParam("code_verifier", oidcState.CodeVerifier),
)
if err != nil {
return nil, "", NewCodeExchangeError(err)
return nil, "", nil, NewCodeExchangeError(err)
}
rawIDToken, ok := token.Extra("id_token").(string)
if !ok {
return nil, "", NewIDTokenMissingError()
return nil, "", nil, NewIDTokenMissingError()
}
claims, err := s.verifyAndParseIDToken(ctx, info, rawIDToken, oidcState.Nonce)
if err != nil {
return nil, "", fmt.Errorf("cannot verify id token: %w", err)
return nil, "", nil, fmt.Errorf("cannot verify id token: %w", err)
}
if claims.Email == "" {
return nil, "", NewMissingEmailClaimError()
return nil, "", nil, NewMissingEmailClaimError()
}
if !info.trustProviderEmail && !claims.isEmailVerified() {
return nil, "", NewEmailNotVerifiedError()
return nil, "", nil, NewEmailNotVerifiedError()
}
if !info.enterpriseChecker(claims) {
return nil, "", NewPersonalAccountNotAllowedError()
return nil, "", nil, NewPersonalAccountNotAllowedError()
}
email, err := mail.ParseAddr(claims.Email)
if err != nil {
return nil, "", fmt.Errorf("cannot parse email from id token: %w", err)
return nil, "", nil, fmt.Errorf("cannot parse email from id token: %w", err)
}
var identity *coredata.Identity
@@ -455,10 +457,10 @@ func (s *Service) HandleCallback(
},
)
if err != nil {
return nil, "", err
return nil, "", nil, err
}
return identity, oidcState.ContinueURL, nil
return identity, oidcState.ContinueURL, oidcState.OrganizationID, nil
}
func (s *Service) verifyAndParseIDToken(ctx context.Context, info *providerInfo, rawIDToken string, expectedNonce string) (*idTokenClaims, error) {

93
pkg/iam/session_service.go
View File

@@ -505,6 +505,99 @@ func (s SessionService) OpenSAMLChildSessionForOrganization(
return childSession, membership, nil
}
// OpenOIDCChildSessionForOrganization creates an OIDC-authenticated child session for the
// given organization under the provided root session.
func (s SessionService) OpenOIDCChildSessionForOrganization(
ctx context.Context,
rootSessionID gid.GID,
organizationID gid.GID,
) (*coredata.Session, *coredata.Membership, error) {
var (
now = time.Now()
rootSession = &coredata.Session{}
identity = &coredata.Identity{}
profile = &coredata.MembershipProfile{}
membership = &coredata.Membership{}
childSession = &coredata.Session{}
scope = coredata.NewScopeFromObjectID(organizationID)
)
err := s.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
err := rootSession.LoadByID(ctx, tx, rootSessionID)
if err != nil {
if err == coredata.ErrResourceNotFound {
return NewSessionNotFoundError(rootSessionID)
}
return fmt.Errorf("cannot load session: %w", err)
}
if !rootSession.IsRootSession() {
return fmt.Errorf("session %q is not a root session", rootSessionID)
}
if rootSession.ExpireReason != nil || now.After(rootSession.ExpiredAt) {
return NewSessionExpiredError(rootSessionID)
}
err = identity.LoadByID(ctx, tx, rootSession.IdentityID)
if err != nil {
return fmt.Errorf("cannot load identity: %w", err)
}
err = profile.LoadByIdentityIDAndOrganizationID(ctx, tx, scope, rootSession.IdentityID, organizationID)
if err != nil {
if err == coredata.ErrResourceNotFound {
return NewProfileNotFoundError(gid.Nil)
}
return fmt.Errorf("cannot load profile: %w", err)
}
if profile.State == coredata.ProfileStateInactive {
return NewUserInactiveError(profile.ID)
}
err = membership.LoadByIdentityIDAndOrganizationID(ctx, tx, scope, rootSession.IdentityID, organizationID)
if err != nil {
if err == coredata.ErrResourceNotFound {
return NewMembershipNotFoundError(organizationID)
}
return fmt.Errorf("cannot load membership: %w", err)
}
tenantID := scope.GetTenantID()
childSession = &coredata.Session{
ID: gid.New(tenantID, coredata.SessionEntityType),
IdentityID: rootSession.IdentityID,
TenantID: &tenantID,
MembershipID: &membership.ID,
ParentSessionID: &rootSession.ID,
AuthMethod: coredata.AuthMethodOIDC,
AuthenticatedAt: now,
ExpiredAt: rootSession.ExpiredAt,
CreatedAt: now,
UpdatedAt: now,
}
err = childSession.Insert(ctx, tx)
if err != nil {
return fmt.Errorf("cannot insert child session: %w", err)
}
return nil
},
)
if err != nil {
return nil, nil, err
}
return childSession, membership, nil
}
func (s SessionService) AssumeOrganizationSession(
ctx context.Context,
sessionID gid.GID,

34
pkg/server/api/connect/v1/oidc_handler.go
View File

@@ -25,6 +25,7 @@ import (
"go.gearno.de/kit/httpserver"
"go.gearno.de/kit/log"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/iam"
"go.probo.inc/probo/pkg/saferedirect"
"go.probo.inc/probo/pkg/securecookie"
@@ -77,7 +78,19 @@ func (h *OIDCHandler) LoginHandler(w http.ResponseWriter, r *http.Request) {
continueURL := r.URL.Query().Get("continue")
authURL, err := h.iam.OIDCService.InitiateLogin(ctx, provider, continueURL)
var organizationID *gid.GID
if organizationIDParam := r.URL.Query().Get("organization_id"); organizationIDParam != "" {
parsedOrganizationID, err := gid.ParseGID(organizationIDParam)
if err != nil {
httpserver.RenderError(w, http.StatusBadRequest, errors.New("invalid organization_id parameter"))
return
}
organizationID = &parsedOrganizationID
}
authURL, err := h.iam.OIDCService.InitiateLogin(ctx, provider, continueURL, organizationID)
if err != nil {
h.logger.ErrorCtx(ctx, "cannot initiate OIDC login", log.Error(err))
httpserver.RenderError(w, http.StatusInternalServerError, errors.New("internal server error"))
@@ -118,7 +131,7 @@ func (h *OIDCHandler) CallbackHandler(w http.ResponseWriter, r *http.Request) {
return
}
identity, continueURL, err := h.iam.OIDCService.HandleCallback(ctx, provider, stateParam, code)
identity, continueURL, organizationID, err := h.iam.OIDCService.HandleCallback(ctx, provider, stateParam, code)
if err != nil {
h.logger.ErrorCtx(ctx, "cannot handle OIDC callback", log.Error(err))
httpserver.RenderError(w, http.StatusUnauthorized, errors.New("authentication failed"))
@@ -155,9 +168,24 @@ func (h *OIDCHandler) CallbackHandler(w http.ResponseWriter, r *http.Request) {
}
}
if organizationID != nil {
_, _, err = h.iam.SessionService.OpenOIDCChildSessionForOrganization(ctx, rootSession.ID, *organizationID)
if err != nil {
h.logger.ErrorCtx(ctx, "cannot open OIDC child session", log.Error(err))
httpserver.RenderError(w, http.StatusInternalServerError, errors.New("internal server error"))
return
}
}
h.sessionCookie.Set(w, rootSession)
redirectURL := h.safeRedirect.GetSafeRedirectURL(ctx, continueURL, "/")
defaultRedirect := "/"
if organizationID != nil {
defaultRedirect = "/organizations/" + organizationID.String()
}
redirectURL := h.safeRedirect.GetSafeRedirectURL(ctx, continueURL, defaultRedirect)
if transferURL, ok := h.buildSessionTransferURL(ctx, redirectURL, rootSession.ID.String()); ok {
http.Redirect(w, r, transferURL, http.StatusFound)