package decodebudget
import (
"bytes"
"encoding/json"
"errors"
"io"
)
type JSON struct {
MaxTokens int
MaxDepth int
}
func (b JSON) Validate(data []byte) error {
if b.MaxTokens <= 0 || b.MaxDepth <= 0 {
return errors.New("JSON budget limits must be positive")
}
dec := json.NewDecoder(bytes.NewReader(data))
dec.UseNumber()
depth, tokens := 0, 0
seen, complete := false, false
for {
token, err := dec.Token()
if err == io.EOF {
if !seen || depth != 0 {
return io.ErrUnexpectedEOF
}
return nil
}
if err != nil {
return err
}
if complete {
return errors.New("unexpected value after JSON document")
}
tokens++
if tokens > b.MaxTokens {
return ErrMessageComplexity
}
seen = true
if d, ok := token.(json.Delim); ok {
switch d {
case '{', '[':
depth++
case '}', ']':
depth--
}
if depth > b.MaxDepth {
return ErrMessageComplexity
}
}
if depth == 0 {
complete = true
}
}
}