decoder.go

// Copyright (c) 2023-2025 Dell Inc. or its subsidiaries.  All rights reserved.
// Licensed to You under the Apache License, Version 2.0.   See the LICENSE file for more details.

// Package jsonstream is a stream-oriented JSON processing library for
// processing, editing, and manipulating JSON
package jsonstream

/*
Modifications
Copyright (c) 2021, Dell Inc  Michael Brown

Original
Copyright (c) 2020, Dave Cheney <dave@cheney.net>
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

/*
Modified from Dave Cheney's original:
 - Our use case is streaming tokenization only. No need for decoding into go
   data structures.
 - Need to bookkeep an accurate path for each token to allow upper layers to
   filter/modify
 - Need to support .Close() so that we can propagate error handling or EOF
   upstream/downstream
	- If a TokenIterator wraps another TokenIterator, ensure that the wrapper
	  calls .Close() on the wraped TokenIterator
*/

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"slices"
	"strconv"
	"sync"

	cheneyjson "github.com/pkg/json"
)

// const defaultBufferSize = 8192
//
//nolint:gochecknoglobals
var constBuffer = []byte{} // The buffer isn't used anymore. This gets rid of an alloc that isn't needed...

// A Decoder decodes JSON values from an input stream.
type Decoder struct {
	stack
	path      JSONPath
	scanner   cheneyjson.Scanner
	state     func(*Decoder) (JSONPath, []byte, error)
	closefunc func() error
}

// NewDecoder returns a new Decoder for the supplied Reader r.
func NewDecoder(r io.ReadCloser) *Decoder {
	return NewDecoderBuffer(r, constBuffer)
}

//nolint:gochecknoglobals
var (
	bufferPool = &sync.Pool{
		New: func() interface{} {
			return bytes.NewBuffer(make([]byte, 0, 128))
		},
	}
)

// NewDecoderBuffer returns a new Decoder for the supplier Reader r, using
// the []byte buf provided for working storage.
func NewDecoderBuffer(r io.ReadCloser, buf []byte) *Decoder {
	return &Decoder{
		scanner:   *cheneyjson.NewScanner(r),
		state:     (*Decoder).stateValue,
		path:      []interface{}{},
		closefunc: r.Close,
	}
}

type stack []bool

func (s *stack) push(v bool) {
	*s = append(*s, v)
}

func (s *stack) pop() bool {
	// skip last element and return the 'currently' last after slicing off last
	// sort of silly
	*s = (*s)[:len(*s)-1]
	if len(*s) == 0 {
		return false
	}
	return (*s)[len(*s)-1]
}

func (s *stack) len() int { return len(*s) }

// TODO: some ideas around adding methods to JSONPath for Equal(), HasPrefix() and then implement the matchfunc off that
type JSONPath []interface{}

// construct a JSONPath 'by hand' and automatically quote strings
// []byte are assumed to be already quoted!
func NewJSONPath(elems ...interface{}) JSONPath {
	ret := make([]interface{}, 0, len(elems))
	for _, elem := range elems {
		ret = append(ret, NewJSONPathElem(elem))
	}
	return ret
}

func NewJSONPathProp(elem string) []byte {
	return []byte(`"` + elem + `"`)
}

func NewJSONPathElem(elem interface{}) interface{} {
	switch t := elem.(type) {
	case int:
		return t
	case []byte:
		return t
	case string:
		return NewJSONPathProp(t)
	}
	panic(fmt.Sprintf("wrong type passed to NewJSONPathElem() = %T", elem))
}

func (s *JSONPath) push(v interface{}) {
	*s = append(*s, v)
}

func (s *JSONPath) pop() interface{} {
	if len(*s) == 0 {
		return nil
	}
	last := (*s)[len(*s)-1]
	(*s)[len(*s)-1] = nil
	*s = (*s)[:len(*s)-1]
	return last
}

func (s JSONPath) Len() int {
	return len(s)
}

// return a newly-allocated path struct that copies the underlying path strings to avoid buffer mutability problems
func (s JSONPath) Clone() JSONPath {
	q := make([]interface{}, len(s))
	for i := range s {
		switch t := s[i].(type) {
		case []byte:
			q[i] = slices.Clone(t)
		default:
			q[i] = s[i]
		}
	}
	return q
}

// PathPropElemFromEnd will return the 'last - fromEnd' element requested, if
// the last path element is a json property string. return nil, false if the
// path element is an array or fromEnd is out of bounds.
// THE PROPERTY STRING RETURNED HAS THE QUOTES REMOVED (ease of use)
//
//	ex: fromEnd=0  ==  return final path element.
func (s JSONPath) PathPropFromEnd(fromEnd int) (ret []byte, ok bool) {
	return s.PathProp(len(s) - fromEnd - 1)
}

