feat: add custom attributes to saml response (#56)

* feat: add custom attributes to saml response * fix: correct unit test with logic changes * fix: correct unit test with logic changes
zitadel · Aug 14, 2023 · afbec8e · afbec8e
1 parent cc9378c
commit afbec8e
Show file tree

Hide file tree

Showing 9 changed files with 229 additions and 19 deletions.
diff --git a/pkg/provider/attributes.go b/pkg/provider/attributes.go
@@ -14,13 +14,20 @@ const (
 	AttributeUserID
 )
 
+type CustomAttribute struct {
+	FriendlyName   string
+	NameFormat     string
+	AttributeValue []string
+}
+
 type Attributes struct {
-	email     string
-	fullName  string
-	givenName string
-	surname   string
-	userID    string
-	username  string
+	email            string
+	fullName         string
+	givenName        string
+	surname          string
+	userID           string
+	username         string
+	customAttributes map[string]*CustomAttribute
 }
 
 var _ models.AttributeSetter = &Attributes{}
@@ -56,6 +63,17 @@ func (a *Attributes) SetUserID(value string) {
 	a.userID = value
 }
 
+func (a *Attributes) SetCustomAttribute(name, friendlyName, nameFormat string, attributeValue []string) {
+	if a.customAttributes == nil {
+		a.customAttributes = make(map[string]*CustomAttribute)
+	}
+	a.customAttributes[name] = &CustomAttribute{
+		FriendlyName:   friendlyName,
+		NameFormat:     nameFormat,
+		AttributeValue: attributeValue,
+	}
+}
+
 func (a *Attributes) GetSAML() []*saml.AttributeType {
 	attrs := make([]*saml.AttributeType, 0)
 	if a.email != "" {
@@ -100,5 +118,13 @@ func (a *Attributes) GetSAML() []*saml.AttributeType {
 			AttributeValue: []string{a.userID},
 		})
 	}
+	for name, attr := range a.customAttributes {
+		attrs = append(attrs, &saml.AttributeType{
+			Name:           name,
+			FriendlyName:   attr.FriendlyName,
+			NameFormat:     attr.NameFormat,
+			AttributeValue: attr.AttributeValue,
+		})
+	}
 	return attrs
 }
