Rewrite identity and access management

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

pkg/iam/saml/attributes.go

@@ -0,0 +1,168 @@
+// 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 saml
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/crewjam/saml"
+	"go.probo.inc/probo/pkg/coredata"
+	"go.probo.inc/probo/pkg/mail"
+)
+
+func extractUserAttributes(assertion *saml.Assertion, config *coredata.SAMLConfiguration) (mail.Addr, string, *coredata.MembershipRole, error) {
+	var (
+		email    mail.Addr
+		fullname string
+		role     *coredata.MembershipRole
+	)
+
+	if len(assertion.AttributeStatements) == 0 {
+		if assertion.Subject != nil && assertion.Subject.NameID != nil {
+			email, err := mail.ParseAddr(assertion.Subject.NameID.Value)
+			if err != nil {
+				return mail.Nil, "", nil, fmt.Errorf("cannot parse email: %w", err)
+			}
+
+			fullname = email.String()
+			role = nil
+			return email, fullname, role, nil
+		}
+
+		return mail.Nil, "", nil, fmt.Errorf("no attribute statement and no NameID in assertion")
+	}
+
+	emailString, err := extractAttributeValue(assertion, config.AttributeEmail)
+	if err != nil {
+		if assertion.Subject != nil && assertion.Subject.NameID != nil {
+			emailString = assertion.Subject.NameID.Value
+		} else {
+			return mail.Nil, "", nil, fmt.Errorf("cannot extract email: %w", err)
+		}
+	}
+
+	email, err = mail.ParseAddr(emailString)
+	if err != nil {
+		return mail.Nil, "", nil, fmt.Errorf("cannot parse email: %w", err)
+	}
+
+	firstname, err := extractAttributeValue(assertion, config.AttributeFirstname)
+	if err != nil {
+		firstname = ""
+	}
+
+	lastname, err := extractAttributeValue(assertion, config.AttributeLastname)
+	if err != nil {
+		lastname = ""
+	}
+
+	if firstname != "" && lastname != "" {
+		fullname = strings.TrimSpace(firstname + " " + lastname)
+	} else if firstname != "" {
+		fullname = firstname
+	} else if lastname != "" {
+		fullname = lastname
+	} else {
+		fullname = email.String()
+	}
+
+	roleString, err := extractAttributeValue(assertion, config.AttributeRole)
+	if err != nil {
+		role = nil
+	}
+
+	if roleString != "" {
+		role = mapSAMLRoleToSystemRole(roleString)
+	}
+
+	return email, fullname, role, nil
+}
+
+func extractAttributeValue(assertion *saml.Assertion, attributeName string) (string, error) {
+	if len(assertion.AttributeStatements) == 0 {
+		return "", fmt.Errorf("no attribute statement in assertion")
+	}
+
+	for _, statement := range assertion.AttributeStatements {
+		for _, attr := range statement.Attributes {
+			if attr.Name == attributeName {
+				if len(attr.Values) == 0 {
+					return "", fmt.Errorf("attribute %q has no values", attributeName)
+				}
+
+				return attr.Values[0].Value, nil
+			}
+		}
+	}
+
+	return "", fmt.Errorf("attribute %q not found in assertion", attributeName)
+}
+
+func extractEmailFromAssertion(assertion *saml.Assertion) (string, error) {
+	commonEmailAttributes := []string{
+		"email",
+		"Email",
+		"emailAddress",
+		"mail",
+		"http://schemas.xmlsoap.org/ws/2005/05/identity/claims/emailaddress",
+		"http://schemas.xmlsoap.org/claims/EmailAddress",
+	}
+
+	for _, attrName := range commonEmailAttributes {
+		email, err := extractAttributeValue(assertion, attrName)
+		if err == nil && email != "" {
+			return email, nil
+		}
+	}
+
+	if assertion.Subject != nil && assertion.Subject.NameID != nil && assertion.Subject.NameID.Value != "" {
+		return assertion.Subject.NameID.Value, nil
+	}
+
+	return "", fmt.Errorf("could not extract email from assertion")
+}
+
+func extractEmailDomain(email string) (string, error) {
+	parts := strings.Split(email, "@")
+	if len(parts) != 2 {
+		return "", fmt.Errorf("invalid email address: %s", email)
+	}
+
+	domain := strings.ToLower(strings.TrimSpace(parts[1]))
+	if domain == "" {
+		return "", fmt.Errorf("empty domain in email address: %s", email)
+	}
+
+	return domain, nil
+}
+
+func mapSAMLRoleToSystemRole(samlRole string) *coredata.MembershipRole {
+	if samlRole != "" && isValidRole(samlRole) {
+		role := coredata.MembershipRole(samlRole)
+		return &role
+	}
+
+	return nil
+}
+
+func isValidRole(role string) bool {
+	switch role {
+	case "OWNER", "ADMIN", "EMPLOYEE", "VIEWER":
+		return true
+	default:
+		return false
+	}
+}

