From 7f1cbf51299baa62314b5d627c7c28d7a2230b38 Mon Sep 17 00:00:00 2001
From: Ilya Radostev <feeltheajf@gmail.com>
Date: Tue, 5 Nov 2024 22:32:28 +0100
Subject: [PATCH] Allow setting custom HTML templates

---
 pkg/provider/identityprovider.go | 33 +++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/pkg/provider/identityprovider.go b/pkg/provider/identityprovider.go
index 3a7c024..8d87057 100644
--- a/pkg/provider/identityprovider.go
+++ b/pkg/provider/identityprovider.go
@@ -41,6 +41,9 @@ type MetadataIDPConfig struct {
 type IdentityProviderConfig struct {
 	MetadataIDPConfig *MetadataIDPConfig
 
+	PostTemplate   *template.Template
+	LogoutTemplate *template.Template
+
 	SignatureAlgorithm  string
 	DigestAlgorithm     string
 	EncryptionAlgorithm string
@@ -79,27 +82,31 @@ type Endpoints struct {
 	attributeEndpoint    Endpoint
 }
 
-func NewIdentityProvider(metadata Endpoint, conf *IdentityProviderConfig, storage IDPStorage) (*IdentityProvider, error) {
-	postTemplate, err := template.New("post").Parse(postTemplate)
-	if err != nil {
-		return nil, err
-	}
-
-	logoutTemplate, err := template.New("logout").Parse(logoutTemplate)
-	if err != nil {
-		return nil, err
-	}
-
+func NewIdentityProvider(metadata Endpoint, conf *IdentityProviderConfig, storage IDPStorage) (_ *IdentityProvider, err error) {
 	idp := &IdentityProvider{
 		storage:          storage,
 		metadataEndpoint: &metadata,
 		conf:             conf,
-		postTemplate:     postTemplate,
-		logoutTemplate:   logoutTemplate,
+		postTemplate:     conf.PostTemplate,
+		logoutTemplate:   conf.LogoutTemplate,
 		endpoints:        endpointConfigToEndpoints(conf.Endpoints),
 		timeFormat:       DefaultTimeFormat,
 	}
 
+	if conf.PostTemplate == nil {
+		idp.postTemplate, err = template.New("post").Parse(postTemplate)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if conf.LogoutTemplate == nil {
+		idp.logoutTemplate, err = template.New("logout").Parse(logoutTemplate)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	if conf.MetadataIDPConfig == nil {
 		conf.MetadataIDPConfig = &MetadataIDPConfig{}
 	}