Rewrite JSON formatter to avoid decoder internals

internal/utils/utils.go

@@ -9,7 +9,6 @@ import (
 	"io"
 	"os"
 	"os/exec"
-	"reflect"
 	"regexp"
 	"strconv"
 	"strings"
@@ -44,18 +43,6 @@ type QueryOptions struct {
 	Colors   int
 }
 
-const (
-	jsonTokenTopValue = iota
-	jsonTokenArrayStart
-	jsonTokenArrayValue
-	jsonTokenArrayComma
-	jsonTokenObjectStart
-	jsonTokenObjectKey
-	jsonTokenObjectColon
-	jsonTokenObjectValue
-	jsonTokenObjectComma
-)
-
 func FormatXml(reader io.Reader, writer io.Writer, indent string, colors int) error {
 	decoder := xml.NewDecoder(reader)
 	decoder.Strict = false
@@ -435,85 +422,144 @@ func FormatJson(reader io.Reader, writer io.Writer, indent string, colors int) e
 	attrColor := color.New(color.FgHiBlue).SprintFunc()
 	valueColor := color.New(color.FgGreen).SprintFunc()
 
-	level := 0
-	suffix := ""
-	prefix := ""
 	newline := "\n"
 	if indent == "" {
 		newline = ""
 	}
+	write := func(args ...any) error {
+		_, err := fmt.Fprint(writer, args...)
+		return err
+	}
 
-	for {
-		token, err := decoder.Token()
+	var formatToken func(json.Token, int) error
+	formatToken = func(token json.Token, level int) error {
+		switch typedToken := token.(type) {
+		case json.Delim:
+			switch rune(typedToken) {
+			case '{':
+				if err := write(tagColor("{")); err != nil {
+					return err
+				}
+				level++
+				index := 0
+				for decoder.More() {
+					if index > 0 {
+						if err := write(",", newline); err != nil {
+							return err
+						}
+					} else if err := write(newline); err != nil {
+						return err
+					}
+					if err := write(strings.Repeat(indent, level)); err != nil {
+						return err
+					}
 
-		if err == io.EOF {
-			break
-		}
+					keyToken, err := decoder.Token()
+					if err != nil {
+						return err
+					}
+					key, ok := keyToken.(string)
+					if !ok {
+						return fmt.Errorf("expected JSON object key, got %T", keyToken)
+					}
 
-		if err != nil {
-			return err
-		}
+					valueToken, err := decoder.Token()
+					if err != nil {
+						return err
+					}
 
-		v := reflect.ValueOf(*decoder)
-		tokenState := v.FieldByName("tokenState").Int()
+					if err := write(attrColor(strconv.Quote(key)), ": "); err != nil {
+						return err
+					}
+					if err := formatToken(valueToken, level); err != nil {
+						return err
+					}
+					index++
+				}
+				level--
 
-		switch tokenType := token.(type) {
-		case json.Delim:
-			switch rune(tokenType) {
-			case '{':
-				_, _ = fmt.Fprint(writer, prefix, tagColor("{"), newline)
-				level++
-				suffix = strings.Repeat(indent, level)
-			case '}':
-				if level > 0 {
-					level--
+				endToken, err := decoder.Token()
+				if err != nil {
+					return err
 				}
-				_, _ = fmt.Fprint(writer, newline, strings.Repeat(indent, level), tagColor("}"))
-				if tokenState == jsonTokenArrayComma {
-					suffix = "," + newline + strings.Repeat(indent, level)
+				if endToken != json.Delim('}') {
+					return fmt.Errorf("expected JSON object end, got %v", endToken)
 				}
+				if index > 0 {
+					return write(newline, strings.Repeat(indent, level), tagColor("}"))
+				}
+				return write(tagColor("}"))
 			case '[':
-				_, _ = fmt.Fprint(writer, prefix, tagColor("["), newline)
+				if err := write(tagColor("[")); err != nil {
+					return err
+				}
 				level++
-				suffix = strings.Repeat(indent, level)
-			case ']':
-				if level > 0 {
-					level--
+				index := 0
+				for decoder.More() {
+					if index > 0 {
+						if err := write(",", newline); err != nil {
+							return err
+						}
+					} else if err := write(newline); err != nil {
+						return err
+					}
+					if err := write(strings.Repeat(indent, level)); err != nil {
+						return err
+					}
+
+					valueToken, err := decoder.Token()
+					if err != nil {
+						return err
+					}
+					if err := formatToken(valueToken, level); err != nil {
+						return err
+					}
+					index++
+				}
+				level--
+
+				endToken, err := decoder.Token()
+				if err != nil {
+					return err
+				}
+				if endToken != json.Delim(']') {
+					return fmt.Errorf("expected JSON array end, got %v", endToken)
 				}
-				_, _ = fmt.Fprint(writer, newline, strings.Repeat(indent, level), tagColor("]"))
+				if index > 0 {
+					return write(newline, strings.Repeat(indent, level), tagColor("]"))
+				}
+				return write(tagColor("]"))
+			default:
+				return fmt.Errorf("unexpected JSON delimiter %q", typedToken)
 			}
 		case string:
-			escapedToken := strconv.Quote(token.(string))
-			value := valueColor(escapedToken)
-			if tokenState == jsonTokenObjectColon {
-				value = attrColor(escapedToken)
-			}
-			_, _ = fmt.Fprintf(writer, "%s%s", prefix, value)
-		case float64:
-			_, _ = fmt.Fprintf(writer, "%s%v", prefix, valueColor(token))
-		case json.Number:
-			_, _ = fmt.Fprintf(writer, "%s%v", prefix, valueColor(token))
-		case bool:
-			_, _ = fmt.Fprintf(writer, "%s%v", prefix, valueColor(token))
+			return write(valueColor(strconv.Quote(typedToken)))
+		case float64, json.Number, bool:
+			return write(valueColor(typedToken))
 		case nil:
-			_, _ = fmt.Fprintf(writer, "%s%s", prefix, valueColor("null"))
+			return write(valueColor("null"))
 		}
 
-		switch tokenState {
-		case jsonTokenObjectColon:
-			suffix = ": "
-		case jsonTokenObjectComma:
-			suffix = "," + newline + strings.Repeat(indent, level)
-		case jsonTokenArrayComma:
-			suffix = "," + newline + strings.Repeat(indent, level)
+		return nil
+	}
+
+	for {
+		token, err := decoder.Token()
+
+		if err == io.EOF {
+			break
 		}
 
-		prefix = suffix
-	}
+		if err != nil {
+			return err
+		}
 
-	_, _ = fmt.Fprint(writer, "\n")
+		if err := formatToken(token, 0); err != nil {
+			return err
+		}
+	}
 
-	return nil
+	return write("\n")
 }
 
 func IsHTML(input string) bool {

internal/utils/utils_test.go

@@ -2,6 +2,7 @@ package utils
 
 import (
 	"bytes"
+	"errors"
 	"io"
 	"os"
 	"path/filepath"
@@ -11,6 +12,28 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+var errWriteFailed = errors.New("write failed")
+
+type failingWriter struct{}
+
+func (failingWriter) Write([]byte) (int, error) {
+	return 0, errWriteFailed
+}
+
+type failOnWriteWriter struct {
+	count  int
+	failOn int
+}
+
+func (writer *failOnWriteWriter) Write(data []byte) (int, error) {
+	writer.count++
+	if writer.count == writer.failOn {
+		return 0, errWriteFailed
+	}
+
+	return len(data), nil
+}
+
 func getFileReader(filename string) io.Reader {
 	reader, err := os.Open(filename)
 
@@ -100,6 +123,32 @@ func TestFormatJson(t *testing.T) {
 		assert.Nil(t, formatErr)
 		assert.Equal(t, expectedJson, output.String())
 	}
+
+	tests := []struct {
+		name     string
+		input    string
+		expected string
+	}{
+		{"empty object", "{}", "{}\n"},
+		{"empty array", "[]", "[]\n"},
+		{"nested empty containers in object", "{\"a\":{},\"b\":[]}", "{\n  \"a\": {},\n  \"b\": []\n}\n"},
+		{"nested empty containers in array", "[{},[]]", "[\n  {},\n  []\n]\n"},
+	}
+
+	for _, testCase := range tests {
+		t.Run(testCase.name, func(t *testing.T) {
+			output := new(strings.Builder)
+			formatErr := FormatJson(strings.NewReader(testCase.input), output, "  ", ColorsDisabled)
+			assert.NoError(t, formatErr)
+			assert.Equal(t, testCase.expected, output.String())
+		})
+	}
+
+	formatErr := FormatJson(strings.NewReader("{}"), failingWriter{}, "  ", ColorsDisabled)
+	assert.ErrorIs(t, formatErr, errWriteFailed)
+
+	formatErr = FormatJson(strings.NewReader("{}"), &failOnWriteWriter{failOn: 3}, "  ", ColorsDisabled)
+	assert.ErrorIs(t, formatErr, errWriteFailed)
 }
 
 func TestXPathQuery(t *testing.T) {