pkg/iam/saml/errors.go

@@ -0,0 +1,90 @@
+// 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 saml
+
+import (
+	"fmt"
+
+	"go.probo.inc/probo/pkg/gid"
+	"go.probo.inc/probo/pkg/mail"
+)
+
+type ErrSAMLConfigurationNotFound struct{ ConfigID gid.GID }
+
+func NewSAMLConfigurationNotFoundError(configID gid.GID) error {
+	return &ErrSAMLConfigurationNotFound{ConfigID: configID}
+}
+
+func (e ErrSAMLConfigurationNotFound) Error() string {
+	return fmt.Sprintf("SAML configuration %q not found", e.ConfigID)
+}
+
+type ErrSAMLDisabled struct{}
+
+func NewSAMLDisabledError() error {
+	return &ErrSAMLDisabled{}
+}
+
+func (e ErrSAMLDisabled) Error() string {
+	return "SAML is disabled for this organization"
+}
+
+type ErrInvalidAssertion struct {
+	AssertionID string
+	Err         error
+}
+
+func NewInvalidAssertionError(assertionID string, err error) error {
+	return &ErrInvalidAssertion{AssertionID: assertionID, Err: err}
+}
+
+func (e ErrInvalidAssertion) Error() string {
+	return fmt.Sprintf("invalid assertion %q: %v", e.AssertionID, e.Err)
+}
+
+type ErrReplayAttackDetected struct {
+	AssertionID string
+}
+
+func NewReplayAttackDetectedError(assertionID string) error {
+	return &ErrReplayAttackDetected{AssertionID: assertionID}
+}
+
+func (e ErrReplayAttackDetected) Error() string {
+	return fmt.Sprintf("replay attack detected for assertion %q", e.AssertionID)
+}
+
+type ErrEmailDomainMismatch struct {
+	Email          mail.Addr
+	ExpectedDomain string
+}
+
+func NewEmailDomainMismatchError(email mail.Addr, expectedDomain string) error {
+	return &ErrEmailDomainMismatch{Email: email, ExpectedDomain: expectedDomain}
+}
+
+func (e ErrEmailDomainMismatch) Error() string {
+	return fmt.Sprintf("email domain mismatch: assertion contains email %q but SAML config is for domain %q", e.Email, e.ExpectedDomain)
+}
+
+type ErrSAMLAutoSignupDisabled struct{ ConfigID gid.GID }
+
+func NewSAMLAutoSignupDisabledError(configID gid.GID) error {
+	return &ErrSAMLAutoSignupDisabled{ConfigID: configID}
+}
+
+func (e ErrSAMLAutoSignupDisabled) Error() string {
+	return fmt.Sprintf("SAML auto-signup is disabled for configuration %q", e.ConfigID)
+}

pkg/iam/saml/gc.go

