Repository: supabase/gotrue
Pull request #837 (2022-11-29 17:57:48)

feat: add `provider` claim to `amr` when the method is `sso/saml`
When the authentication method reference ends with (i.e. the initial authentication method is) `sso/saml`, an additional sub-claim `provider` is added that contains the UUID of the SSO provider that originated the authentication. This is extracted from the single row under the `identities` table.


Patch:
diff --git a/api/token.go b/api/token.go
index 465f1144..ebbfbbd7 100644
--- a/api/token.go
+++ b/api/token.go
@@ -571,7 +571,10 @@ func generateAccessToken(tx *storage.Connection, user *models.User, sessionId *u
 		if terr != nil {
 			return "", terr
 		}
-		aal, amr = session.CalculateAALAndAMR()
+		aal, amr, terr = session.CalculateAALAndAMR(tx)
+		if terr != nil {
+			return "", terr
+		}
 	}
 
 	claims := &GoTrueClaims{
@@ -672,7 +675,11 @@ func (a *API) updateMFASessionAndClaims(r *http.Request, tx *storage.Connection,
 		if terr != nil {
 			return terr
 		}
-		aal, _ := session.CalculateAALAndAMR()
+		aal, _, terr := session.CalculateAALAndAMR(tx)
+		if terr != nil {
+			return terr
+		}
+
 		if err := session.UpdateAssociatedFactor(tx, grantParams.FactorID); err != nil {
 			return err
 		}
diff --git a/models/identity.go b/models/identity.go
index fb5f3a3d..cd2d714b 100644
--- a/models/identity.go
+++ b/models/identity.go
@@ -72,10 +72,10 @@ func FindIdentityByIdAndProvider(tx *storage.Connection, providerId, provider st
 	return identity, nil
 }
 
-// FindIdentitiesByUser returns all identities associated to a user
-func FindIdentitiesByUser(tx *storage.Connection, user *User) ([]*Identity, error) {
+// FindIdentitiesByUserID returns all identities associated to a user ID.
+func FindIdentitiesByUserID(tx *storage.Connection, userID uuid.UUID) ([]*Identity, error) {
 	identities := []*Identity{}
-	if err := tx.Q().Where("user_id = ?", user.ID).All(&identities); err != nil {
+	if err := tx.Q().Where("user_id = ?", userID).All(&identities); err != nil {
 		if errors.Cause(err) == sql.ErrNoRows {
 			return identities, nil
 		}
diff --git a/models/sessions.go b/models/sessions.go
index aebf853c..29cbdc3a 100644
--- a/models/sessions.go
+++ b/models/sessions.go
@@ -3,6 +3,7 @@ package models
 import (
 	"database/sql"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/gobuffalo/pop/v5"
@@ -36,6 +37,7 @@ func (aal AuthenticatorAssuranceLevel) String() string {
 type AMREntry struct {
 	Method    string `json:"method"`
 	Timestamp int64  `json:"timestamp"`
+	Provider  string `json:"provider,omitempty"`
 }
 
 type sortAMREntries struct {
@@ -149,7 +151,7 @@ func (s *Session) UpdateAssociatedAAL(tx *storage.Connection, aal string) error
 	return tx.Update(s)
 }
 
-func (s *Session) CalculateAALAndAMR() (aal string, amr []AMREntry) {
+func (s *Session) CalculateAALAndAMR(tx *storage.Connection) (aal string, amr []AMREntry, err error) {
 	amr, aal = []AMREntry{}, AAL1.String()
 	for _, claim := range s.AMRClaims {
 		if *claim.AuthenticationMethod == TOTPSignIn.String() {
@@ -170,7 +172,30 @@ func (s *Session) CalculateAALAndAMR() (aal string, amr []AMREntry) {
 		Array: amr,
 	})
 
-	return aal, amr
+	lastIndex := len(amr) - 1
+
+	if lastIndex > -1 && amr[lastIndex].Method == SSOSAML.String() {
+		// initial AMR claim is from sso/saml, we need to add information
+		// about the provider that was used for the authentication
+		identities, err := FindIdentitiesByUserID(tx, s.UserID)
+		if err != nil {
+			return aal, amr, err
+		}
+
+		if len(identities) == 1 {
+			identity := identities[0]
+
+			if strings.HasPrefix(identity.Provider, "sso:") {
+				amr[lastIndex].Provider = strings.TrimPrefix(identity.Provider, "sso:")
+			}
+		}
+
+		// otherwise we can't identify that this user account has only
+		// one SSO identity, so we are not encoding the provider at
+		// this time
+	}
+
+	return aal, amr, nil
 }
 
 func (s *Session) GetAAL() string {


Test patch:
diff --git a/models/identity_test.go b/models/identity_test.go
index 37dfe96f..b28e1dfb 100644
--- a/models/identity_test.go
+++ b/models/identity_test.go
@@ -53,7 +53,7 @@ func (ts *IdentityTestSuite) TestNewIdentity() {
 
 func (ts *IdentityTestSuite) TestFindUserIdentities() {
 	u := ts.createUserWithIdentity("test@supabase.io")
-	identities, err := FindIdentitiesByUser(ts.db, u)
+	identities, err := FindIdentitiesByUserID(ts.db, u.ID)
 	require.NoError(ts.T(), err)
 
 	require.Len(ts.T(), identities, 1)
diff --git a/models/sessions_test.go b/models/sessions_test.go
index 7e4eb647..b0991514 100644
--- a/models/sessions_test.go
+++ b/models/sessions_test.go
@@ -57,7 +57,8 @@ func (ts *SessionsTestSuite) TestCalculateAALAndAMR() {
 	session, err = FindSessionByID(ts.db, session.ID)
 	require.NoError(ts.T(), err)
 
-	aal, amr := session.CalculateAALAndAMR()
+	aal, amr, err := session.CalculateAALAndAMR(ts.db)
+	require.NoError(ts.T(), err)
 	require.Equal(ts.T(), AAL2.String(), aal)
 	require.Equal(ts.T(), totalDistinctClaims, len(amr))
 
@@ -67,7 +68,8 @@ func (ts *SessionsTestSuite) TestCalculateAALAndAMR() {
 	session, err = FindSessionByID(ts.db, session.ID)
 	require.NoError(ts.T(), err)
 
-	aal, amr = session.CalculateAALAndAMR()
+	aal, amr, err = session.CalculateAALAndAMR(ts.db)
+	require.NoError(ts.T(), err)
 
 	require.Equal(ts.T(), AAL2.String(), aal)
 	require.Equal(ts.T(), totalDistinctClaims, len(amr))