// Check the element of a jsonpath at an index from the end of the path.
// Returns the byte array name of that element (if valid), the integer index
// of that element (if valid), and a boolean indicating whether anything was found.
// THE PROPERTY STRING RETURNED HAS THE QUOTES REMOVED (ease of use)
func (s JSONPath) PathPropOrIndexFromEnd(fromEnd int) (ret []byte, idx int, ok bool) {
	// Default idx to -1 since 0 is a valid index
	idx = -1
	// Out-of-bounds check
	if fromEnd < 0 || fromEnd > len(s)-1 {
		return
	}
	idxFromEnd := len(s) - 1 - fromEnd
	// Check if the element is a byte array
	ret, ok = s[idxFromEnd].([]byte)
	if ok {
		// Check if the byte array is properly quoted
		if len(ret) > 1 {
			ret = ret[1 : len(ret)-1]
		} else {
			ok = false
		}
		return
	}
	// Check if the element is an index
	idx, ok = s[idxFromEnd].(int)
	if ok {
		// Check if the index is valid
		ok = idx >= 0
	}
	return
}

func (s JSONPath) PathProp(elemIdx int) (ret []byte, ok bool) {
	if elemIdx < 0 || elemIdx > len(s)-1 {
		return nil, false
	}
	ret, ok = s[elemIdx].([]byte)
	if ok && len(ret) > 1 {
		ret = ret[1 : len(ret)-1]
	}
	return ret, ok
}

func (s JSONPath) PathArr(elemIdx int) (ret int, ok bool) {
	if elemIdx < 0 || elemIdx > len(s)-1 {
		return 0, false
	}
	ret, ok = s[elemIdx].(int)
	return ret, ok
}

func (s JSONPath) String() string {
	buffer := bufferPool.Get().(*bytes.Buffer)
	defer bufferPool.Put(buffer)
	buffer.Reset()
	for idx := range s {
		switch t := s[idx].(type) {
		case int:
			buffer.WriteRune('/')
			buf := buffer.AvailableBuffer()
			buf = strconv.AppendInt(buf, int64(t), 10)
			buffer.Write(buf)
		case []byte:
			buffer.WriteRune('/')
			if len(t) > 2 {
				buffer.Write(t[1 : len(t)-1])
				continue
			}
			buffer.WriteString("<nil>")
		}
	}
	return buffer.String()
}

func NewTokenFromString(p string) []byte {
	return []byte(`"` + p + `"`)
}

// NextToken breaks down the JSON stream into tuples that (hopefully?)
// efficiently represent the information in the JSON stream. The tuple is a
// (path, token) pair modelled (with a small exception) off the output from
// 'jq --stream'.
//
// The []byte is valid until NextToken is called again. DO NOT RETAIN
// REFERENCES TO THE RETURNED JSONPath or []byte across calls to .NextToken()!
// Make a copy if you need to save them!
//
// DIFFERENCES from encoding/json and pkg/json:
// NextToken does not return delimiters {, [, ], }, nor does it return commas
// or colons. The full JSON output can be reconstituted by using the path and
// token information.
//
// NextToken guarantees that the incoming JSON is valid with properly nested
// delimiters. Any invalid JSON will return an error at the point of the
// invalidity.
//
// Interpreting the path:
//
//	  path is an []interface{} that contains either []byte that contain JSON
//	    'property' strings, or int that contain the array index of JSON array
//	    values.
//	  from one call to NextToken to the next:
//		- only the last element can change from the previous call
//		- the path can shrink by exactly one element
//		- the path can grow arbitrarily: the last element of previous call can
//	       change and an arbitrary number of elements may be added
//
//	  When the path grows, that represents the beginning of an arbitrarily deep nesting of new objects.
//	  When the token returned is a nil []byte, that is the end of the current path object/array.
//	  An empty array is represented by a path that is longer than the previous path, with a nil token
//	  An empty object is represented by a path that is longer than the previous path, with a nil token
func (d *Decoder) NextToken() (JSONPath, []byte, error) {
	return d.state(d)
}

func (d *Decoder) Close() error {
	d.state = (*Decoder).stateEnd
	return d.closefunc()
}

var ErrMissingKey = errors.New("stateObjectString: missing string key")
var ErrObjectNoColon = errors.New("stateObjectColon: expecting colon")
var ErrExpectingComma = errors.New("stateObjectComma: expecting comma")
var ErrUnexpectedComma = errors.New("stateArrayValue: unexpected comma")
var ErrValueUnexpectedComma = errors.New("stateValue: unexpected comma")

func (d *Decoder) stateObjectString() (JSONPath, []byte, error) {
	tok := d.scanner.Next()
	if len(tok) < 1 {
		return d.path, nil, io.ErrUnexpectedEOF
	}

	d.path.pop()

	switch tok[0] {
	case cheneyjson.ObjectEnd:
		inObj := d.pop()
		switch {
		case d.len() == 0:
			d.state = (*Decoder).stateEnd
		case inObj:
			d.state = (*Decoder).stateObjectComma
		case !inObj:
			d.state = (*Decoder).stateArrayComma
		}
		return append(d.path, []byte(nil)), nil, nil // caller must not retain reference to any returned values
	case cheneyjson.String:
		d.state = (*Decoder).stateObjectColon
		newPathElem := slices.Clone(tok)
		d.path.push(newPathElem)
		return d.NextToken()
	default:
		return d.path, nil, ErrMissingKey
	}
}