@@ -0,0 +1,97 @@
+// 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 saml
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"go.gearno.de/kit/log"
+	"go.gearno.de/kit/pg"
+	"go.probo.inc/probo/pkg/coredata"
+)
+
+const (
+	DefaultGarbageCollectionInterval = 1 * time.Hour
+)
+
+type (
+	GarbageCollector struct {
+		pg       *pg.Client
+		interval time.Duration
+		logger   *log.Logger
+	}
+)
+
+func NewGarbageCollector(
+	pg *pg.Client,
+	interval time.Duration,
+	logger *log.Logger,
+) *GarbageCollector {
+	return &GarbageCollector{
+		pg:       pg,
+		interval: interval,
+		logger:   logger.Named("saml.garbage_collector").With(log.Duration("interval", interval)),
+	}
+}
+
+func (gc *GarbageCollector) Run(ctx context.Context) error {
+	gc.logger.InfoCtx(ctx, "saml garbage collector starting")
+
+	if err := gc.cleanup(ctx); err != nil {
+		gc.logger.ErrorCtx(ctx, "cannot run initial cleanup", log.Error(err))
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			gc.logger.InfoCtx(ctx, "saml garbage collector shutting down")
+			return ctx.Err()
+		case <-time.After(gc.interval):
+			if err := gc.cleanup(ctx); err != nil {
+				gc.logger.ErrorCtx(ctx, "cannot run periodic cleanup", log.Error(err))
+			}
+		}
+	}
+}
+
+func (gc *GarbageCollector) cleanup(ctx context.Context) error {
+	now := time.Now()
+
+	return gc.pg.WithTx(
+		ctx,
+		func(tx pg.Conn) error {
+			assertionsDeleted, err := coredata.DeleteExpiredSAMLAssertions(ctx, tx, now)
+			if err != nil {
+				return fmt.Errorf("cannot delete expired saml assertions: %w", err)
+			}
+
+			requestsDeleted, err := coredata.DeleteExpiredSAMLRequests(ctx, tx, now)
+			if err != nil {
+				return fmt.Errorf("cannot delete expired saml requests: %w", err)
+			}
+
+			gc.logger.InfoCtx(
+				ctx,
+				"saml garbage collector cleaned up expired assertions and requests",
+				log.Int64("assertions", assertionsDeleted),
+				log.Int64("requests", requestsDeleted),
+			)
+
+			return nil
+		},
+	)
+}

pkg/iam/saml/saml_metadata.go

@@ -0,0 +1,87 @@
+// 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 saml
+
+import (
+	"crypto/x509"
+	"encoding/base64"
+	"encoding/xml"
+	"fmt"
+
+	"github.com/crewjam/saml"
+)
+
+func ParseIdpMetadata(metadataXML []byte) (string, string, *x509.Certificate, error) {
+	var entityDescriptor saml.EntityDescriptor
+	err := xml.Unmarshal(metadataXML, &entityDescriptor)
+	if err != nil {
+		return "", "", nil, fmt.Errorf("cannot parse metadata XML: %w", err)
+	}
+
+	if len(entityDescriptor.IDPSSODescriptors) == 0 {
+		return "", "", nil, fmt.Errorf("no IDPSSODescriptor found in metadata")
+	}
+
+	idpDescriptor := entityDescriptor.IDPSSODescriptors[0]
+
+	ssoURL, err := getSsoURLFromMetadata(idpDescriptor)
+	if err != nil {
+		return "", "", nil, fmt.Errorf("cannot get SSO URL from metadata: %w", err)
+	}
+
+	cert, err := getCertificateFromMetadata(idpDescriptor)
+	if err != nil {
+		return "", "", nil, fmt.Errorf("cannot get certificate from metadata: %w", err)
+	}
+
+	return entityDescriptor.EntityID, ssoURL, cert, nil
+}
+
+func getSsoURLFromMetadata(idpDescriptor saml.IDPSSODescriptor) (string, error) {
+	for _, sso := range idpDescriptor.SingleSignOnServices {
+		if sso.Binding == saml.HTTPPostBinding || sso.Binding == saml.HTTPRedirectBinding {
+			return sso.Location, nil
+		}
+	}
+
+	if len(idpDescriptor.SingleSignOnServices) > 0 {
+		return idpDescriptor.SingleSignOnServices[0].Location, nil
+	}
+
+	return "", fmt.Errorf("no SingleSignOnService found in metadata")
+}
+
+func getCertificateFromMetadata(idpDescriptor saml.IDPSSODescriptor) (*x509.Certificate, error) {
+	for _, keyDescriptor := range idpDescriptor.KeyDescriptors {
+		if keyDescriptor.Use == "signing" || keyDescriptor.Use == "" {
+			if len(keyDescriptor.KeyInfo.X509Data.X509Certificates) > 0 {
+				certData := keyDescriptor.KeyInfo.X509Data.X509Certificates[0].Data
+				certDER, err := base64.StdEncoding.DecodeString(certData)
+				if err != nil {
+					return nil, fmt.Errorf("cannot decode certificate: %w", err)
+				}
+
+				cert, err := x509.ParseCertificate(certDER)
+				if err != nil {
+					return nil, fmt.Errorf("cannot parse certificate: %w", err)
+				}
+
+				return cert, nil
+			}
+		}
+	}
+
+	return nil, fmt.Errorf("no signing certificate found in metadata")
+}

