Add assume org session on api

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

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

@@ -59,7 +59,7 @@ type Query {
 }
 
 type Mutation {
-  signIn(input: SignInInput!): SignInPayload! @session(required: NONE)
+  signIn(input: SignInInput!): SignInPayload! @session(required: OPTIONAL)
   signUp(input: SignUpInput!): SignUpPayload! @session(required: NONE)
   signOut: SignOutPayload! @session(required: PRESENT)
   signUpFromInvitation(
@@ -75,6 +75,9 @@ type Mutation {
     @session(required: PRESENT)
   changeEmail(input: ChangeEmailInput!): ChangeEmailPayload!
     @session(required: PRESENT)
+  assumeOrganizationSession(
+    input: AssumeOrganizationSessionInput!
+  ): AssumeOrganizationSessionPayload! @session(required: PRESENT)
 
   updateIdentityProfile(
     input: UpdateIdentityProfileInput!
@@ -417,6 +420,12 @@ enum ProvisioningSource {
   SAML
 }
 
+enum ReauthenticationReason {
+  SESSION_EXPIRED
+  SENSITIVE_ACTION
+  POLICY_REQUIREMENT
+}
+
 enum MembershipOrderField
   @goModel(model: "go.probo.inc/probo/pkg/coredata.MembershipOrderField") {
   ROLE
@@ -579,6 +588,10 @@ input ChangeEmailInput {
   password: String!
 }
 
+input AssumeOrganizationSessionInput {
+  organizationId: ID!
+}
+
 input DeactivateAccountInput {
   password: String!
 }
@@ -709,6 +722,7 @@ input DeleteSAMLConfigurationInput {
 
 type SignInPayload {
   identity: Identity
+  session: Session
 }
 
 type SignUpPayload {
@@ -743,6 +757,29 @@ type ChangeEmailPayload {
   success: Boolean!
 }
 
+union AssumeOrganizationSessionResult =
+  | OrganizationSessionCreated
+  | PasswordRequired
+  | SAMLAuthenticationRequired
+
+type OrganizationSessionCreated {
+  session: Session!
+  membership: Membership!
+}
+
+type PasswordRequired {
+  reason: ReauthenticationReason!
+}
+
+type SAMLAuthenticationRequired {
+  reason: ReauthenticationReason!
+  redirectUrl: String!
+}
+
+type AssumeOrganizationSessionPayload {
+  result: AssumeOrganizationSessionResult!
+}
+
 type DeactivateAccountPayload {
   success: Boolean!
 }

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

@@ -16,6 +16,10 @@ import (
 	"go.probo.inc/probo/pkg/page"
 )
 
+type AssumeOrganizationSessionResult interface {
+	IsAssumeOrganizationSessionResult()
+}
+
 type Node interface {
 	IsNode()
 	GetID() gid.GID
@@ -42,6 +46,14 @@ type Application struct {
 	AvailableAccessLevels []AccessLevel `json:"availableAccessLevels"`
 }
 
+type AssumeOrganizationSessionInput struct {
+	OrganizationID gid.GID `json:"organizationId"`
+}
+
+type AssumeOrganizationSessionPayload struct {
+	Result AssumeOrganizationSessionResult `json:"result"`
+}
+
 type ChangeEmailInput struct {
 	NewEmail mail.Addr `json:"newEmail"`
 	Password string    `json:"password"`
@@ -263,6 +275,13 @@ type Organization struct {
 func (Organization) IsNode()             {}
 func (this Organization) GetID() gid.GID { return this.ID }
 
+type OrganizationSessionCreated struct {
+	Session    *Session    `json:"session"`
+	Membership *Membership `json:"membership"`
+}
+
+func (OrganizationSessionCreated) IsAssumeOrganizationSessionResult() {}
+
 type PageInfo struct {
 	HasNextPage     bool            `json:"hasNextPage"`
 	HasPreviousPage bool            `json:"hasPreviousPage"`
@@ -270,6 +289,12 @@ type PageInfo struct {
 	EndCursor       *page.CursorKey `json:"endCursor,omitempty"`
 }
 
+type PasswordRequired struct {
+	Reason ReauthenticationReason `json:"reason"`
+}
+
+func (PasswordRequired) IsAssumeOrganizationSessionResult() {}
+
 type Permission struct {
 	ID            gid.GID       `json:"id"`
 	CreatedAt     time.Time     `json:"createdAt"`
@@ -365,6 +390,13 @@ type SAMLAttributeMappingsInput struct {
 	Role      *string `json:"role,omitempty"`
 }
 
+type SAMLAuthenticationRequired struct {
+	Reason      ReauthenticationReason `json:"reason"`
+	RedirectURL string                 `json:"redirectUrl"`
+}
+
+func (SAMLAuthenticationRequired) IsAssumeOrganizationSessionResult() {}
+
 type SAMLConfiguration struct {
 	ID                      gid.GID                        `json:"id"`
 	EmailDomain             string                         `json:"emailDomain"`
@@ -442,6 +474,7 @@ type SignInInput struct {
 
 type SignInPayload struct {
 	Identity *Identity `json:"identity,omitempty"`
+	Session  *Session  `json:"session,omitempty"`
 }
 
 type SignOutPayload struct {
@@ -820,6 +853,63 @@ func (e ProvisioningSource) MarshalJSON() ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
+type ReauthenticationReason string
+
+const (
+	ReauthenticationReasonSessionExpired    ReauthenticationReason = "SESSION_EXPIRED"
+	ReauthenticationReasonSensitiveAction   ReauthenticationReason = "SENSITIVE_ACTION"
+	ReauthenticationReasonPolicyRequirement ReauthenticationReason = "POLICY_REQUIREMENT"
+)
+
+var AllReauthenticationReason = []ReauthenticationReason{
+	ReauthenticationReasonSessionExpired,
+	ReauthenticationReasonSensitiveAction,
+	ReauthenticationReasonPolicyRequirement,
+}
+
+func (e ReauthenticationReason) IsValid() bool {
+	switch e {
+	case ReauthenticationReasonSessionExpired, ReauthenticationReasonSensitiveAction, ReauthenticationReasonPolicyRequirement:
+		return true
+	}
+	return false
+}
+
+func (e ReauthenticationReason) String() string {
+	return string(e)
+}
+
+func (e *ReauthenticationReason) UnmarshalGQL(v any) error {
+	str, ok := v.(string)
+	if !ok {
+		return fmt.Errorf("enums must be strings")
+	}
+
+	*e = ReauthenticationReason(str)
+	if !e.IsValid() {
+		return fmt.Errorf("%s is not a valid ReauthenticationReason", str)
+	}
+	return nil
+}
+
+func (e ReauthenticationReason) MarshalGQL(w io.Writer) {
+	fmt.Fprint(w, strconv.Quote(e.String()))
+}
+
+func (e *ReauthenticationReason) UnmarshalJSON(b []byte) error {
+	s, err := strconv.Unquote(string(b))
+	if err != nil {
+		return err
+	}
+	return e.UnmarshalGQL(s)
+}
+
+func (e ReauthenticationReason) MarshalJSON() ([]byte, error) {
+	var buf bytes.Buffer
+	e.MarshalGQL(&buf)
+	return buf.Bytes(), nil
+}
+
 type SessionRequirement string
 
 const (

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

@@ -215,13 +215,8 @@ func (r *mutationResolver) SignIn(ctx context.Context, input types.SignInInput)
 	)
 
 	return &types.SignInPayload{
-		Identity: &types.Identity{
-			ID:            user.ID,
-			Email:         user.EmailAddress,
-			EmailVerified: user.EmailAddressVerified,
-			CreatedAt:     user.CreatedAt,
-			UpdatedAt:     user.UpdatedAt,
-		},
+		Identity: types.NewIdentity(user),
+		Session:  types.NewSession(session),
 	}, nil
 }
 
@@ -465,6 +460,50 @@ func (r *mutationResolver) ChangeEmail(ctx context.Context, input types.ChangeEm
 	}, nil
 }
 
+// AssumeOrganizationSession is the resolver for the assumeOrganizationSession field.
+func (r *mutationResolver) AssumeOrganizationSession(ctx context.Context, input types.AssumeOrganizationSessionInput) (*types.AssumeOrganizationSessionPayload, error) {
+	rootSession := SessionFromContext(ctx)
+
+	childSession, membership, err := r.iam.SessionService.AssumeOrganizationSession(ctx, rootSession.ID, input.OrganizationID)
+	if err != nil {
+		var (
+			errMembershipNotFound         *iam.ErrMembershipNotFound
+			errPasswordRequired           *iam.ErrPasswordRequired
+			errSAMLAuthenticationRequired *iam.ErrSAMLAuthenticationRequired
+		)
+
+		switch {
+		case errors.As(err, &errMembershipNotFound):
+			return nil, gqlutils.NotFound(err)
+
+		case errors.As(err, &errPasswordRequired):
+			return &types.AssumeOrganizationSessionPayload{
+				Result: types.PasswordRequired{
+					Reason: types.ReauthenticationReason(errPasswordRequired.Reason),
+				},
+			}, nil
+
+		case errors.As(err, &errSAMLAuthenticationRequired):
+			return &types.AssumeOrganizationSessionPayload{
+				Result: types.SAMLAuthenticationRequired{
+					Reason:      types.ReauthenticationReason(errSAMLAuthenticationRequired.Reason),
+					RedirectURL: errSAMLAuthenticationRequired.RedirectURL,
+				},
+			}, nil
+
+		default:
+			panic(fmt.Errorf("cannot assume organization session: %w", err))
+		}
+	}
+
+	return &types.AssumeOrganizationSessionPayload{
+		Result: types.OrganizationSessionCreated{
+			Session:    types.NewSession(childSession),
+			Membership: types.NewMembership(membership),
+		},
+	}, nil
+}
+
 // UpdateIdentityProfile is the resolver for the updateIdentityProfile field.
 func (r *mutationResolver) UpdateIdentityProfile(ctx context.Context, input types.UpdateIdentityProfileInput) (*types.UpdateIdentityProfilePayload, error) {
 	panic(fmt.Errorf("not implemented: UpdateIdentityProfile - updateIdentityProfile"))
@@ -1061,3 +1100,15 @@ type personalAPIKeyConnectionResolver struct{ *Resolver }
 type queryResolver struct{ *Resolver }
 type sAMLConfigurationConnectionResolver struct{ *Resolver }
 type sessionConnectionResolver struct{ *Resolver }
+
+// !!! WARNING !!!
+// The code below was going to be deleted when updating resolvers. It has been copied here so you have
+// one last chance to move it out of harms way if you want. There are two reasons this happens:
+//  - When renaming or deleting a resolver the old code will be put in here. You can safely delete
+//    it when you're done.
+//  - You have helper methods in this file. Move them out to keep these resolver files clean.
+/*
+	func (r *mutationResolver) SignInWithSession(ctx context.Context, input types.SignInWithSessionInput) (*types.SignInWithSessionPayload, error) {
+	panic(fmt.Errorf("not implemented: SignInWithSession - signInWithSession"))
+}
+*/