crewjam · sword-jin · Nov 21, 2023 · Nov 21, 2023
diff --git a/identity_provider.go b/identity_provider.go
@@ -107,8 +107,12 @@ type IdentityProvider struct {
 	AssertionMaker          AssertionMaker
 	SignatureMethod         string
 	ValidDuration           *time.Duration
+	EntityIDConstructor     EntityIDConstructor
 }
 
+// EntityIDConstructor is a function that returns the entityID for customization.
+type EntityIDConstructor func() string
+
 // Metadata returns the metadata structure for this identity provider.
 func (idp *IdentityProvider) Metadata() *EntityDescriptor {
 	certStr := base64.StdEncoding.EncodeToString(idp.Certificate.Raw)
@@ -121,7 +125,7 @@ func (idp *IdentityProvider) Metadata() *EntityDescriptor {
 	}
 
 	ed := &EntityDescriptor{
-		EntityID:      idp.MetadataURL.String(),
+		EntityID:      idp.getEntityID(),
 		ValidUntil:    TimeNow().Add(validDuration),
 		CacheDuration: validDuration,
 		IDPSSODescriptors: []IDPSSODescriptor{
@@ -334,6 +338,18 @@ func (idp *IdentityProvider) ServeIDPInitiated(w http.ResponseWriter, r *http.Re
 	}
 }
 
+// createDefaultEntityIDConstructor creates a function to return entityID from metadataURL.
+func createDefaultEntityIDConstructor(metadataURL url.URL) func() string {
+	return metadataURL.String
+}
+
+func (idp *IdentityProvider) getEntityID() string {
+	if idp.EntityIDConstructor == nil {
+		return createDefaultEntityIDConstructor(idp.MetadataURL)()
+	}
+	return idp.EntityIDConstructor()
+}
+
 // IdpAuthnRequest is used by IdentityProvider to handle a single authentication request.
 type IdpAuthnRequest struct {
 	IDP                     *IdentityProvider
@@ -1019,7 +1035,7 @@ func (req *IdpAuthnRequest) MakeResponse() error {
 		Version:      "2.0",
 		Issuer: &Issuer{
 			Format: "urn:oasis:names:tc:SAML:2.0:nameid-format:entity",
-			Value:  req.IDP.MetadataURL.String(),
+			Value:  req.IDP.getEntityID(),
 		},
 		Status: Status{
 			StatusCode: StatusCode{

diff --git a/identity_provider_test.go b/identity_provider_test.go
@@ -38,11 +38,12 @@ type IdentityProviderTest struct {
 	SPCertificate *x509.Certificate
 	SP            ServiceProvider
 
-	Key             crypto.PrivateKey
-	Signer          crypto.Signer
-	Certificate     *x509.Certificate
-	SessionProvider SessionProvider
-	IDP             IdentityProvider
+	Key              crypto.PrivateKey
+	Signer           crypto.Signer
+	Certificate      *x509.Certificate
+	SessionProvider  SessionProvider
+	IDP              IdentityProvider
+	ExpectedFilename string
 }
 
 func mustParseURL(s string) url.URL {
@@ -98,6 +99,24 @@ var applySigner = idpTestOpts{
 	},
 }
 
+// applyEntityIDConstructor will set the entity ID constructor for the identity provider.
+func applyEntityIDConstructor(c EntityIDConstructor) idpTestOpts {
+	return idpTestOpts{
+		apply: func(_ *testing.T, test *IdentityProviderTest) {
+			test.IDP.EntityIDConstructor = c
+		},
+	}
+}
+
+// applyExpectedFilename will set the expected filename for the identity provider.
+func applyExpectedFilename(filename string) idpTestOpts {
+	return idpTestOpts{
+		apply: func(_ *testing.T, test *IdentityProviderTest) {
+			test.ExpectedFilename = filename
+		},
+	}
+}
+
 func NewIdentityProviderTest(t *testing.T, opts ...idpTestOpts) *IdentityProviderTest {
 	test := IdentityProviderTest{}
 	TimeNow = func() time.Time {
@@ -139,6 +158,7 @@ func NewIdentityProviderTest(t *testing.T, opts ...idpTestOpts) *IdentityProvide
 			},
 		},
 	}
+	test.ExpectedFilename = "TestIDPMakeResponse_response.xml"
 
 	// apply the test options
 	for _, opt := range opts {
@@ -772,7 +792,7 @@ func testMakeResponse(t *testing.T, test *IdentityProviderTest) {
 	doc.Indent(2)
 	responseStr, err := doc.WriteToString()
 	assert.Check(t, err)
-	golden.Assert(t, responseStr, "TestIDPMakeResponse_response.xml")
+	golden.Assert(t, responseStr, test.ExpectedFilename)
 }
 
 func TestIDPWriteResponse(t *testing.T) {
@@ -1130,3 +1150,18 @@ func TestIDPHTTPCanHandleSSORequest(t *testing.T) {
 		assert.Check(t, is.Equal(http.StatusBadRequest, w.Code))
 	}
 }
+
+func TestIdentityProviderCustomEntityID(t *testing.T) {
+	customEntityID := "https://idp.example.com/entity-id"
+	test := NewIdentityProviderTest(
+		t,
+		applyKey,
+		applyEntityIDConstructor(func() string {
+			return customEntityID
+		}),
+		applyExpectedFilename("TestIDPMakeResponse_response_with_custom_entity_id.xml"),
+	)
+
+	assert.Equal(t, customEntityID, test.IDP.Metadata().EntityID)
+	testMakeResponse(t, test)
+}
diff --git a/samlidp/samlidp.go b/samlidp/samlidp.go
@@ -18,12 +18,13 @@ import (
 
 // Options represent the parameters to New() for creating a new IDP server
 type Options struct {
-	URL         url.URL
-	Key         crypto.PrivateKey
-	Signer      crypto.Signer
-	Logger      logger.Interface
-	Certificate *x509.Certificate
-	Store       Store
+	URL                 url.URL
+	Key                 crypto.PrivateKey
+	Signer              crypto.Signer
+	Logger              logger.Interface
+	Certificate         *x509.Certificate
+	Store               Store
+	EntityIDConstructor saml.EntityIDConstructor
 }
 
 // Server represents an IDP server. The server provides the following URLs:
@@ -59,12 +60,13 @@ func New(opts Options) (*Server, error) {
 	s := &Server{
 		serviceProviders: map[string]*saml.EntityDescriptor{},
 		IDP: saml.IdentityProvider{
-			Key:         opts.Key,
-			Signer:      opts.Signer,
-			Logger:      logr,
-			Certificate: opts.Certificate,
-			MetadataURL: metadataURL,
-			SSOURL:      ssoURL,
+			Key:                 opts.Key,
+			Signer:              opts.Signer,
+			Logger:              logr,
+			Certificate:         opts.Certificate,
+			MetadataURL:         metadataURL,
+			SSOURL:              ssoURL,
+			EntityIDConstructor: opts.EntityIDConstructor,
 		},
 		logger: logr,
 		Store:  opts.Store,

diff --git a/samlidp/samlidp_test.go b/samlidp/samlidp_test.go
@@ -66,6 +66,15 @@ func mustParseCertificate(pemStr []byte) *x509.Certificate {
 	return cert
 }
 
+func setupTestVariables() {
+	saml.TimeNow = func() time.Time {
+		rv, _ := time.Parse("Mon Jan 2 15:04:05 MST 2006", "Mon Dec 1 01:57:09 UTC 2015")
+		return rv
+	}
+	jwt.TimeFunc = saml.TimeNow
+	saml.RandReader = &testRandomReader{}
+}
+
 type ServerTest struct {
 	SPKey         *rsa.PrivateKey
 	SPCertificate *x509.Certificate
@@ -79,12 +88,7 @@ type ServerTest struct {
 
 func NewServerTest(t *testing.T) *ServerTest {
 	test := ServerTest{}
-	saml.TimeNow = func() time.Time {
-		rv, _ := time.Parse("Mon Jan 2 15:04:05 MST 2006", "Mon Dec 1 01:57:09 UTC 2015")
-		return rv
-	}
-	jwt.TimeFunc = saml.TimeNow
-	saml.RandReader = &testRandomReader{}
+	setupTestVariables()
 
 	test.SPKey = mustParsePrivateKey(golden.Get(t, "sp_key.pem")).(*rsa.PrivateKey)
 	test.SPCertificate = mustParseCertificate(golden.Get(t, "sp_cert.pem"))
@@ -143,3 +147,28 @@ func TestHTTPCanSSORequest(t *testing.T) {
 		w.Body.String())
 	golden.Assert(t, w.Body.String(), "http_sso_response.html")
 }
+
+func TestHTTPMetadataResponseWithCustomEntityID(t *testing.T) {
+	setupTestVariables()
+
+	server, err := New(Options{
+		Certificate: mustParseCertificate(golden.Get(t, "idp_cert.pem")),
+		Key:         mustParsePrivateKey(golden.Get(t, "idp_key.pem")).(*rsa.PrivateKey),
+		Logger:      logger.DefaultLogger,
+		URL:         url.URL{Scheme: "https", Host: "idp.example.com"},
+		Store:       &MemoryStore{},
+		EntityIDConstructor: func() string {
+			return "https://idp.example.com/idp-id"
+		},
+	})
+	assert.Check(t, err)
+
+	w := httptest.NewRecorder()
+	r, _ := http.NewRequest("GET", "https://idp.example.com/metadata", nil)
+	server.ServeHTTP(w, r)
+	assert.Check(t, is.Equal(http.StatusOK, w.Code))
+	assert.Check(t,
+		strings.HasPrefix(w.Body.String(), "<EntityDescriptor"),
+		w.Body.String())
+	golden.Assert(t, w.Body.String(), "http_metadata_response_with_custom_entity_id.html")
+}
diff --git a/samlidp/testdata/http_metadata_response_with_custom_entity_id.html b/samlidp/testdata/http_metadata_response_with_custom_entity_id.html
@@ -0,0 +1,25 @@
+<EntityDescriptor xmlns="urn:oasis:names:tc:SAML:2.0:metadata" validUntil="2015-12-03T01:57:09Z" cacheDuration="PT48H" entityID="https://idp.example.com/idp-id">
+  <IDPSSODescriptor xmlns="urn:oasis:names:tc:SAML:2.0:metadata" protocolSupportEnumeration="urn:oasis:names:tc:SAML:2.0:protocol">
+    <KeyDescriptor use="signing">
+      <KeyInfo xmlns="http://www.w3.org/2000/09/xmldsig#">
+        <X509Data xmlns="http://www.w3.org/2000/09/xmldsig#">
+          <X509Certificate xmlns="http://www.w3.org/2000/09/xmldsig#">MIIB7zCCAVgCCQDFzbKIp7b3MTANBgkqhkiG9w0BAQUFADA8MQswCQYDVQQGEwJVUzELMAkGA1UECAwCR0ExDDAKBgNVBAoMA2ZvbzESMBAGA1UEAwwJbG9jYWxob3N0MB4XDTEzMTAwMjAwMDg1MVoXDTE0MTAwMjAwMDg1MVowPDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkdBMQwwCgYDVQQKDANmb28xEjAQBgNVBAMMCWxvY2FsaG9zdDCBnzANBgkqhkiG9w0BAQEFAAOBjQAwgYkCgYEA1PMHYmhZj308kWLhZVT4vOulqx/9ibm5B86fPWwUKKQ2i12MYtz07tzukPymisTDhQaqyJ8Kqb/6JjhmeMnEOdTvSPmHO8m1ZVveJU6NoKRn/mP/BD7FW52WhbrUXLSeHVSKfWkNk6S4hk9MV9TswTvyRIKvRsw0X/gfnqkroJcCAwEAATANBgkqhkiG9w0BAQUFAAOBgQCMMlIO+GNcGekevKgkakpMdAqJfs24maGb90DvTLbRZRD7Xvn1MnVBBS9hzlXiFLYOInXACMW5gcoRFfeTQLSouMM8o57h0uKjfTmuoWHLQLi6hnF+cvCsEFiJZ4AbF+DgmO6TarJ8O05t8zvnOwJlNCASPZRH/JmF8tX0hoHuAQ==</X509Certificate>
+        </X509Data>
+      </KeyInfo>
+    </KeyDescriptor>
+    <KeyDescriptor use="encryption">
+      <KeyInfo xmlns="http://www.w3.org/2000/09/xmldsig#">
+        <X509Data xmlns="http://www.w3.org/2000/09/xmldsig#">
+          <X509Certificate xmlns="http://www.w3.org/2000/09/xmldsig#">MIIB7zCCAVgCCQDFzbKIp7b3MTANBgkqhkiG9w0BAQUFADA8MQswCQYDVQQGEwJVUzELMAkGA1UECAwCR0ExDDAKBgNVBAoMA2ZvbzESMBAGA1UEAwwJbG9jYWxob3N0MB4XDTEzMTAwMjAwMDg1MVoXDTE0MTAwMjAwMDg1MVowPDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkdBMQwwCgYDVQQKDANmb28xEjAQBgNVBAMMCWxvY2FsaG9zdDCBnzANBgkqhkiG9w0BAQEFAAOBjQAwgYkCgYEA1PMHYmhZj308kWLhZVT4vOulqx/9ibm5B86fPWwUKKQ2i12MYtz07tzukPymisTDhQaqyJ8Kqb/6JjhmeMnEOdTvSPmHO8m1ZVveJU6NoKRn/mP/BD7FW52WhbrUXLSeHVSKfWkNk6S4hk9MV9TswTvyRIKvRsw0X/gfnqkroJcCAwEAATANBgkqhkiG9w0BAQUFAAOBgQCMMlIO+GNcGekevKgkakpMdAqJfs24maGb90DvTLbRZRD7Xvn1MnVBBS9hzlXiFLYOInXACMW5gcoRFfeTQLSouMM8o57h0uKjfTmuoWHLQLi6hnF+cvCsEFiJZ4AbF+DgmO6TarJ8O05t8zvnOwJlNCASPZRH/JmF8tX0hoHuAQ==</X509Certificate>
+        </X509Data>
+      </KeyInfo>
+      <EncryptionMethod Algorithm="http://www.w3.org/2001/04/xmlenc#aes128-cbc"></EncryptionMethod>
+      <EncryptionMethod Algorithm="http://www.w3.org/2001/04/xmlenc#aes192-cbc"></EncryptionMethod>
+      <EncryptionMethod Algorithm="http://www.w3.org/2001/04/xmlenc#aes256-cbc"></EncryptionMethod>
+      <EncryptionMethod Algorithm="http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p"></EncryptionMethod>
+    </KeyDescriptor>
+    <NameIDFormat>urn:oasis:names:tc:SAML:2.0:nameid-format:transient</NameIDFormat>
+    <SingleSignOnService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect" Location="https://idp.example.com/sso"></SingleSignOnService>
+    <SingleSignOnService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST" Location="https://idp.example.com/sso"></SingleSignOnService>
+  </IDPSSODescriptor>
+</EntityDescriptor>
diff --git a/testdata/TestIDPMakeResponse_response_with_custom_entity_id.xml b/testdata/TestIDPMakeResponse_response_with_custom_entity_id.xml
@@ -0,0 +1,27 @@
+<samlp:Response xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion" xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" xmlns:xs="http://www.w3.org/2001/XMLSchema" ID="id-282a2c2e30323436383a3c3e40424446484a4c4e" InResponseTo="id-00020406080a0c0e10121416181a1c1e" Version="2.0" IssueInstant="2015-12-01T01:57:09Z" Destination="https://sp.example.com/saml2/acs">
+  <saml:Issuer Format="urn:oasis:names:tc:SAML:2.0:nameid-format:entity">https://idp.example.com/entity-id</saml:Issuer>
+  <ds:Signature xmlns:ds="http://www.w3.org/2000/09/xmldsig#">
+    <ds:SignedInfo>
+      <ds:CanonicalizationMethod Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
+      <ds:SignatureMethod Algorithm="http://www.w3.org/2000/09/xmldsig#rsa-sha1"/>
+      <ds:Reference URI="#id-282a2c2e30323436383a3c3e40424446484a4c4e">
+        <ds:Transforms>
+          <ds:Transform Algorithm="http://www.w3.org/2000/09/xmldsig#enveloped-signature"/>
+          <ds:Transform Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
+        </ds:Transforms>
+        <ds:DigestMethod Algorithm="http://www.w3.org/2000/09/xmldsig#sha1"/>
+        <ds:DigestValue>5bBiRThV9gjcTNlKa+y00Gnzkh8=</ds:DigestValue>
+      </ds:Reference>
+    </ds:SignedInfo>
+    <ds:SignatureValue>A9fzgSO00HntRcx32qCEVHoTR8YiisGk6tkeAbhRKzXoIOw3UE4nhoBIYPTYj5G+mMjnB/eEw84kuUSZ9mLV+EIAMQuR6ctJyO6xdxy65l+iC0IBSk65wqCb6C4IRB5OaxN/QC0yTJ8Ps2+s1WRJSLLcmQU6Xatpe25vzk+hQ+4=</ds:SignatureValue>
+    <ds:KeyInfo>
+      <ds:X509Data>
+        <ds:X509Certificate>MIIB7zCCAVgCCQDFzbKIp7b3MTANBgkqhkiG9w0BAQUFADA8MQswCQYDVQQGEwJVUzELMAkGA1UECAwCR0ExDDAKBgNVBAoMA2ZvbzESMBAGA1UEAwwJbG9jYWxob3N0MB4XDTEzMTAwMjAwMDg1MVoXDTE0MTAwMjAwMDg1MVowPDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkdBMQwwCgYDVQQKDANmb28xEjAQBgNVBAMMCWxvY2FsaG9zdDCBnzANBgkqhkiG9w0BAQEFAAOBjQAwgYkCgYEA1PMHYmhZj308kWLhZVT4vOulqx/9ibm5B86fPWwUKKQ2i12MYtz07tzukPymisTDhQaqyJ8Kqb/6JjhmeMnEOdTvSPmHO8m1ZVveJU6NoKRn/mP/BD7FW52WhbrUXLSeHVSKfWkNk6S4hk9MV9TswTvyRIKvRsw0X/gfnqkroJcCAwEAATANBgkqhkiG9w0BAQUFAAOBgQCMMlIO+GNcGekevKgkakpMdAqJfs24maGb90DvTLbRZRD7Xvn1MnVBBS9hzlXiFLYOInXACMW5gcoRFfeTQLSouMM8o57h0uKjfTmuoWHLQLi6hnF+cvCsEFiJZ4AbF+DgmO6TarJ8O05t8zvnOwJlNCASPZRH/JmF8tX0hoHuAQ==</ds:X509Certificate>
+      </ds:X509Data>
+    </ds:KeyInfo>
+  </ds:Signature>
+  <samlp:Status>
+    <samlp:StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>
+  </samlp:Status>
+  <this-is-an-encrypted-assertion/>
+</samlp:Response>