fix(config): Fix comment removal in TOML files

config/config.go

@@ -20,8 +20,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/compose-spec/compose-go/template"
-	"github.com/compose-spec/compose-go/utils"
 	"github.com/coreos/go-semver/semver"
 	"github.com/influxdata/toml"
 	"github.com/influxdata/toml/ast"
@@ -790,103 +788,6 @@ func parseConfig(contents []byte) (*ast.Table, error) {
 	return toml.Parse(outputBytes)
 }
 
-func removeComments(contents []byte) ([]byte, error) {
-	tomlReader := bytes.NewReader(contents)
-
-	// Initialize variables for tracking state
-	var inQuote, inComment, escaped bool
-	var quoteChar byte
-
-	// Initialize buffer for modified TOML data
-	var output bytes.Buffer
-
-	buf := make([]byte, 1)
-	// Iterate over each character in the file
-	for {
-		_, err := tomlReader.Read(buf)
-		if err != nil {
-			if errors.Is(err, io.EOF) {
-				break
-			}
-			return nil, err
-		}
-		char := buf[0]
-
-		// Toggle the escaped state at backslash to we have true every odd occurrence.
-		if char == '\\' {
-			escaped = !escaped
-		}
-
-		if inComment {
-			// If we're currently in a comment, check if this character ends the comment
-			if char == '\n' {
-				// End of line, comment is finished
-				inComment = false
-				_, _ = output.WriteRune('\n')
-			}
-		} else if inQuote {
-			// If we're currently in a quote, check if this character ends the quote
-			if char == quoteChar && !escaped {
-				// End of quote, we're no longer in a quote
-				inQuote = false
-			}
-			output.WriteByte(char)
-		} else {
-			// Not in a comment or a quote
-			if (char == '"' || char == '\'') && !escaped {
-				// Start of quote
-				inQuote = true
-				quoteChar = char
-				output.WriteByte(char)
-			} else if char == '#' && !escaped {
-				// Start of comment
-				inComment = true
-			} else {
-				// Not a comment or a quote, just output the character
-				output.WriteByte(char)
-			}
-		}
-
-		// Reset escaping if any other character occurred
-		if char != '\\' {
-			escaped = false
-		}
-	}
-	return output.Bytes(), nil
-}
-
-func substituteEnvironment(contents []byte, oldReplacementBehavior bool) ([]byte, error) {
-	options := []template.Option{
-		template.WithReplacementFunction(func(s string, m template.Mapping, cfg *template.Config) (string, error) {
-			result, applied, err := template.DefaultReplacementAppliedFunc(s, m, cfg)
-			if err == nil && !applied {
-				// Keep undeclared environment-variable patterns to reproduce
-				// pre-v1.27 behavior
-				return s, nil
-			}
-			if err != nil && strings.HasPrefix(err.Error(), "Invalid template:") {
-				// Keep invalid template patterns to ignore regexp substitutions
-				// like ${1}
-				return s, nil
-			}
-			return result, err
-		}),
-		template.WithoutLogging,
-	}
-	if oldReplacementBehavior {
-		options = append(options, template.WithPattern(oldVarRe))
-	}
-
-	envMap := utils.GetAsEqualsMap(os.Environ())
-	retVal, err := template.SubstituteWithOptions(string(contents), func(k string) (string, bool) {
-		if v, ok := envMap[k]; ok {
-			return v, ok
-		}
-		return "", false
-	}, options...)
-	return []byte(retVal), err
-}
-
 func (c *Config) addAggregator(name string, table *ast.Table) error {
 	creator, ok := aggregators.Aggregators[name]
 	if !ok {

config/envvar.go

@@ -0,0 +1,253 @@
+package config
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/compose-spec/compose-go/template"
+	"github.com/compose-spec/compose-go/utils"
+)
+
+type trimmer struct {
+	input  *bytes.Reader
+	output bytes.Buffer
+}
+
+func removeComments(buf []byte) ([]byte, error) {
+	t := &trimmer{
+		input:  bytes.NewReader(buf),
+		output: bytes.Buffer{},
+	}
+	err := t.process()
+	return t.output.Bytes(), err
+}
+
+func (t *trimmer) process() error {
+	for {
+		// Read the next byte until EOF
+		c, err := t.input.ReadByte()
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			return err
+		}
+
+		// Switch states if we need to
+		switch c {
+		case '\\':
+			_ = t.input.UnreadByte()
+			err = t.escape()
+		case '\'':
+			_ = t.input.UnreadByte()
+			if t.hasNQuotes(c, 3) {
+				err = t.tripleSingleQuote()
+			} else {
+				err = t.singleQuote()
+			}
+		case '"':
+			_ = t.input.UnreadByte()
+			if t.hasNQuotes(c, 3) {
+				err = t.tripleDoubleQuote()
+			} else {
+				err = t.doubleQuote()
+			}
+		case '#':
+			err = t.comment()
+		default:
+			if err := t.output.WriteByte(c); err != nil {
+				return err
+			}
+			continue
+		}
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			return err
+		}
+	}
+	return nil
+}
+
+func (t *trimmer) hasNQuotes(ref byte, limit int64) bool {
+	var count int64
+	// Look ahead check if the next characters are what we expect
+	for count = 0; count < limit; count++ {
+		c, err := t.input.ReadByte()
+		if err != nil || c != ref {
+			break
+		}
+	}
+	// We also need to unread the non-matching character
+	offset := -count
+	if count < limit {
+		offset--
+	}
+	// Unread the matched characters
+	_, _ = t.input.Seek(offset, io.SeekCurrent)
+	return count >= limit
+}
+
+func (t *trimmer) readWriteByte() (byte, error) {
+	c, err := t.input.ReadByte()
+	if err != nil {
+		return 0, err
+	}
+	return c, t.output.WriteByte(c)
+}
+
+func (t *trimmer) escape() error {
+	// Consumer the known starting backslash and quote
+	_, _ = t.readWriteByte()
+
+	// Read the next character which is the escaped one and exit
+	_, err := t.readWriteByte()
+	return err
+}
+
+func (t *trimmer) singleQuote() error {
+	// Consumer the known starting quote
+	_, _ = t.readWriteByte()
+
+	// Read bytes until EOF, line end or another single quote
+	for {
+		if c, err := t.readWriteByte(); err != nil || c == '\'' || c == '\n' {
+			return err
+		}
+	}
+}
+
+func (t *trimmer) tripleSingleQuote() error {
+	for i := 0; i < 3; i++ {
+		// Consumer the known starting quotes
+		_, _ = t.readWriteByte()
+	}
+
+	// Read bytes until EOF or another set of triple single quotes
+	for {
+		c, err := t.readWriteByte()
+		if err != nil {
+			return err
+		}
+
+		if c == '\'' && t.hasNQuotes('\'', 2) {
+			// Consumer the two additional ending quotes
+			_, _ = t.readWriteByte()
+			_, _ = t.readWriteByte()
+			return nil
+		}
+	}
+}
+
+func (t *trimmer) doubleQuote() error {
+	// Consumer the known starting quote
+	_, _ = t.readWriteByte()
+
+	// Read bytes until EOF, line end or another double quote
+	for {
+		c, err := t.input.ReadByte()
+		if err != nil {
+			return err
+		}
+		switch c {
+		case '\\':
+			// Found escaped character
+			_ = t.input.UnreadByte()
+			if err := t.escape(); err != nil {
+				return err
+			}
+			continue
+		case '"', '\n':
+			// Found terminator
+			return t.output.WriteByte(c)
+		}
+		if err := t.output.WriteByte(c); err != nil {
+			return err
+		}
+	}
+}
+
+func (t *trimmer) tripleDoubleQuote() error {
+	for i := 0; i < 3; i++ {
+		// Consumer the known starting quotes
+		_, _ = t.readWriteByte()
+	}
+
+	// Read bytes until EOF or another set of triple double quotes
+	for {
+		c, err := t.input.ReadByte()
+		if err != nil {
+			return err
+		}
+		switch c {
+		case '\\':
+			// Found escaped character
+			_ = t.input.UnreadByte()
+			if err := t.escape(); err != nil {
+				return err
+			}
+			continue
+		case '"':
+			_ = t.output.WriteByte(c)
+			if t.hasNQuotes('"', 2) {
+				// Consumer the two additional ending quotes
+				_, _ = t.readWriteByte()
+				_, _ = t.readWriteByte()
+				return nil
+			}
+			continue
+		}
+		if err := t.output.WriteByte(c); err != nil {
+			return err
+		}
+	}
+}
+
+func (t *trimmer) comment() error {
+	// Read bytes until EOF or a line break
+	for {
+		c, err := t.input.ReadByte()
+		if err != nil {
+			return err
+		}
+		if c == '\n' {
+			return t.output.WriteByte(c)
+		}
+	}
+}
+
+func substituteEnvironment(contents []byte, oldReplacementBehavior bool) ([]byte, error) {
+	options := []template.Option{
+		template.WithReplacementFunction(func(s string, m template.Mapping, cfg *template.Config) (string, error) {
+			result, applied, err := template.DefaultReplacementAppliedFunc(s, m, cfg)
+			if err == nil && !applied {
+				// Keep undeclared environment-variable patterns to reproduce
+				// pre-v1.27 behavior
+				return s, nil
+			}
+			if err != nil && strings.HasPrefix(err.Error(), "Invalid template:") {
+				// Keep invalid template patterns to ignore regexp substitutions
+				// like ${1}
+				return s, nil
+			}
+			return result, err
+		}),
+		template.WithoutLogging,
+	}
+	if oldReplacementBehavior {
+		options = append(options, template.WithPattern(oldVarRe))
+	}
+
+	envMap := utils.GetAsEqualsMap(os.Environ())
+	retVal, err := template.SubstituteWithOptions(string(contents), func(k string) (string, bool) {
+		if v, ok := envMap[k]; ok {
+			return v, ok
+		}
+		return "", false
+	}, options...)
+	return []byte(retVal), err
+}