replace manual parsing of head comments

This leaves more to the YAML parser so we have less chance of
accidentally introducing a bug or security issue in our own logic.

Also adds preservation of footer comments as a bonus!

Signed-off-by: Ashley Davis <[email protected]>

pkg/mod/index.go

@@ -1,7 +1,6 @@
 package mod
 
 import (
-	"bufio"
 	"io"
 	"path/filepath"
 	"sort"
@@ -45,54 +44,26 @@ func (w WorkDir) editKloneFile(fn func(*kloneFile) error) error {
 	if err != nil {
 		return err
 	}
+
 	defer file.Close()
 
-	index := kloneFile{}
+	var rawDocument yaml.Node
 
 	// decode current contents of index file
-	if err := yaml.NewDecoder(file).Decode(&index); err != nil && err != io.EOF {
+	if err := yaml.NewDecoder(file).Decode(&rawDocument); err != nil && err != io.EOF {
 		return err
 	}
 
-	// update index
-	if err := fn(&index); err != nil {
+	index := kloneFile{}
+
+	err = rawDocument.Decode(&index)
+	if err != nil {
 		return err
 	}
 
-	topComments := ""
-
-	{
-		// go back to the beginning of the file
-		if _, err := file.Seek(0, 0); err != nil {
-			return err
-		}
-
-		comments := strings.Builder{}
-
-		// read lines until the first non-comment line
-		reader := bufio.NewReader(file)
-		for {
-			line, isPrefix, err := reader.ReadLine()
-			if err != nil {
-				return err
-			}
-
-			if !isPrefix && (len(line) > 0 && line[0] != '#') {
-				break
-			}
-
-			if _, err := comments.Write(line); err != nil {
-				return err
-			}
-
-			if !isPrefix {
-				if _, err := comments.WriteRune('\n'); err != nil {
-					return err
-				}
-			}
-		}
-
-		topComments = comments.String()
+	// update index
+	if err := fn(&index); err != nil {
+		return err
 	}
 
 	// truncate file
@@ -103,9 +74,9 @@ func (w WorkDir) editKloneFile(fn func(*kloneFile) error) error {
 		return err
 	}
 
-	// write comments
-	if _, err := file.WriteString(topComments); err != nil {
-		return err
+	trimmedHeadComment := strings.TrimSpace(rawDocument.HeadComment)
+	if len(trimmedHeadComment) > 0 {
+		_, _ = file.WriteString(trimmedHeadComment + "\n\n")
 	}
 
 	encoder := yaml.NewEncoder(file)
@@ -115,6 +86,11 @@ func (w WorkDir) editKloneFile(fn func(*kloneFile) error) error {
 		return err
 	}
 
+	trimmedFootComment := strings.TrimSpace(rawDocument.FootComment)
+	if len(trimmedFootComment) > 0 {
+		_, _ = file.WriteString("\n\n" + trimmedFootComment)
+	}
+
 	return nil
 }