From 6098f6c0e7bbe1ec75187f11e85c2b5f52d943ca Mon Sep 17 00:00:00 2001
From: Manuel Alejandro de Brito Fontes <aledbf@gmail.com>
Date: Thu, 29 Nov 2018 22:45:32 -0300
Subject: [PATCH] Fix data size validations

---
 .../ingress/controller/template/template.go   | 29 +++-------
 .../controller/template/template_test.go      | 57 +++++++------------
 2 files changed, 29 insertions(+), 57 deletions(-)

diff --git a/internal/ingress/controller/template/template.go b/internal/ingress/controller/template/template.go
index 6aaf43e2e1..41b443f9f0 100644
--- a/internal/ingress/controller/template/template.go
+++ b/internal/ingress/controller/template/template.go
@@ -26,7 +26,7 @@ import (
 	"net/url"
 	"os"
 	"os/exec"
-	"strconv"
+	"regexp"
 	"strings"
 	text_template "text/template"
 	"time"
@@ -726,6 +726,10 @@ func buildNextUpstream(i, r interface{}) string {
 	return strings.Join(nextUpstreamCodes, " ")
 }
 
+var sizeUnitRegex = regexp.MustCompile("^[0-9]+[kKmMgG]{0,1}$")
+
+// isValidByteSize validates size units valid in nginx
+// http://nginx.org/en/docs/syntax.html
 func isValidByteSize(input interface{}) bool {
 	s, ok := input.(string)
 	if !ok {
@@ -733,32 +737,13 @@ func isValidByteSize(input interface{}) bool {
 		return false
 	}
 
+	s = strings.TrimSpace(s)
 	if s == "" {
 		glog.V(2).Info("empty byte size, hence it will not be set")
 		return false
 	}
 
-	_, err := strconv.Atoi(s)
-	if err != nil {
-		sLowercase := strings.ToLower(s)
-
-		check := strings.TrimSuffix(sLowercase, "k")
-		_, err := strconv.Atoi(check)
-		if err == nil {
-			return true
-		}
-
-		mCheck := strings.TrimSuffix(sLowercase, "m")
-		_, err = strconv.Atoi(mCheck)
-		if err == nil {
-			return true
-		}
-
-		glog.Errorf("incorrect byte size format '%v', hence it will not be set.", s)
-		return false
-	}
-
-	return true
+	return sizeUnitRegex.MatchString(s)
 }
 
 type ingressInformation struct {
diff --git a/internal/ingress/controller/template/template_test.go b/internal/ingress/controller/template/template_test.go
index 385ed7359f..02ee41af82 100644
--- a/internal/ingress/controller/template/template_test.go
+++ b/internal/ingress/controller/template/template_test.go
@@ -558,41 +558,28 @@ func TestBuildDenyVariable(t *testing.T) {
 }
 
 func TestBuildByteSize(t *testing.T) {
-	a := isValidByteSize("1000")
-	if !a {
-		t.Errorf("Expected '%v' but returned '%v'", true, a)
-	}
-	b := isValidByteSize("1000k")
-	if !b {
-		t.Errorf("Expected '%v' but returned '%v'", true, b)
-	}
-	c := isValidByteSize("1000m")
-	if !c {
-		t.Errorf("Expected '%v' but returned '%v'", true, c)
-	}
-	d := isValidByteSize("1000km")
-	if d {
-		t.Errorf("Expected '%v' but returned '%v'", false, d)
-	}
-	e := isValidByteSize("1000mk")
-	if e {
-		t.Errorf("Expected '%v' but returned '%v'", false, e)
-	}
-	f := isValidByteSize("1000kk")
-	if f {
-		t.Errorf("Expected '%v' but returned '%v'", false, f)
-	}
-	g := isValidByteSize("1000mm")
-	if g {
-		t.Errorf("Expected '%v' but returned '%v'", false, g)
-	}
-	h := isValidByteSize(nil)
-	if h {
-		t.Errorf("Expected '%v' but returned '%v'", false, h)
-	}
-	i := isValidByteSize("")
-	if i {
-		t.Errorf("Expected '%v' but returned '%v'", false, i)
+	cases := []struct {
+		input    interface{}
+		expected bool
+	}{
+		{"1000", true},
+		{"1000k", true},
+		{"1m", true},
+		{"10g", true},
+		{" 1m ", true},
+		{"1000kk", false},
+		{"1000km", false},
+		{"1mm", false},
+		{nil, false},
+		{"", false},
+		{"    ", false},
+	}
+
+	for _, tc := range cases {
+		val := isValidByteSize(tc.input)
+		if tc.expected != val {
+			t.Errorf("Expected '%v' but returned '%v'", tc.expected, val)
+		}
 	}
 }