From e3fb7729ddf934840a58cad77287c50e99196bbe Mon Sep 17 00:00:00 2001
From: Qiu Jian <qiujian@yunionyun.com>
Date: Fri, 22 Jan 2021 00:52:00 +0800
Subject: [PATCH] fix(cloudid): saml navigation template support i18n

saml navigation template support i18n
---
 pkg/cloudid/saml/initidp.go        |  4 ++--
 pkg/util/samlutils/demo/service.go |  3 ++-
 pkg/util/samlutils/idp/idp.go      | 24 ++++++++++++++----------
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/pkg/cloudid/saml/initidp.go b/pkg/cloudid/saml/initidp.go
index c9683c91f7..cab8b41516 100644
--- a/pkg/cloudid/saml/initidp.go
+++ b/pkg/cloudid/saml/initidp.go
@@ -28,6 +28,7 @@ import (
 	"yunion.io/x/onecloud/pkg/cloudid/models"
 	"yunion.io/x/onecloud/pkg/cloudid/options"
 	"yunion.io/x/onecloud/pkg/httperrors"
+	"yunion.io/x/onecloud/pkg/i18n"
 	"yunion.io/x/onecloud/pkg/mcclient/auth"
 	"yunion.io/x/onecloud/pkg/util/httputils"
 	"yunion.io/x/onecloud/pkg/util/samlutils"
@@ -111,8 +112,7 @@ func initSAMLIdp(app *appsrv.Application, prefix string) error {
 	}
 
 	idpInst.AddHandlers(app, prefix, auth.Authenticate)
-	idpInst.SetHtmlTemplate("zh-CN", `<!DOCTYPE html><html lang="zh"><head><meta charset="utf-8"><meta http-equiv="Content-Type" content="text/html; charset=utf-8"></head><body><h1>正在跳转到云控制台，请等待。。。</h1>$FORM$</body></html>`)
-	idpInst.SetHtmlTemplate("en", `<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta http-equiv="Content-Type" content="text/html; charset=utf-8"></head><body><h1>Jumping to the console, please wait。。。</h1>$FORM$</body></html>`)
+	idpInst.SetHtmlTemplate(i18n.NewTableEntry().CN(`<!DOCTYPE html><html lang="zh"><head><meta charset="utf-8"><meta http-equiv="Content-Type" content="text/html; charset=utf-8"></head><body><h1>正在跳转到云控制台，请等待。。。</h1>$FORM$</body></html>`).EN(`<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta http-equiv="Content-Type" content="text/html; charset=utf-8"></head><body><h1>Login into to the cloud console, please wait。。。</h1>$FORM$</body></html>`))
 
 	idpInstance = idpInst
 
diff --git a/pkg/util/samlutils/demo/service.go b/pkg/util/samlutils/demo/service.go
index 81e8fcd4ff..a9d7be083b 100644
--- a/pkg/util/samlutils/demo/service.go
+++ b/pkg/util/samlutils/demo/service.go
@@ -29,6 +29,7 @@ import (
 
 	"yunion.io/x/onecloud/pkg/appsrv"
 	"yunion.io/x/onecloud/pkg/httperrors"
+	"yunion.io/x/onecloud/pkg/i18n"
 	"yunion.io/x/onecloud/pkg/util/fileutils2"
 	"yunion.io/x/onecloud/pkg/util/httputils"
 	"yunion.io/x/onecloud/pkg/util/samlutils"
@@ -371,7 +372,7 @@ func prepareServer() error {
 		}
 	}
 	idpInst.AddHandlers(app, "SAML/idp", nil)
-	idpInst.SetHtmlTemplate("zh-CN", `<!DOCTYPE html><html lang="zh_CN"><head><meta charset="utf-8"><meta http-equiv="Content-Type" content="text/html; charset=utf-8"></head><body><h1>正在跳转到云控制台，请等待。。。</h1>$FORM$</body></html>`)
+	idpInst.SetHtmlTemplate(i18n.NewTableEntry().CN(`<!DOCTYPE html><html lang="zh_CN"><head><meta charset="utf-8"><meta http-equiv="Content-Type" content="text/html; charset=utf-8"></head><body><h1>正在跳转到云控制台，请等待。。。</h1>$FORM$</body></html>`))
 
 	app.AddHandler("GET", "SAML/idp", func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
 		idpInitUrl := httputils.JoinPath(options.Entity, "SAML/idp/sso")
diff --git a/pkg/util/samlutils/idp/idp.go b/pkg/util/samlutils/idp/idp.go
index 2a4a7af24d..1d0908b484 100644
--- a/pkg/util/samlutils/idp/idp.go
+++ b/pkg/util/samlutils/idp/idp.go
@@ -28,12 +28,15 @@ import (
 	"yunion.io/x/onecloud/pkg/appctx"
 	"yunion.io/x/onecloud/pkg/appsrv"
 	"yunion.io/x/onecloud/pkg/httperrors"
+	"yunion.io/x/onecloud/pkg/i18n"
 	"yunion.io/x/onecloud/pkg/util/httputils"
 	"yunion.io/x/onecloud/pkg/util/samlutils"
 )
 
 const (
 	IDP_ID_KEY = "<idp_id>"
+
+	langTemplateKey = "lang_template_key"
 )
 
 type OnSpInitiatedLogin func(ctx context.Context, idpId string, sp *SSAMLServiceProvider) (samlutils.SSAMLSpInitiatedLoginData, error)
@@ -54,7 +57,7 @@ type SSAMLIdpInstance struct {
 	onIdpInitiatedLogin OnIdpInitiatedLogin
 	onLogout            OnLogout
 
-	htmlTemplate map[string]string
+	htmlTemplate i18n.Table
 }
 
 func NewIdpInstance(saml *samlutils.SSAMLInstance, spLoginFunc OnSpInitiatedLogin, idpLoginFunc OnIdpInitiatedLogin, logoutFunc OnLogout) *SSAMLIdpInstance {
@@ -63,7 +66,7 @@ func NewIdpInstance(saml *samlutils.SSAMLInstance, spLoginFunc OnSpInitiatedLogi
 		onSpInitiatedLogin:  spLoginFunc,
 		onIdpInitiatedLogin: idpLoginFunc,
 		onLogout:            logoutFunc,
-		htmlTemplate:        map[string]string{},
+		htmlTemplate:        i18n.Table{},
 	}
 }
 
@@ -97,11 +100,13 @@ func (idp *SSAMLIdpInstance) AddHandlers(app *appsrv.Application, prefix string,
 	log.Infof("IDP initated SSO: %s", idp.getIdpInitiatedSSOUrl())
 }
 
-func (idp *SSAMLIdpInstance) SetHtmlTemplate(lang, tmp string) error {
-	if strings.Index(tmp, samlutils.HTML_SAML_FORM_TOKEN) < 0 {
-		return errors.Wrapf(httperrors.ErrInvalidFormat, "no %s found", samlutils.HTML_SAML_FORM_TOKEN)
+func (idp *SSAMLIdpInstance) SetHtmlTemplate(entry i18n.TableEntry) error {
+	for _, tmp := range entry {
+		if strings.Index(tmp, samlutils.HTML_SAML_FORM_TOKEN) < 0 {
+			return errors.Wrapf(httperrors.ErrInvalidFormat, "no %s found", samlutils.HTML_SAML_FORM_TOKEN)
+		}
 	}
-	idp.htmlTemplate[lang] = tmp
+	idp.htmlTemplate.Set(langTemplateKey, entry)
 	return nil
 }
 
@@ -263,10 +268,9 @@ func (idp *SSAMLIdpInstance) samlResponse2Form(ctx context.Context, url string,
 		"RelayState":   state,
 	})
 	template := samlutils.DEFAULT_HTML_TEMPLATE
-	if len(idp.htmlTemplate) > 0 {
-		if _temp, ok := idp.htmlTemplate[appctx.AppContextLang(ctx)]; ok {
-			template = _temp
-		}
+	_temp := idp.htmlTemplate.Lookup(ctx, langTemplateKey)
+	if _temp != langTemplateKey {
+		template = _temp
 	}
 	form = strings.Replace(template, samlutils.HTML_SAML_FORM_TOKEN, form, 1)
 	return form, nil