Reuse existing session when signing in in for assume

Signed-off-by: Émile Ré <emile@getprobo.com>

pkg/iam/auth_service.go

@@ -472,19 +472,20 @@ func (s AuthService) OpenSessionWithSAML(ctx context.Context, identityID gid.GID
 	return session, nil
 }
 
-func (s AuthService) OpenSessionWithPassword(ctx context.Context, email mail.Addr, password string) (*coredata.Identity, *coredata.Session, error) {
+func (s AuthService) CheckCredentials(
+	ctx context.Context,
+	email mail.Addr,
+	password string,
+) (*coredata.Identity, error) {
 	v := validator.New()
 	v.Check(password, "password", PasswordValidator())
 
 	err := v.Error()
 	if err != nil {
-		return nil, nil, NewInvalidPasswordError("invalid password")
+		return nil, NewInvalidPasswordError("invalid password")
 	}
 
-	var (
+	identity := &coredata.Identity{}
-		identity = &coredata.Identity{}
-		session  = &coredata.Session{}
-	)
 
 	err = s.pg.WithTx(
 		ctx,
@@ -513,7 +514,20 @@ func (s AuthService) OpenSessionWithPassword(ctx context.Context, email mail.Add
 				return NewInvalidCredentialsError("invalid email or password")
 			}
 
-			session = coredata.NewRootSession(identity.ID, coredata.AuthMethodPassword, s.sessionDuration)
+			return nil
+		},
+	)
+
+	return identity, err
+}
+
+func (s AuthService) OpenSessionWithPassword(ctx context.Context, identityID gid.GID) (*coredata.Session, error) {
+	session := &coredata.Session{}
+
+	err := s.pg.WithTx(
+		ctx,
+		func(conn pg.Conn) (err error) {
+			session = coredata.NewRootSession(identityID, coredata.AuthMethodPassword, s.sessionDuration)
 			err = session.Insert(ctx, conn)
 			if err != nil {
 				return fmt.Errorf("cannot insert session: %w", err)
@@ -523,7 +537,11 @@ func (s AuthService) OpenSessionWithPassword(ctx context.Context, email mail.Add
 		},
 	)
 
-	return identity, session, err
+	if err != nil {
+		return nil, err
+	}
+
+	return session, nil
 }
 
 func (s AuthService) SendMagicLink(ctx context.Context, req *SendMagicLinkRequest) error {

pkg/iam/session_service.go

@@ -315,6 +315,93 @@ func (s SessionService) GetActiveSessionForMembership(ctx context.Context, rootS
 	return childSession, nil
 }
 
+func (s SessionService) OpenPasswordChildSessionForOrganization(
+	ctx context.Context,
+	rootSessionID gid.GID,
+	organizationID gid.GID,
+) (*coredata.Session, *coredata.Membership, error) {
+	var (
+		now          = time.Now()
+		rootSession  = &coredata.Session{}
+		identity     = &coredata.Identity{}
+		membership   = &coredata.Membership{}
+		childSession = &coredata.Session{}
+		scope        = coredata.NewScopeFromObjectID(organizationID)
+	)
+
+	err := s.pg.WithTx(
+		ctx,
+		func(tx pg.Conn) 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 = membership.LoadByIdentityInOrganization(ctx, tx, rootSession.IdentityID, organizationID)
+			if err != nil {
+				if err == coredata.ErrResourceNotFound {
+					return NewMembershipNotFoundError(organizationID)
+				}
+				return fmt.Errorf("cannot load membership: %w", err)
+			}
+
+			if membership.State == coredata.MembershipStateInactive {
+				return NewMembershipInactiveError(membership.ID)
+			}
+
+			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.AuthMethodPassword,
+				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)
+			}
+
+			// Change root session auth method to password
+			rootSession.UpdatedAt = now
+			rootSession.AuthMethod = coredata.AuthMethodPassword
+
+			if err := rootSession.Update(ctx, tx); err != nil {
+				return fmt.Errorf("cannot update root session: %w", err)
+			}
+
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return childSession, membership, nil
+}
+
 // OpenSAMLChildSessionForOrganization creates a SAML-authenticated child session for the given
 // organization under the provided root session.
 //

pkg/server/api/connect/v1/schema.graphql

@@ -626,6 +626,8 @@ type PageInfo {
 }
 
 input SignInInput {
+  # When assuming an org with a password auth method
+  organizationId: ID
   email: EmailAddr!
   password: String!
 }

pkg/server/api/connect/v1/schema/schema.go

@@ -3068,6 +3068,8 @@ type PageInfo {
 }
 
 input SignInInput {
+  # When assuming an org with a password auth method
+  organizationId: ID
   email: EmailAddr!
   password: String!
 }
@@ -15772,13 +15774,20 @@ func (ec *executionContext) unmarshalInputSignInInput(ctx context.Context, obj a
 		asMap[k] = v
 	}
 
-	fieldsInOrder := [...]string{"email", "password"}
+	fieldsInOrder := [...]string{"organizationId", "email", "password"}
 	for _, k := range fieldsInOrder {
 		v, ok := asMap[k]
 		if !ok {
 			continue
 		}
 		switch k {
+		case "organizationId":
+			ctx := graphql.WithPathContext(ctx, graphql.NewPathWithField("organizationId"))
+			data, err := ec.unmarshalOID2ᚖgoᚗproboᚗincᚋproboᚋpkgᚋgidᚐGID(ctx, v)
+			if err != nil {
+				return it, err
+			}
+			it.OrganizationID = data
 		case "email":
 			ctx := graphql.WithPathContext(ctx, graphql.NewPathWithField("email"))
 			data, err := ec.unmarshalNEmailAddr2goᚗproboᚗincᚋproboᚋpkgᚋmailᚐAddr(ctx, v)

pkg/server/api/connect/v1/types/types.go

@@ -485,6 +485,7 @@ type SessionOrder struct {
 }
 
 type SignInInput struct {
+	OrganizationID *gid.GID  `json:"organizationId,omitempty"`
 	Email          mail.Addr `json:"email"`
 	Password       string    `json:"password"`
 }

pkg/server/api/connect/v1/v1_resolver.go

@@ -356,9 +356,7 @@ func (r *membershipProfileResolver) Permission(ctx context.Context, obj *types.M
 
 // SignIn is the resolver for the signIn field.
 func (r *mutationResolver) SignIn(ctx context.Context, input types.SignInInput) (*types.SignInPayload, error) {
-	// TODO: handle existing session to only open child session and chnage root session auth method to PASSWORD
+	identity, err := r.iam.AuthService.CheckCredentials(ctx, input.Email, input.Password)
-
-	identity, session, err := r.iam.AuthService.OpenSessionWithPassword(ctx, input.Email, input.Password)
 	if err != nil {
 		var errInvalidPassword *iam.ErrInvalidPassword
 		if errors.As(err, &errInvalidPassword) {
@@ -375,10 +373,59 @@ func (r *mutationResolver) SignIn(ctx context.Context, input types.SignInInput)
 			}
 		}
 
-		r.logger.ErrorCtx(ctx, "cannot sign in", log.Error(err))
+		r.logger.ErrorCtx(ctx, "cannot check credentials", log.Error(err))
 		return nil, gqlutils.Internal(ctx)
 	}
 
+	session := authn.SessionFromContext(ctx)
+
+	switch {
+	case session == nil:
+		var err error
+		session, err = r.iam.AuthService.OpenSessionWithPassword(
+			ctx,
+			identity.ID,
+		)
+		if err != nil {
+			r.logger.ErrorCtx(ctx, "cannot create session", log.Error(err))
+			return nil, gqlutils.Internal(ctx)
+		}
+	case session.IdentityID != identity.ID:
+		if err := r.iam.SessionService.CloseSession(ctx, session.ID); err != nil {
+			r.logger.ErrorCtx(ctx, "cannot close session", log.Error(err))
+			return nil, gqlutils.Internal(ctx)
+		}
+
+		session, err = r.iam.AuthService.OpenSessionWithPassword(
+			ctx,
+			identity.ID,
+		)
+		if err != nil {
+			r.logger.ErrorCtx(ctx, "cannot create session", log.Error(err))
+			return nil, gqlutils.Internal(ctx)
+		}
+	}
+
+	if input.OrganizationID != nil {
+		var err error
+		session, _, err = r.iam.SessionService.OpenPasswordChildSessionForOrganization(ctx, session.ID, *input.OrganizationID)
+		if err != nil {
+			var errSessionExpired *iam.ErrSessionExpired
+			var errMembershipNotFound *iam.ErrMembershipNotFound
+			var errMembershipInactive *iam.ErrMembershipInactive
+
+			if errors.As(err, errSessionExpired) {
+				return nil, gqlutils.Unauthenticated(ctx, err)
+			}
+			if errors.As(err, errMembershipNotFound) || errors.As(err, errMembershipInactive) {
+				return nil, gqlutils.Forbidden(ctx, err)
+			}
+
+			r.logger.ErrorCtx(ctx, "cannot assume organization", log.Error(err))
+			return nil, gqlutils.Internal(ctx)
+		}
+	}
+
 	w := gqlutils.HTTPResponseWriterFromContext(ctx)
 	r.sessionCookie.Set(w, session)