diff --git a/pkg/provider/attributes_test.go b/pkg/provider/attributes_test.go
@@ -0,0 +1,183 @@
+package provider
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/zitadel/saml/pkg/provider/xml/saml"
+)
+
+func TestSSO_Attributes(t *testing.T) {
+	type args struct {
+		email            string
+		fullName         string
+		givenName        string
+		surname          string
+		userID           string
+		username         string
+		customAttributes map[string]*CustomAttribute
+	}
+	tests := []struct {
+		name string
+		args args
+		res  []*saml.AttributeType
+	}{
+		{
+			"empty attributes",
+			args{},
+			[]*saml.AttributeType{},
+		},
+		{
+			"full attributes",
+			args{
+				email:            "email",
+				fullName:         "fullname",
+				givenName:        "givenname",
+				surname:          "surname",
+				userID:           "userid",
+				username:         "username",
+				customAttributes: nil,
+			},
+			[]*saml.AttributeType{
+				{
+					Name:           "Email",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"email"},
+				},
+				{
+					Name:           "SurName",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"surname"},
+				},
+				{
+					Name:           "FirstName",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"givenname"},
+				},
+				{
+					Name:           "FullName",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"fullname"},
+				},
+				{
+					Name:           "UserName",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"username"},
+				},
+				{
+					Name:           "UserID",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"userid"},
+				},
+			},
+		},
+		{
+			"full attributes with custom",
+			args{
+				email:     "email",
+				fullName:  "fullname",
+				givenName: "givenname",
+				surname:   "surname",
+				userID:    "userid",
+				username:  "username",
+				customAttributes: map[string]*CustomAttribute{
+					"empty": {
+						FriendlyName:   "fname",
+						NameFormat:     "nameformat",
+						AttributeValue: []string{""},
+					},
+					"key1": {
+						FriendlyName:   "fname1",
+						NameFormat:     "nameformat1",
+						AttributeValue: []string{"first"},
+					},
+					"key2": {
+						FriendlyName:   "fname2",
+						NameFormat:     "nameformat2",
+						AttributeValue: []string{"first", "second"},
+					},
+					"key3": {
+						FriendlyName:   "fname3",
+						NameFormat:     "nameformat3",
+						AttributeValue: []string{"first", "second", "third"},
+					},
+				},
+			},
+			[]*saml.AttributeType{
+				{
+					Name:           "Email",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"email"},
+				},
+				{
+					Name:           "SurName",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"surname"},
+				},
+				{
+					Name:           "FirstName",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"givenname"},
+				},
+				{
+					Name:           "FullName",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"fullname"},
+				},
+				{
+					Name:           "UserName",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"username"},
+				},
+				{
+					Name:           "UserID",
+					NameFormat:     "urn:oasis:names:tc:SAML:2.0:attrname-format:basic",
+					AttributeValue: []string{"userid"},
+				},
+				{
+					Name:           "empty",
+					NameFormat:     "nameformat",
+					FriendlyName:   "fname",
+					AttributeValue: []string{""},
+				},
+				{
+					Name:           "key1",
+					NameFormat:     "nameformat1",
+					FriendlyName:   "fname1",
+					AttributeValue: []string{"first"},
+				},
+				{
+					Name:           "key2",
+					NameFormat:     "nameformat2",
+					FriendlyName:   "fname2",
+					AttributeValue: []string{"first", "second"},
+				},
+				{
+					Name:           "key3",
+					NameFormat:     "nameformat3",
+					FriendlyName:   "fname3",
+					AttributeValue: []string{"first", "second", "third"},
+				},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			attrs := &Attributes{
+				email:            tt.args.email,
+				fullName:         tt.args.fullName,
+				givenName:        tt.args.givenName,
+				surname:          tt.args.surname,
+				userID:           tt.args.userID,
+				username:         tt.args.username,
+				customAttributes: tt.args.customAttributes,
+			}
+			samlResponseAttributes := attrs.GetSAML()
+			for _, item := range tt.res {
+				assert.Contains(t, samlResponseAttributes, item)
+			}
+		})
+	}
+}
diff --git a/pkg/provider/login.go b/pkg/provider/login.go
@@ -57,7 +57,7 @@ func (p *IdentityProvider) callbackHandleFunc(w http.ResponseWriter, r *http.Req
 	response.Audience = entityID
 
 	attrs := &Attributes{}
-	if err := p.storage.SetUserinfoWithUserID(ctx, attrs, authRequest.GetUserID(), []int{}); err != nil {
+	if err := p.storage.SetUserinfoWithUserID(ctx, authRequest.GetApplicationID(), attrs, authRequest.GetUserID(), []int{}); err != nil {
 		logging.Error(err)
 		http.Error(w, fmt.Errorf("failed to get userinfo: %w", err).Error(), http.StatusInternalServerError)
 		return

diff --git a/pkg/provider/login_test.go b/pkg/provider/login_test.go
@@ -297,7 +297,7 @@ func idpStorageWithResponseCertAndApp(
 ) *mock.MockIDPStorage {
 	mockStorage := idpStorageWithResponseCert(t, cert, pKey)
 	mockStorage.EXPECT().GetEntityIDByAppID(gomock.Any(), appID).Return(entityID, spErr).MinTimes(0).MaxTimes(1)
-	mockStorage.EXPECT().SetUserinfoWithUserID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).MinTimes(0).MaxTimes(1)
+	mockStorage.EXPECT().SetUserinfoWithUserID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).MinTimes(0).MaxTimes(1)
 
 	request := mock.NewMockAuthRequestInt(gomock.NewController(t))
 	request.EXPECT().GetAuthRequestID().Return(samlAuthRequestID).MinTimes(0).MaxTimes(1)
@@ -306,7 +306,7 @@ func idpStorageWithResponseCertAndApp(
 	request.EXPECT().GetAccessConsumerServiceURL().Return(acsURL).MinTimes(0).MaxTimes(1)
 	request.EXPECT().GetUserID().Return(userID).MinTimes(0).MaxTimes(1)
 	request.EXPECT().Done().Return(done).MinTimes(0).MaxTimes(1)
-	request.EXPECT().GetApplicationID().Return(appID).MinTimes(0).MaxTimes(1)
+	request.EXPECT().GetApplicationID().Return(appID).MinTimes(0).MaxTimes(2)
 	mockStorage.EXPECT().AuthRequestByID(gomock.Any(), authRequestID).Return(request, nil).MinTimes(0).MaxTimes(1)
 
 	return mockStorage

diff --git a/pkg/provider/metadata.go b/pkg/provider/metadata.go
@@ -71,7 +71,7 @@ func (p *IdentityProviderConfig) getMetadata(
 	}
 
 	attrs := &Attributes{
-		"empty", "empty", "empty", "empty", "empty", "empty",
+		"empty", "empty", "empty", "empty", "empty", "empty", nil,
 	}
 	attrsSaml := attrs.GetSAML()
 	for _, attr := range attrsSaml {

diff --git a/pkg/provider/mock/idpstorage.mock.go b/pkg/provider/mock/idpstorage.mock.go
diff --git a/pkg/provider/mock/storage.mock.go b/pkg/provider/mock/storage.mock.go
diff --git a/pkg/provider/models/models.go b/pkg/provider/models/models.go
@@ -24,4 +24,5 @@ type AttributeSetter interface {
 	SetSurname(string)
 	SetUserID(string)
 	SetUsername(string)
+	SetCustomAttribute(name string, friendlyName string, nameFormat string, attributeValue []string)
 }
diff --git a/pkg/provider/storage.go b/pkg/provider/storage.go
@@ -26,6 +26,6 @@ type AuthStorage interface {
 }
 
 type UserStorage interface {
-	SetUserinfoWithUserID(ctx context.Context, userinfo models.AttributeSetter, userID string, attributes []int) (err error)
+	SetUserinfoWithUserID(ctx context.Context, applicationID string, userinfo models.AttributeSetter, userID string, attributes []int) (err error)
 	SetUserinfoWithLoginName(ctx context.Context, userinfo models.AttributeSetter, loginName string, attributes []int) (err error)
 }