From dd15f836dd533988ea4e6d79e46d8d15fa9ff25e Mon Sep 17 00:00:00 2001
From: skrishna <skrishna@cloudera.com>
Date: Fri, 15 Sep 2023 12:56:45 +0530
Subject: [PATCH] [YUNIKORN-1875] requested changes incorporated

Import ordering changes, error handling and logging for gzipReader close, added testcases for decompress.
---
 pkg/common/utils/utils_test.go |  2 +-
 pkg/conf/schedulerconf.go      |  8 ++++++++
 pkg/conf/schedulerconf_test.go | 25 ++++++++++++++++++++++---
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/pkg/common/utils/utils_test.go b/pkg/common/utils/utils_test.go
index 0d478d7b6..e921f3806 100644
--- a/pkg/common/utils/utils_test.go
+++ b/pkg/common/utils/utils_test.go
@@ -23,7 +23,6 @@ import (
 	"compress/gzip"
 	"encoding/base64"
 	"fmt"
-	"github.com/apache/yunikorn-core/pkg/common/configs"
 	"testing"
 	"time"
 
@@ -33,6 +32,7 @@ import (
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
+	"github.com/apache/yunikorn-core/pkg/common/configs"
 	"github.com/apache/yunikorn-k8shim/pkg/common"
 	"github.com/apache/yunikorn-k8shim/pkg/common/constants"
 	"github.com/apache/yunikorn-k8shim/pkg/conf"
diff --git a/pkg/conf/schedulerconf.go b/pkg/conf/schedulerconf.go
index 9dddb96e5..5a99b4421 100644
--- a/pkg/conf/schedulerconf.go
+++ b/pkg/conf/schedulerconf.go
@@ -466,6 +466,7 @@ func Decompress(key string, value []byte) (string, string) {
 	n, err := base64.StdEncoding.Decode(decodedValue, value)
 	if err != nil {
 		log.Log(log.ShimConfig).Error("failed to decode schedulerConfig entry", zap.Error(err))
+		return "", ""
 	}
 	decodedValue = decodedValue[:n]
 	splitKey := strings.Split(key, ".")
@@ -475,10 +476,17 @@ func Decompress(key string, value []byte) (string, string) {
 		gzReader, err := gzip.NewReader(reader)
 		if err != nil {
 			log.Log(log.ShimConfig).Error("failed to decompress decoded schedulerConfig entry", zap.Error(err))
+			return "", ""
 		}
+		defer func() {
+			if err := gzReader.Close(); err != nil {
+				log.Log(log.ShimConfig).Debug("gzip Reader could not be closed ", zap.Error(err))
+			}
+		}()
 		decompressedBytes, err := io.ReadAll(gzReader)
 		if err != nil {
 			log.Log(log.ShimConfig).Error("failed to decompress decoded schedulerConfig entry", zap.Error(err))
+			return "", ""
 		}
 		uncompressedData = string(decompressedBytes)
 	}
diff --git a/pkg/conf/schedulerconf_test.go b/pkg/conf/schedulerconf_test.go
index 49eedfe9e..f79c108a1 100644
--- a/pkg/conf/schedulerconf_test.go
+++ b/pkg/conf/schedulerconf_test.go
@@ -22,14 +22,15 @@ import (
 	"compress/gzip"
 	"encoding/base64"
 	"fmt"
-	"github.com/apache/yunikorn-core/pkg/common/configs"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
 
 	"gotest.tools/v3/assert"
 	v1 "k8s.io/api/core/v1"
 
+	"github.com/apache/yunikorn-core/pkg/common/configs"
 	"github.com/apache/yunikorn-k8shim/pkg/common/constants"
 )
 
@@ -80,17 +81,35 @@ func TestDecompress(t *testing.T) {
 	var b bytes.Buffer
 	gzWriter := gzip.NewWriter(&b)
 	if _, err := gzWriter.Write([]byte(configs.DefaultSchedulerConfig)); err != nil {
-		t.Fatal("expected nil, got error while compressing test schedulerConfig")
+		assert.NilError(t, err, "expected nil, got error while compressing test schedulerConfig")
 	}
 	if err := gzWriter.Close(); err != nil {
+		assert.NilError(t, err, "expected nil, got error")
 		t.Fatal("expected nil, got error")
 	}
 	encodedConfigString := make([]byte, base64.StdEncoding.EncodedLen(len(b.Bytes())))
 	base64.StdEncoding.Encode(encodedConfigString, b.Bytes())
-	_, decodedConfigString := Decompress("queues.yaml."+constants.GzipSuffix, encodedConfigString)
+	key, decodedConfigString := Decompress("queues.yaml."+constants.GzipSuffix, encodedConfigString)
+	assert.Equal(t, "queues.yaml", key)
 	assert.Equal(t, configs.DefaultSchedulerConfig, decodedConfigString)
 }
 
+func TestDecompressUnkownKey(t *testing.T) {
+	encodedConfigString := make([]byte, base64.StdEncoding.EncodedLen(len([]byte(configs.DefaultSchedulerConfig))))
+	base64.StdEncoding.Encode(encodedConfigString, []byte(configs.DefaultSchedulerConfig))
+	key, decodedConfigString := Decompress("queues.yaml.bin", encodedConfigString)
+	assert.Equal(t, "queues.yaml", key)
+	assert.Assert(t, len(decodedConfigString) == 0, "expected decodedConfigString to be nil")
+}
+
+func TestDecompressBadCompression(t *testing.T) {
+	encodedConfigString := make([]byte, base64.StdEncoding.EncodedLen(len([]byte(configs.DefaultSchedulerConfig))))
+	base64.StdEncoding.Encode(encodedConfigString, []byte(configs.DefaultSchedulerConfig))
+	key, decodedConfigString := Decompress("queues.yaml."+constants.GzipSuffix, encodedConfigString)
+	assert.Equal(t, "", key)
+	assert.Assert(t, !strings.EqualFold(configs.DefaultSchedulerConfig, decodedConfigString), "expected decodedConfigString to be nil")
+}
+
 func TestParseConfigMap(t *testing.T) {
 	testCases := []struct {
 		name  string