optimized yaml utility function

tenable · Nov 12, 2020 · 09d81f8 · 09d81f8
1 parent 442e39a
commit 09d81f8
Showing 1 changed file with 49 additions and 105 deletions.
diff --git a/pkg/utils/yaml.go b/pkg/utils/yaml.go
@@ -26,118 +26,76 @@ var (
 func LoadYAML(filePath string) ([]*IacDocument, error) {
 	iacDocumentList := make([]*IacDocument, 0)
 
-	// First pass determines line number data
-	{ // Limit the scope for Close()
-		file, err := os.Open(filePath)
-		if err != nil {
-			return iacDocumentList, err
-		}
-		defer file.Close()
-
-		scanner := bufio.NewScanner(file)
-		startLineNumber := 1
-		currentLineNumber := 1
-		for scanner.Scan() {
-			if strings.HasPrefix(scanner.Text(), "---") {
-				// We've found the end-of-directives marker, so record results for the current document
-				iacDocumentList = append(iacDocumentList, &IacDocument{
-					Type:      YAMLDoc,
-					StartLine: startLineNumber,
-					EndLine:   currentLineNumber,
-					FilePath:  filePath,
-				})
-				startLineNumber = currentLineNumber + 1
-			}
-			currentLineNumber++
-		}
-
-		// Add the very last entry
-		iacDocumentList = append(iacDocumentList, &IacDocument{
-			Type:      YAMLDoc,
-			StartLine: startLineNumber,
-			EndLine:   currentLineNumber,
-			FilePath:  filePath,
-		})
-
-		if err = scanner.Err(); err != nil {
-			return iacDocumentList, err
-		}
+	file, err := os.Open(filePath)
+	if err != nil {
+		return iacDocumentList, err
 	}
+	defer file.Close()
 
-	// Second pass extracts all YAML documents and saves it in the document struct
 	fileBytes, err := ioutil.ReadFile(filePath)
 	if err != nil {
 		return iacDocumentList, err
 	}
 
-	dec := yaml.NewDecoder(bytes.NewReader(fileBytes))
-	i := 0
-	for {
-		// each iteration extracts and marshals one yaml document
-		var value interface{}
-		err = dec.Decode(&value)
-		if err == io.EOF {
-			break
-		}
-		if err != nil {
-			return iacDocumentList, err
-		}
-		if i > (len(iacDocumentList) - 1) {
-			return iacDocumentList, errHighDocumentCount
-		}
-
-		var documentBytes []byte
-		documentBytes, err = yaml.Marshal(value)
-		if err != nil {
-			return iacDocumentList, err
-		}
-		iacDocumentList[i].Data = documentBytes
-		i++
-	}
-
-	return iacDocumentList, nil
+	return getIacDocumentList(bufio.NewScanner(file), fileBytes, filePath)
 }
 
 // LoadYAMLString loads a YAML String. Can return one or more IaC Documents.
 // Besides reading in file data, its main purpose is to determine and store line number and filename metadata
 func LoadYAMLString(data, absFilePath string) ([]*IacDocument, error) {
+	return getIacDocumentList(bufio.NewScanner(strings.NewReader(data)), []byte(data), absFilePath)
+}
+
+// ReadYamlFile reads a yaml file and load content in a map[string]interface{} type
+func ReadYamlFile(path string) (map[string]interface{}, error) {
+	dat, err := ioutil.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+	var output map[string]interface{}
+	err = yaml.Unmarshal(dat, &output)
+	if err != nil {
+		return nil, err
+	}
+	return output, nil
+}
+
+// getIacDocumentList provides one or more IaC Documents.
+// Besides reading in file data, its main purpose is to determine and store line number and filename metadata
+func getIacDocumentList(scanner *bufio.Scanner, bytearray []byte, filePath string) ([]*IacDocument, error) {
 	iacDocumentList := make([]*IacDocument, 0)
 
 	// First pass determines line number data
-	{ // Limit the scope for Close()
-		scanner := bufio.NewScanner(strings.NewReader(data))
-		startLineNumber := 1
-		currentLineNumber := 1
-		for scanner.Scan() {
-			if strings.HasPrefix(scanner.Text(), "---") {
-				// We've found the end-of-directives marker, so record results for the current document
-				iacDocumentList = append(iacDocumentList, &IacDocument{
-					Type:      YAMLDoc,
-					StartLine: startLineNumber,
-					EndLine:   currentLineNumber,
-					FilePath:  absFilePath,
-				})
-				startLineNumber = currentLineNumber + 1
-			}
-			currentLineNumber++
+	startLineNumber := 1
+	currentLineNumber := 1
+	for scanner.Scan() {
+		if strings.HasPrefix(scanner.Text(), "---") {
+			// We've found the end-of-directives marker, so record results for the current document
+			iacDocumentList = append(iacDocumentList, &IacDocument{
+				Type:      YAMLDoc,
+				StartLine: startLineNumber,
+				EndLine:   currentLineNumber,
+				FilePath:  filePath,
+			})
+			startLineNumber = currentLineNumber + 1
 		}
+		currentLineNumber++
+	}
 
-		// Add the very last entry
-		iacDocumentList = append(iacDocumentList, &IacDocument{
-			Type:      YAMLDoc,
-			StartLine: startLineNumber,
-			EndLine:   currentLineNumber,
-			FilePath:  "",
-		})
+	// Add the very last entry
+	iacDocumentList = append(iacDocumentList, &IacDocument{
+		Type:      YAMLDoc,
+		StartLine: startLineNumber,
+		EndLine:   currentLineNumber,
+		FilePath:  "",
+	})
 
-		if err := scanner.Err(); err != nil {
-			return iacDocumentList, err
-		}
+	if err := scanner.Err(); err != nil {
+		return iacDocumentList, err
 	}
 
-	// Second pass extracts all YAML documents and saves it in the document struct
+	dec := yaml.NewDecoder(bytes.NewReader(bytearray))
 
-	dec := yaml.NewDecoder(bytes.NewReader([]byte(data)))
 	i := 0
 	for {
 		// each iteration extracts and marshals one yaml document
@@ -164,17 +122,3 @@ func LoadYAMLString(data, absFilePath string) ([]*IacDocument, error) {
 
 	return iacDocumentList, nil
 }
-
-// ReadYamlFile reads a yaml file and load content in a map[string]interface{} type
-func ReadYamlFile(path string) (map[string]interface{}, error) {
-	dat, err := ioutil.ReadFile(path)
-	if err != nil {
-		return nil, err
-	}
-	var output map[string]interface{}
-	err = yaml.Unmarshal(dat, &output)
-	if err != nil {
-		return nil, err
-	}
-	return output, nil
-}