func (d *Decoder) stateObjectColon() (JSONPath, []byte, error) {
	tok := d.scanner.Next()
	if len(tok) < 1 {
		return d.path, nil, io.ErrUnexpectedEOF
	}
	switch tok[0] {
	case cheneyjson.Colon:
		d.state = (*Decoder).stateObjectValue
		return d.NextToken()
	default:
		return d.path, tok, ErrObjectNoColon
	}
}

func (d *Decoder) stateObjectValue() (JSONPath, []byte, error) {
	tok := d.scanner.Next()
	if len(tok) < 1 {
		return d.path, nil, io.ErrUnexpectedEOF
	}
	switch tok[0] {
	case cheneyjson.ObjectStart:
		d.push(true)
		d.path.push([]byte(nil))
		d.state = (*Decoder).stateObjectString
		return d.NextToken()
	case cheneyjson.ArrayStart:
		d.push(false)
		d.path.push(0)
		d.state = (*Decoder).stateArrayValue
		return d.NextToken()
	default:
		d.state = (*Decoder).stateObjectComma
		return d.path, tok, nil
	}
}

func (d *Decoder) stateObjectComma() (JSONPath, []byte, error) {
	tok := d.scanner.Next()
	if len(tok) < 1 {
		return d.path, nil, io.ErrUnexpectedEOF
	}

	switch tok[0] {
	case cheneyjson.ObjectEnd:
		inObj := d.pop()
		switch {
		case d.len() == 0:
			d.state = (*Decoder).stateEnd
		case inObj:
			d.state = (*Decoder).stateObjectComma
		case !inObj:
			d.state = (*Decoder).stateArrayComma
		}

		d.path.pop()
		return append(d.path, []byte(nil)), nil, nil // caller must not retain reference to any returned values
	case cheneyjson.Comma:
		d.state = (*Decoder).stateObjectString
		return d.NextToken()
	default:
		return d.path, tok, ErrExpectingComma
	}
}

func (d *Decoder) stateArrayValue() (JSONPath, []byte, error) {
	tok := d.scanner.Next()
	if len(tok) < 1 {
		return d.path, nil, io.ErrUnexpectedEOF
	}
	switch tok[0] {
	case cheneyjson.ObjectStart:
		d.push(true)
		d.path.push([]byte(nil))
		d.state = (*Decoder).stateObjectString
		return d.NextToken()
	case cheneyjson.ArrayStart:
		d.push(false)
		d.path.push(0)
		d.state = (*Decoder).stateArrayValue
		return d.NextToken()
	case cheneyjson.ArrayEnd:
		inObj := d.pop()
		switch {
		case d.len() == 0:
			d.state = (*Decoder).stateEnd
		case inObj:
			d.state = (*Decoder).stateObjectComma
		case !inObj:
			d.state = (*Decoder).stateArrayComma
		}

		d.path.pop()
		return append(d.path, -1), nil, nil // caller must not retain reference to any returned values
	case cheneyjson.Comma:
		return d.path, nil, ErrUnexpectedComma
	default:
		d.state = (*Decoder).stateArrayComma
		return d.path, tok, nil
	}
}

func (d *Decoder) stateArrayComma() (JSONPath, []byte, error) {
	tok := d.scanner.Next()
	if len(tok) < 1 {
		return d.path, nil, io.ErrUnexpectedEOF
	}
	switch tok[0] {
	case cheneyjson.ArrayEnd:
		inObj := d.pop()
		switch {
		case d.len() == 0:
			d.state = (*Decoder).stateEnd
		case inObj:
			d.state = (*Decoder).stateObjectComma
		case !inObj:
			d.state = (*Decoder).stateArrayComma
		}

		d.path.pop()
		return append(d.path, -1), nil, nil // caller must not retain reference to any returned values
	case cheneyjson.Comma:
		idx := d.path.pop().(int) + 1
		d.path.push(idx)
		d.state = (*Decoder).stateArrayValue
		return d.NextToken()
	default:
		return d.path, nil, fmt.Errorf("stateArrayComma: expected comma, %v", d.stack)
	}
}

func (d *Decoder) stateValue() (JSONPath, []byte, error) {
	tok := d.scanner.Next()
	if len(tok) < 1 {
		return d.path, nil, io.ErrUnexpectedEOF
	}
	switch tok[0] {
	case cheneyjson.ObjectStart:
		d.push(true)
		d.path.push([]byte(nil))
		d.state = (*Decoder).stateObjectString
		return d.NextToken()
	case cheneyjson.ArrayStart:
		d.push(false)
		d.path.push(0)
		d.state = (*Decoder).stateArrayValue
		return d.NextToken()
	case cheneyjson.Comma:
		return d.path, nil, ErrValueUnexpectedComma
	default:
		d.state = (*Decoder).stateEnd
		return d.path, tok, nil
	}
}

func (d *Decoder) stateEnd() (JSONPath, []byte, error) { return nil, nil, io.EOF }