pkg/iam/saml/service.go

@@ -0,0 +1,402 @@
+// 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 saml
+
+import (
+	"context"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/base64"
+	"encoding/xml"
+	"errors"
+	"fmt"
+	"net/url"
+	"strings"
+	"time"
+
+	"github.com/crewjam/saml"
+	"go.gearno.de/kit/log"
+	"go.gearno.de/kit/pg"
+	"go.probo.inc/probo/pkg/baseurl"
+	"go.probo.inc/probo/pkg/coredata"
+	"go.probo.inc/probo/pkg/crypto/cipher"
+	"go.probo.inc/probo/pkg/gid"
+)
+
+type (
+	Service struct {
+		pg            *pg.Client
+		encryptionKey cipher.EncryptionKey
+		baseURL       string
+		certificate   *x509.Certificate
+		privateKey    *rsa.PrivateKey
+		logger        *log.Logger
+	}
+
+	UserInfo struct {
+		Email          string
+		FullName       string
+		Role           *coredata.MembershipRole
+		SAMLSubject    string
+		OrganizationID gid.GID
+		SAMLConfigID   gid.GID
+	}
+)
+
+func NewService(
+	pg *pg.Client,
+	encryptionKey cipher.EncryptionKey,
+	baseURL string,
+	certificate *x509.Certificate,
+	privateKey *rsa.PrivateKey,
+	logger *log.Logger,
+) (*Service, error) {
+	return &Service{
+		pg:            pg,
+		encryptionKey: encryptionKey,
+		baseURL:       baseURL,
+		certificate:   certificate,
+		privateKey:    privateKey,
+		logger:        logger,
+	}, nil
+}
+
+func (s *Service) Run(ctx context.Context) error {
+	gc := NewGarbageCollector(s.pg, DefaultGarbageCollectionInterval, s.logger)
+
+	gcCtx, stopGC := context.WithCancel(ctx)
+	defer stopGC()
+
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- gc.Run(gcCtx)
+	}()
+
+	select {
+	case <-ctx.Done():
+		stopGC()
+		<-errCh
+		return ctx.Err()
+	case err := <-errCh:
+		if err != nil {
+			s.logger.ErrorCtx(ctx, "saml garbage collector failed", log.Error(err))
+			return err
+		}
+
+		return nil
+	}
+}
+
+func (s *Service) GenerateSpMetadata() ([]byte, error) {
+	sp := s.baseServiceProvider()
+	return xml.MarshalIndent(sp.Metadata(), "", "  ")
+}
+
+func (s *Service) InitiateLogin(
+	ctx context.Context,
+	configID gid.GID,
+) (*url.URL, error) {
+	var (
+		now           = time.Now()
+		requestExpiry = now.Add(10 * time.Minute)
+		redirect      *url.URL
+	)
+
+	err := s.pg.WithTx(
+		ctx,
+		func(tx pg.Conn) error {
+			config := &coredata.SAMLConfiguration{}
+			err := config.LoadByID(ctx, tx, coredata.NewNoScope(), configID)
+			if err != nil {
+				if err == coredata.ErrResourceNotFound {
+					return NewSAMLConfigurationNotFoundError(configID)
+				}
+
+				return fmt.Errorf("cannot load SAML configuration: %w", err)
+			}
+
+			if config.EnforcementPolicy == coredata.SAMLEnforcementPolicyOff {
+				return NewSAMLDisabledError()
+			}
+
+			sp, err := s.serviceProvider(ctx, config)
+			if err != nil {
+				return fmt.Errorf("cannot build service provider: %w", err)
+			}
+
+			req, err := sp.MakeAuthenticationRequest(config.IdPSsoURL, saml.HTTPRedirectBinding, saml.HTTPPostBinding)
+			if err != nil {
+				return fmt.Errorf("cannot create authentication request: %w", err)
+			}
+
+			samlRequest := coredata.SAMLRequest{
+				ID:             req.ID,
+				OrganizationID: config.OrganizationID,
+				CreatedAt:      now,
+				ExpiresAt:      requestExpiry,
+			}
+
+			if err := samlRequest.Insert(ctx, tx); err != nil {
+				return fmt.Errorf("cannot insert SAML request: %w", err)
+			}
+
+			redirect, err = req.Redirect(config.ID.String(), sp)
+			if err != nil {
+				return fmt.Errorf("cannot generate redirect URL: %w", err)
+			}
+
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return redirect, nil
+}
+
+func (s *Service) HandleAssertion(
+	ctx context.Context,
+	samlResponse string,
+	configID gid.GID,
+) (*coredata.User, *coredata.Membership, error) {
+	var (
+		now        = time.Now()
+		user       = &coredata.User{}
+		membership = &coredata.Membership{}
+	)
+
+	err := s.pg.WithTx(
+		ctx,
+		func(tx pg.Conn) error {
+			config := &coredata.SAMLConfiguration{}
+
+			err := config.LoadByID(ctx, tx, coredata.NewNoScope(), configID)
+			if err != nil {
+				if err == coredata.ErrResourceNotFound {
+					return NewSAMLConfigurationNotFoundError(configID)
+				}
+
+				return fmt.Errorf("cannot load SAML configuration: %w", err)
+			}
+
+			if config.EnforcementPolicy == coredata.SAMLEnforcementPolicyOff {
+				return NewSAMLDisabledError()
+			}
+
+			sp, err := s.serviceProvider(ctx, config)
+			if err != nil {
+				return fmt.Errorf("cannot create service provider: %w", err)
+			}
+
+			possibleRequestIDs, err := coredata.LoadValidRequestIDsForOrganization(ctx, tx, config.OrganizationID, now)
+			if err != nil {
+				return fmt.Errorf("cannot load valid request IDs: %w", err)
+			}
+
+			decodedResponse, err := base64.StdEncoding.DecodeString(samlResponse)
+			if err != nil {
+				return fmt.Errorf("cannot decode SAML response: %w", err)
+			}
+
+			assertion, err := sp.ParseXMLResponse(decodedResponse, possibleRequestIDs, sp.AcsURL)
+			if err != nil {
+				return fmt.Errorf("cannot parse SAML response: %w", err)
+			}
+
+			err = s.validateAssertion(assertion, config, now)
+			if err != nil {
+				return NewInvalidAssertionError(assertion.ID, err)
+			}
+
+			expiresAt := now.Add(24 * time.Hour)
+			if assertion.Conditions.NotOnOrAfter.IsZero() {
+				expiresAt = assertion.Conditions.NotOnOrAfter
+			}
+
+			samlAssertion := coredata.SAMLAssertion{
+				ID:             assertion.ID,
+				OrganizationID: config.OrganizationID,
+				UsedAt:         now,
+				ExpiresAt:      expiresAt,
+			}
+
+			err = samlAssertion.Insert(ctx, tx)
+			if err != nil {
+				if err == coredata.ErrResourceAlreadyExists {
+					return NewReplayAttackDetectedError(samlAssertion.ID)
+				}
+
+				return fmt.Errorf("cannot insert SAML assertion: %w", err)
+			}
+
+			email, fullname, role, err := extractUserAttributes(assertion, config)
+			if err != nil {
+				return fmt.Errorf("cannot extract user attributes: %w", err)
+			}
+
+			if !strings.EqualFold(email.Domain(), config.EmailDomain) {
+				return NewEmailDomainMismatchError(email, config.EmailDomain)
+			}
+
+			err = user.LoadByEmail(ctx, tx, email)
+			if err == coredata.ErrResourceNotFound && !config.AutoSignupEnabled {
+				return NewSAMLAutoSignupDisabledError(config.ID)
+			} else if err == coredata.ErrResourceNotFound && config.AutoSignupEnabled {
+				*user = coredata.User{
+					ID:                   gid.New(gid.NilTenant, coredata.UserEntityType),
+					EmailAddress:         email,
+					HashedPassword:       nil,
+					EmailAddressVerified: true,
+					FullName:             fullname,
+					CreatedAt:            now,
+					UpdatedAt:            now,
+				}
+
+				err := user.Insert(ctx, tx)
+				if err != nil {
+					return fmt.Errorf("cannot insert user: %w", err)
+				}
+			} else if err != nil {
+				return fmt.Errorf("cannot load user: %w", err)
+			} else {
+				user.SAMLSubject = &assertion.Subject.NameID.Value
+				user.FullName = fullname
+				user.EmailAddress = email
+				user.EmailAddressVerified = true
+				user.UpdatedAt = now
+
+				err = user.Update(ctx, tx)
+				if err != nil {
+					return fmt.Errorf("cannot update user: %w", err)
+				}
+			}
+
+			err = membership.LoadByUserAndOrg(ctx, tx, coredata.NewNoScope(), user.ID, config.OrganizationID)
+			if err != nil && err != coredata.ErrResourceNotFound {
+				return fmt.Errorf("cannot load membership: %w", err)
+			}
+
+			isMember := membership.ID != gid.Nil
+			if !isMember {
+				membership = &coredata.Membership{
+					ID:             gid.New(config.ID.TenantID(), coredata.MembershipEntityType),
+					UserID:         user.ID,
+					OrganizationID: config.OrganizationID,
+					Role:           coredata.MembershipRoleViewer,
+					CreatedAt:      now,
+					UpdatedAt:      now,
+				}
+
+				err = membership.Insert(ctx, tx, coredata.NewNoScope())
+				if err != nil {
+					return fmt.Errorf("cannot insert membership: %w", err)
+				}
+			}
+
+			if role != nil {
+				membership.Role = *role
+				membership.UpdatedAt = now
+
+				err = membership.Update(ctx, tx, coredata.NewNoScope())
+				if err != nil {
+					return fmt.Errorf("cannot update membership: %w", err)
+				}
+			}
+
+			return nil
+		},
+	)
+
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return user, membership, nil
+}
+
+func (s *Service) validateAssertion(assertion *saml.Assertion, config *coredata.SAMLConfiguration, now time.Time) error {
+	const clockSkewTolerance = 5 * time.Minute
+
+	if assertion.ID == "" {
+		return errors.New("assertion ID is required")
+	}
+
+	if assertion.Subject == nil || assertion.Subject.NameID == nil {
+		return fmt.Errorf("subject or NameID missing")
+	}
+
+	if assertion.Issuer.Value != config.IdPEntityID {
+		return fmt.Errorf("assertion issuer %q does not match expected issuer %q",
+			assertion.Issuer.Value, config.IdPEntityID)
+	}
+
+	if assertion.Conditions == nil {
+		return errors.New("assertion conditions are required")
+	}
+
+	if assertion.Conditions.NotOnOrAfter.IsZero() {
+		return errors.New("assertion NotOnOrAfter condition is required")
+	}
+
+	if !assertion.Conditions.NotBefore.IsZero() {
+		if now.Add(clockSkewTolerance).Before(assertion.Conditions.NotBefore) {
+			return fmt.Errorf("assertion not yet valid (NotBefore: %v, now: %v)",
+				assertion.Conditions.NotBefore, now)
+		}
+	}
+
+	if now.Add(-clockSkewTolerance).After(assertion.Conditions.NotOnOrAfter) {
+		return fmt.Errorf("assertion expired (NotOnOrAfter: %v, now: %v)",
+			assertion.Conditions.NotOnOrAfter, now)
+	}
+
+	if len(assertion.Conditions.AudienceRestrictions) == 0 {
+		return errors.New("assertion audience restriction is required")
+	}
+
+	expectedAudience := baseurl.MustParse(s.baseURL).WithPath("/api/connect/v1/saml/2.0/metadata").MustString()
+
+	audienceValid := false
+	for _, restriction := range assertion.Conditions.AudienceRestrictions {
+		if restriction.Audience.Value == expectedAudience {
+			audienceValid = true
+			break
+		}
+	}
+
+	if !audienceValid {
+		return fmt.Errorf("assertion audience %q does not match expected %q",
+			assertion.Conditions.AudienceRestrictions, expectedAudience)
+	}
+
+	return nil
+}
+
+func (s *Service) baseServiceProvider() *saml.ServiceProvider {
+	baseURL := baseurl.MustParse(s.baseURL)
+	metadataURL := baseURL.WithPath("/api/connect/v1/saml/2.0/metadata").URL()
+	acsURL := baseURL.WithPath("/api/connect/v1/saml/2.0/consume").URL()
+
+	return &saml.ServiceProvider{
+		EntityID:          metadataURL.String(),
+		Key:               s.privateKey,
+		Certificate:       s.certificate,
+		MetadataURL:       metadataURL,
+		AcsURL:            acsURL,
+		SloURL:            acsURL,
+		AllowIDPInitiated: true,
+	}
+}

pkg/iam/saml/sp.go

@@ -0,0 +1,68 @@
+// 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 saml
+
+import (
+	"context"
+	"encoding/base64"
+	"fmt"
+
+	"github.com/crewjam/saml"
+	"go.probo.inc/probo/pkg/coredata"
+)
+
+func (s *Service) serviceProvider(
+	ctx context.Context,
+	config *coredata.SAMLConfiguration,
+) (*saml.ServiceProvider, error) {
+	cert, err := config.GetIdPCertificate()
+	if err != nil {
+		return nil, fmt.Errorf("cannot parse IdP certificate: %w", err)
+	}
+
+	sp := s.baseServiceProvider()
+	sp.IDPMetadata = &saml.EntityDescriptor{
+		EntityID: config.IdPEntityID,
+		IDPSSODescriptors: []saml.IDPSSODescriptor{
+			{
+				SSODescriptor: saml.SSODescriptor{
+					RoleDescriptor: saml.RoleDescriptor{
+						ProtocolSupportEnumeration: "urn:oasis:names:tc:SAML:2.0:protocol",
+						KeyDescriptors: []saml.KeyDescriptor{
+							{
+								Use: "signing",
+								KeyInfo: saml.KeyInfo{
+									X509Data: saml.X509Data{
+										X509Certificates: []saml.X509Certificate{
+											{Data: base64.StdEncoding.EncodeToString(cert.Raw)},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+				SingleSignOnServices: []saml.Endpoint{
+					{
+						Binding:  saml.HTTPRedirectBinding,
+						Location: config.IdPSsoURL,
+					},
+				},
+			},
+		},
+	}
+
+	return sp, nil
+}