From b403b364fef5b0758d6dd7e576ccb8b2dd14ae28 Mon Sep 17 00:00:00 2001
From: Roberto Rojas <robertojrojas@gmail.com>
Date: Thu, 5 Mar 2020 10:34:44 -0500
Subject: [PATCH] Get Kubernetes Version: Adds Retry Logic  (#593)

* Closes #551

* Closes #551

* Update cmd/kubernetes_version.go

Co-Authored-By: Liz Rice <liz@lizrice.com>

* Closes #551

Co-authored-by: Liz Rice <liz@lizrice.com>
---
 cmd/kubernetes_version.go      | 21 ++++++++++++++-
 cmd/kubernetes_version_test.go | 48 ++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/cmd/kubernetes_version.go b/cmd/kubernetes_version.go
index be0ff8b..2c55ed1 100644
--- a/cmd/kubernetes_version.go
+++ b/cmd/kubernetes_version.go
@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/golang/glog"
 )
@@ -30,7 +31,7 @@ func getKubeVersionFromRESTAPI() (string, error) {
 	}
 	token := strings.TrimSpace(string(tb))
 
-	data, err := getWebData(k8sVersionURL, token, tlsCert)
+	data, err := getWebDataWithRetry(k8sVersionURL, token, tlsCert)
 	if err != nil {
 		return "", err
 	}
@@ -42,6 +43,24 @@ func getKubeVersionFromRESTAPI() (string, error) {
 	return k8sVersion, nil
 }
 
+// The idea of this function is so if Kubernetes DNS is not completely seetup and the
+// Container where kube-bench is running needs time for DNS configure.
+// Basically try 10 times, waiting 1 second until either it is successful or it fails.
+func getWebDataWithRetry(k8sVersionURL, token string, cacert *tls.Certificate) (data []byte, err error) {
+	tries := 0
+	// We retry a few times in case the DNS service has not had time to come up
+	for tries < 10 {
+		data, err = getWebData(k8sVersionURL, token, cacert)
+		if err == nil {
+			return
+		}
+		tries++
+		time.Sleep(1 * time.Second)
+	}
+
+	return
+}
+
 func extractVersion(data []byte) (string, error) {
 	type versionResponse struct {
 		Major        string
diff --git a/cmd/kubernetes_version_test.go b/cmd/kubernetes_version_test.go
index 1513f62..ea3f4bf 100644
--- a/cmd/kubernetes_version_test.go
+++ b/cmd/kubernetes_version_test.go
@@ -126,7 +126,55 @@ func TestGetWebData(t *testing.T) {
 	}
 
 }
+func TestGetWebDataWithRetry(t *testing.T) {
+	okfn := func(w http.ResponseWriter, r *http.Request) {
+		_, _ = fmt.Fprintln(w, `{
+			"major": "1",
+			"minor": "15"}`)
+	}
+	errfn := func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, http.StatusText(http.StatusInternalServerError),
+			http.StatusInternalServerError)
+	}
+	token := "dummyToken"
+	var tlsCert tls.Certificate
+
+	cases := []struct {
+		fn   http.HandlerFunc
+		fail bool
+	}{
+		{
+			fn:   okfn,
+			fail: false,
+		},
+		{
+			fn:   errfn,
+			fail: true,
+		},
+	}
+
+	for id, c := range cases {
+		t.Run(strconv.Itoa(id), func(t *testing.T) {
+			ts := httptest.NewServer(c.fn)
+			defer ts.Close()
+			data, err := getWebDataWithRetry(ts.URL, token, &tlsCert)
+			if !c.fail {
+				if err != nil {
+					t.Errorf("unexpected error: %v", err)
+				}
 
+				if len(data) == 0 {
+					t.Errorf("missing data")
+				}
+			} else {
+				if err == nil {
+					t.Errorf("Expected error")
+				}
+			}
+		})
+	}
+
+}
 func TestExtractVersion(t *testing.T) {
 	okJSON := []byte(`{
 	"major": "1",