Hotfix/config init validation (flyteorg#250)

* fix config endpoint validation Signed-off-by: Yuvraj <[email protected]>
austin362667 · May 7, 2024 · 0744ff8 · 0744ff8
1 parent 8581491
commit 0744ff8
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 21 deletions.
diff --git a/flytectl/cmd/configuration/configuration.go b/flytectl/cmd/configuration/configuration.go
@@ -5,9 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"net"
 	"os"
-	"regexp"
 	"strings"
 
 	"github.com/flyteorg/flytestdlib/logger"
@@ -20,8 +18,8 @@ import (
 	cmdcore "github.com/flyteorg/flytectl/cmd/core"
 	cmdUtil "github.com/flyteorg/flytectl/pkg/commandutils"
 	"github.com/flyteorg/flytestdlib/config/viper"
+	"github.com/go-ozzo/ozzo-validation/v4/is"
 	"github.com/manifoldco/promptui"
-
 	"github.com/spf13/cobra"
 )
 
@@ -93,9 +91,8 @@ func initFlytectlConfig(ctx context.Context, reader io.Reader) error {
 	templateStr := configutil.GetSandboxTemplate()
 
 	if len(initConfig.DefaultConfig.Host) > 0 {
-		trimHost := trim(initConfig.DefaultConfig.Host)
-		host := strings.Split(trimHost, ":")
-		if !validateEndpointName(host[0]) {
+		trimHost := trimEndpoint(initConfig.DefaultConfig.Host)
+		if !validateEndpointName(trimHost) {
 			return errors.New("Please use a valid endpoint")
 		}
 		templateValues.Host = fmt.Sprintf("dns://%s", trimHost)
@@ -132,22 +129,36 @@ func initFlytectlConfig(ctx context.Context, reader io.Reader) error {
 	return nil
 }
 
-func trim(hostname string) string {
+func trimEndpoint(hostname string) string {
 	for _, prefix := range endpointPrefix {
 		hostname = strings.TrimPrefix(hostname, prefix)
 	}
 	return hostname
 
 }
 
-func validateEndpointName(domain string) bool {
-	RegExp := regexp.MustCompile(`^(([a-zA-Z]{1})|([a-zA-Z]{1}[a-zA-Z]{1})|([a-zA-Z]{1}[0-9]{1})|([0-9]{1}[a-zA-Z]{1})|([a-zA-Z0-9][a-zA-Z0-9-_]{1,61}[a-zA-Z0-9]))\.([a-zA-Z]{2,6}|[a-zA-Z0-9-]{2,30}\.[a-zA-Z
- ]{2,3})$`)
-	if RegExp.MatchString(domain) || domain == "localhost" {
+func validateEndpointName(endPoint string) bool {
+	var validate = false
+	if endPoint == "localhost" {
 		return true
 	}
-	if net.ParseIP(domain) == nil {
+	if err := is.URL.Validate(endPoint); err != nil {
 		return false
 	}
-	return true
+	endPointParts := strings.Split(endPoint, ":")
+	if len(endPointParts) <= 2 && len(endPointParts) > 0 {
+		if err := is.DNSName.Validate(endPointParts[0]); !errors.Is(err, is.ErrDNSName) && err == nil {
+			validate = true
+		}
+		if err := is.IP.Validate(endPointParts[0]); !errors.Is(err, is.ErrIP) && err == nil {
+			validate = true
+		}
+		if len(endPointParts) == 2 {
+			if err := is.Port.Validate(endPointParts[1]); err != nil {
+				return false
+			}
+		}
+	}
+
+	return validate
 }
diff --git a/flytectl/cmd/configuration/configuration_test.go b/flytectl/cmd/configuration/configuration_test.go
@@ -65,15 +65,20 @@ func TestSetupConfigFunc(t *testing.T) {
 }
 
 func TestTrimFunc(t *testing.T) {
-	assert.Equal(t, trim("dns://localhost"), "localhost")
-	assert.Equal(t, trim("http://localhost"), "localhost")
-	assert.Equal(t, trim("https://localhost"), "localhost")
+	assert.Equal(t, trimEndpoint("dns://localhost"), "localhost")
+	assert.Equal(t, trimEndpoint("http://localhost"), "localhost")
+	assert.Equal(t, trimEndpoint("https://localhost"), "localhost")
 }
 
 func TestValidateEndpointName(t *testing.T) {
-	assert.Equal(t, validateEndpointName("127.0.0.1"), true)
-	assert.Equal(t, validateEndpointName("127.0.0.1.0"), false)
-	assert.Equal(t, validateEndpointName("localhost"), true)
-	assert.Equal(t, validateEndpointName("flyte.org"), true)
-	assert.Equal(t, validateEndpointName("flyte"), false)
+	assert.Equal(t, true, validateEndpointName("8093405779.ap-northeast-2.elb.amazonaws.com:81"))
+	assert.Equal(t, true, validateEndpointName("8093405779.ap-northeast-2.elb.amazonaws.com"))
+	assert.Equal(t, false, validateEndpointName("8093405779.ap-northeast-2.elb.amazonaws.com:81/console"))
+	assert.Equal(t, true, validateEndpointName("localhost"))
+	assert.Equal(t, true, validateEndpointName("127.0.0.1"))
+	assert.Equal(t, true, validateEndpointName("127.0.0.1:30086"))
+	assert.Equal(t, true, validateEndpointName("112.11.1.1"))
+	assert.Equal(t, true, validateEndpointName("112.11.1.1:8080"))
+	assert.Equal(t, false, validateEndpointName("112.11.1.1:8080/console"))
+	assert.Equal(t, false, validateEndpointName("flyte"))
 }
diff --git a/flytectl/go.mod b/flytectl/go.mod
@@ -14,6 +14,7 @@ require (
 	github.com/flyteorg/flyteidl v0.21.5
 	github.com/flyteorg/flytestdlib v0.4.0
 	github.com/ghodss/yaml v1.0.0
+	github.com/go-ozzo/ozzo-validation/v4 v4.3.0
 	github.com/golang/protobuf v1.5.0
 	github.com/google/go-github v17.0.0+incompatible
 	github.com/google/go-github/v37 v37.0.0

diff --git a/flytectl/go.sum b/flytectl/go.sum
@@ -123,6 +123,8 @@ github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmV
 github.com/armon/go-radix v0.0.0-20180808171621-7fddfc383310/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
 github.com/aryann/difflib v0.0.0-20170710044230-e206f873d14a/go.mod h1:DAHtR1m6lCRdSC2Tm3DSWRPvIPr6xNKyeHdqDQSQT+A=
 github.com/asaskevich/govalidator v0.0.0-20190424111038-f61b66f89f4a/go.mod h1:lB+ZfQJz7igIIfQNfa7Ml4HSf2uFQQRzpGGRXenZAgY=
+github.com/asaskevich/govalidator v0.0.0-20200108200545-475eaeb16496 h1:zV3ejI06GQ59hwDQAvmK1qxOQGB3WuVTRoY0okPTAv0=
+github.com/asaskevich/govalidator v0.0.0-20200108200545-475eaeb16496/go.mod h1:oGkLhpf+kjZl6xBf758TQhh5XrAeiJv/7FRz/2spLIg=
 github.com/avast/retry-go v3.0.0+incompatible h1:4SOWQ7Qs+oroOTQOYnAHqelpCO0biHSxpiH9JdtuBj0=
 github.com/avast/retry-go v3.0.0+incompatible/go.mod h1:XtSnn+n/sHqQIpZ10K1qAevBhOOCWBLXXy3hyiqqBrY=
 github.com/awalterschulze/gographviz v2.0.3+incompatible h1:9sVEXJBJLwGX7EQVhLm2elIKCm7P2YHFC8v6096G09E=
@@ -391,6 +393,8 @@ github.com/go-openapi/jsonreference v0.19.3/go.mod h1:rjx6GuL8TTa9VaixXglHmQmIL9
 github.com/go-openapi/spec v0.19.3/go.mod h1:FpwSN1ksY1eteniUU7X0N/BgJ7a4WvBFVA8Lj9mJglo=
 github.com/go-openapi/swag v0.19.2/go.mod h1:POnQmlKehdgb5mhVOsnJFsivZCEZ/vjK9gh66Z9tfKk=
 github.com/go-openapi/swag v0.19.5/go.mod h1:POnQmlKehdgb5mhVOsnJFsivZCEZ/vjK9gh66Z9tfKk=
+github.com/go-ozzo/ozzo-validation/v4 v4.3.0 h1:byhDUpfEwjsVQb1vBunvIjh2BHQ9ead57VkAEY4V+Es=
+github.com/go-ozzo/ozzo-validation/v4 v4.3.0/go.mod h1:2NKgrcHl3z6cJs+3Oo940FPRiTzuqKbvfrL2RxCj6Ew=
 github.com/go-sql-driver/mysql v1.4.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-test/deep v1.0.7 h1:/VSMRlnY/JSyqxQUzQLKVMAskpY/NZKFA5j2P+0pP2M=