From 9d8092d3809b4b3297b47f83c6275a28f015e0de Mon Sep 17 00:00:00 2001
From: "Dat. Ba Dao" <datdao@ariga.io>
Date: Mon, 2 Oct 2023 21:15:27 +0700
Subject: [PATCH] api: remove special characters from url (#97)

---
 controllers/atlasschema_controller_test.go | 26 ++++++++++++++++++++++
 controllers/common.go                      | 12 ++++++++++
 controllers/templates/atlas_migration.tmpl |  2 +-
 controllers/templates/atlas_schema.tmpl    |  6 ++---
 4 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/controllers/atlasschema_controller_test.go b/controllers/atlasschema_controller_test.go
index dc0d5139..f512bca9 100644
--- a/controllers/atlasschema_controller_test.go
+++ b/controllers/atlasschema_controller_test.go
@@ -444,6 +444,32 @@ env {
 	require.EqualValues(t, expected, buf.String())
 }
 
+func TestTemplate_Func_RemoveSpecialChars(t *testing.T) {
+	var buf bytes.Buffer
+	tmpl, err := tmpl.New("specialChars").Parse(`
+	{{- removeSpecialChars .Text -}}
+	{{- removeSpecialChars .URL -}}
+`)
+	require.NoError(t, err)
+	var textWithSpecialChars = "a\tb\rc\n"
+	err = tmpl.ExecuteTemplate(&buf, `specialChars`, struct {
+		Text string
+		URL  *url.URL
+	}{
+		Text: textWithSpecialChars,
+		URL:  &url.URL{},
+	})
+	require.NoError(t, err)
+	require.EqualValues(t, "abc", buf.String())
+	// invalid data type
+	err = tmpl.ExecuteTemplate(&buf, `specialChars`, struct {
+		Text int
+	}{
+		Text: 0,
+	})
+	require.ErrorContains(t, err, "unsupported type int")
+}
+
 func conditionReconciling() *dbv1alpha1.AtlasSchema {
 	return &dbv1alpha1.AtlasSchema{
 		ObjectMeta: objmeta(),
diff --git a/controllers/common.go b/controllers/common.go
index 88f8b027..b2e71e83 100644
--- a/controllers/common.go
+++ b/controllers/common.go
@@ -20,6 +20,7 @@ import (
 	"errors"
 	"fmt"
 	"io/fs"
+	"regexp"
 	"strings"
 	"testing/fstest"
 	"text/template"
@@ -58,6 +59,17 @@ var (
 				b.WriteRune(']')
 				return b.String()
 			},
+			"removeSpecialChars": func(s interface{}) (string, error) {
+				r := regexp.MustCompile("[\t\r\n]")
+				switch s := s.(type) {
+				case string:
+					return r.ReplaceAllString(s, ""), nil
+				case fmt.Stringer:
+					return r.ReplaceAllString(s.String(), ""), nil
+				default:
+					return "", fmt.Errorf("unsupported type %T", s)
+				}
+			},
 		}).
 		ParseFS(tmpls, "templates/*.tmpl"),
 	)
diff --git a/controllers/templates/atlas_migration.tmpl b/controllers/templates/atlas_migration.tmpl
index d8661595..155b2089 100644
--- a/controllers/templates/atlas_migration.tmpl
+++ b/controllers/templates/atlas_migration.tmpl
@@ -23,7 +23,7 @@ data "remote_dir" "this" {
 {{- end }}
 env {
   name = atlas.env
-  url = "{{ .URL }}"
+  url = "{{ removeSpecialChars .URL }}"
   migration {
 {{- if .Cloud.HasRemoteDir }}
     dir = data.remote_dir.this.url
diff --git a/controllers/templates/atlas_schema.tmpl b/controllers/templates/atlas_schema.tmpl
index 8c698650..66213a0c 100644
--- a/controllers/templates/atlas_schema.tmpl
+++ b/controllers/templates/atlas_schema.tmpl
@@ -68,9 +68,9 @@ lint {
 {{- end }}
 env {
   name = atlas.env
-  src  = "{{ .Source }}"
-  url  = "{{ .URL }}"
-  dev  = "{{ .DevURL }}"
+  src  = "{{ removeSpecialChars .Source }}"
+  url  = "{{ removeSpecialChars .URL }}"
+  dev  = "{{ removeSpecialChars .DevURL }}"
   schemas = {{ slides .Schemas }}
   exclude = {{ slides .Exclude }}
 }