pkg/jsplugin/json_state.go
Rama main · réplica de lectura del repositorio original.
package jsplugin
import (
"context"
"fmt"
"math"
"strings"
"unicode/utf8"
"github.com/Calcium-Ion/moejs/engine"
"github.com/QuantumNous/new-api/common"
)
const (
CapabilityJSONClone = "json-clone@1"
CapabilitySubmitSSEDelta = "submit-sse-delta@1"
CapabilityJSONOrder = "json-order@1"
CapabilityDurationAuto = "duration-auto@1"
MaxJSONToolBytes = 1 << 20
maxJSONToolDepth = 32
maxJSONToolNodes = 32768
maxJSONChanges = 256
)
// HasCapability describes host APIs independently of a plugin's mutable globals.
func HasCapability(name string) bool {
return name == CapabilityJSONClone || name == CapabilitySubmitSSEDelta || name == CapabilityJSONOrder || name == CapabilityDurationAuto
}
// ValidateJSONPath checks a path into a JSON document by the rules of a JSON
// change path: at most 32 segments, each a string key or a non-negative
// integer array index.
func ValidateJSONPath(path []any) error {
if len(path) > maxJSONToolDepth {
return fmt.Errorf("JSON path must have at most %d segments", maxJSONToolDepth)
}
for _, segment := range path {
switch typed := segment.(type) {
case string:
if !utf8.ValidString(typed) {
return fmt.Errorf("JSON path keys must be valid strings")
}
case int:
if typed < 0 {
return fmt.Errorf("JSON path indices must be non-negative integers")
}
case int64:
if typed < 0 {
return fmt.Errorf("JSON path indices must be non-negative integers")
}
case float64:
if math.IsNaN(typed) || typed < 0 || typed > math.MaxInt32 || math.Trunc(typed) != typed {
return fmt.Errorf("JSON path indices must be non-negative integers")
}
default:
return fmt.Errorf("JSON path segments must be strings or integers")
}
}
return nil
}
// JSONState owns a request-local JSON result. Appended strings stay in Go and
// are materialized only when the stream finishes, independently of JS runtimes.
type JSONState struct {
root *jsonStateNode
limit int
failed error
}
type jsonStateNode struct {
scalar any
object map[string]*jsonStateNode
array []*jsonStateNode
text *strings.Builder
bytes int
nodes int
}
type jsonStateBudget struct {
ctx context.Context
bytes int
nodes int
// realm reads JavaScript values; only json.clone passes them.
realm *engine.Realm
}
func (b *jsonStateBudget) spend(bytes, nodes int) error {
if err := b.ctx.Err(); err != nil {
return err
}
if bytes > b.bytes || nodes > b.nodes {
return fmt.Errorf("JSON state exceeds size or node limit")
}
b.bytes -= bytes
b.nodes -= nodes
return nil
}
func newJSONStateNode(value any, depth int, budget *jsonStateBudget) (*jsonStateNode, error) {
if depth > maxJSONToolDepth {
return nil, fmt.Errorf("JSON state exceeds depth limit")
}
if js, ok := value.(engine.Value); ok {
if js.IsUndefined() {
return nil, fmt.Errorf("undefined is not a JSON value")
}
if !js.IsObject() {
value = budget.realm.ToGo(js)
} else {
// Getters and proxy traps run while the value is read; what they
// throw is returned as it is, not as a JSON error.
object := js.AsObject()
switch {
case object.ClassName() == "Object" && !object.IsCallable():
keys, err := budget.realm.EnumerableOwnKeys(object)
if err != nil {
return nil, err
}
if len(keys) > budget.nodes {
return nil, fmt.Errorf("JSON object exceeds node limit")
}
fields := make(map[string]any, len(keys))
for _, key := range keys {
if fields[key.GoString()], err = object.GetProp(budget.realm, key); err != nil {
return nil, err
}
}
value = fields
case object.ClassName() == "Array":
length, err := budget.realm.LengthOfArrayLike(object)
if err != nil {
return nil, err
}
if length > int64(budget.nodes) {
return nil, fmt.Errorf("JSON array exceeds node limit")
}
items := make([]any, int(length))
for index := range items {
key := engine.IndexKey(uint32(index))
present, err := budget.realm.HasOwn(object, key)
if err != nil {
return nil, err
}
if !present {
return nil, fmt.Errorf("sparse arrays are not JSON values")
}
if items[index], err = object.GetProp(budget.realm, key); err != nil {
return nil, err
}
}
value = items
default:
return nil, fmt.Errorf("json.clone accepts only plain objects, arrays and JSON scalars")
}
}
}
beforeBytes, beforeNodes := budget.bytes, budget.nodes
if err := budget.spend(0, 1); err != nil {
return nil, err
}
switch typed := value.(type) {
case map[string]any:
if typed == nil {
value = nil
}
case []any:
if typed == nil {
value = nil
}
}
node := &jsonStateNode{}
switch typed := value.(type) {
case map[string]any:
if len(typed) > budget.nodes {
return nil, fmt.Errorf("JSON object exceeds node limit")
}
if err := budget.spend(2, 0); err != nil {
return nil, err
}
node.object = make(map[string]*jsonStateNode, len(typed))
for key, child := range typed {
if !utf8.ValidString(key) || len(key) > budget.bytes {
return nil, fmt.Errorf("invalid or oversized JSON object key")
}
encoded, err := common.Marshal(key)
if err != nil {
return nil, err
}
cost := len(encoded) + 1
if len(node.object) != 0 {
cost++
}
if err = budget.spend(cost, 0); err != nil {
return nil, err
}
node.object[key], err = newJSONStateNode(child, depth+1, budget)
if err != nil {
return nil, err
}
}
case []any:
if len(typed) > budget.nodes {
return nil, fmt.Errorf("JSON array exceeds node limit")
}
if err := budget.spend(2+max(0, len(typed)-1), 0); err != nil {
return nil, err
}
node.array = make([]*jsonStateNode, 0, len(typed))
for _, child := range typed {
item, err := newJSONStateNode(child, depth+1, budget)
if err != nil {
return nil, err
}
node.array = append(node.array, item)
}
default:
switch number := value.(type) {
case int:
value = float64(number)
case int64:
value = float64(number)
}
switch scalar := value.(type) {
case nil, bool:
case float64:
if math.IsNaN(scalar) || math.IsInf(scalar, 0) {
return nil, fmt.Errorf("JSON numbers must be finite")
}
case string:
if len(scalar) > budget.bytes {
return nil, fmt.Errorf("JSON string exceeds size limit")
}
if !utf8.ValidString(scalar) {
value = string([]rune(scalar))
}
default:
return nil, fmt.Errorf("value must contain only JSON objects, arrays and scalars")
}
encoded, err := common.Marshal(value)
if err != nil {
return nil, err
}
if err = budget.spend(len(encoded), 0); err != nil {
return nil, err
}
node.scalar = value
}
node.bytes, node.nodes = beforeBytes-budget.bytes, beforeNodes-budget.nodes
return node, nil
}
func NewJSONState(limit int) *JSONState {
return &JSONState{root: &jsonStateNode{bytes: 4, nodes: 1}, limit: max(0, min(limit, MaxJSONToolBytes))}
}
// Apply accepts a bounded batch of set, append and appendText operations.
// Failure invalidates the stream: callers must discard the state rather than
// continue after an error. No partially updated result may be published.
func (s *JSONState) Apply(ctx context.Context, changes any) (err error) {
if s.failed != nil {
return s.failed
}
defer func() {
if err != nil {
s.failed = err
}
}()
if err = ctx.Err(); err != nil {
return err
}
items, ok := changes.([]any)
if !ok || len(items) > maxJSONChanges {
return fmt.Errorf("changes must be an array of at most %d operations", maxJSONChanges)
}
for _, item := range items {
if err := ctx.Err(); err != nil {
return err
}
change, ok := item.(map[string]any)
if !ok || len(change) != 3 {
return fmt.Errorf("each change must contain only op, path and value")
}
op, _ := change["op"].(string)
path, ok := change["path"].([]any)
value, hasValue := change["value"]
if !ok || !hasValue || len(path) > maxJSONToolDepth {
return fmt.Errorf("change requires a JSON value and a path of at most %d segments", maxJSONToolDepth)
}
if op != "set" && op != "append" && op != "appendText" {
return fmt.Errorf("unsupported JSON change operation %q", op)
}
ancestors := make([]*jsonStateNode, 0, len(path))
target := s.root
var key string
var index int
for offset, segment := range path {
if target == nil {
return fmt.Errorf("JSON change parent does not exist")
}
ancestors = append(ancestors, target)
switch {
case target.object != nil:
key, ok = segment.(string)
if !ok || !utf8.ValidString(key) || len(key) > s.limit {
return fmt.Errorf("JSON object paths require valid string keys")
}
target = target.object[key]
case target.array != nil:
var number float64
switch typed := segment.(type) {
case int:
number = float64(typed)
case int64:
number = float64(typed)
case float64:
number = typed
default:
return fmt.Errorf("JSON array paths require integer indices")
}
if math.IsNaN(number) || number < 0 || number >= float64(len(target.array)) || math.Trunc(number) != number {
return fmt.Errorf("JSON array index is out of range")
}
index = int(number)
target = target.array[index]
default:
return fmt.Errorf("JSON change path traverses a scalar")
}
if target == nil && (offset != len(path)-1 || op != "set") {
return fmt.Errorf("JSON change target does not exist")
}
}
addedBytes, addedNodes := 0, 0
switch op {
case "set":
oldBytes, oldNodes := 0, 0
if target != nil {
oldBytes, oldNodes = target.bytes, target.nodes
} else {
encodedKey, err := common.Marshal(key)
if err != nil {
return err
}
addedBytes = len(encodedKey) + 1
if len(ancestors[len(ancestors)-1].object) != 0 {
addedBytes++
}
}
budget := &jsonStateBudget{ctx: ctx, bytes: s.limit - s.root.bytes + oldBytes - addedBytes, nodes: maxJSONToolNodes - s.root.nodes + oldNodes}
next, err := newJSONStateNode(value, len(path), budget)
if err != nil {
return err
}
addedBytes += next.bytes - oldBytes
addedNodes = next.nodes - oldNodes
if len(ancestors) == 0 {
s.root = next
} else if parent := ancestors[len(ancestors)-1]; parent.object != nil {
parent.object[key] = next
} else {
parent.array[index] = next
}
case "append":
if target.array == nil {
return fmt.Errorf("append requires an array target")
}
if len(target.array) != 0 {
addedBytes = 1
}
budget := &jsonStateBudget{ctx: ctx, bytes: s.limit - s.root.bytes - addedBytes, nodes: maxJSONToolNodes - s.root.nodes}
next, err := newJSONStateNode(value, len(path)+1, budget)
if err != nil {
return err
}
addedBytes += next.bytes
addedNodes = next.nodes
target.array = append(target.array, next)
target.bytes += addedBytes
target.nodes += addedNodes
case "appendText":
previous, isText := target.scalar.(string)
text, ok := value.(string)
if !isText || !ok || len(text) > s.limit-s.root.bytes {
return fmt.Errorf("appendText requires strings within the state size limit")
}
if !utf8.ValidString(text) {
text = string([]rune(text))
}
encoded, err := common.Marshal(text)
if err != nil {
return err
}
addedBytes = len(encoded) - 2
if addedBytes > s.limit-s.root.bytes {
return fmt.Errorf("JSON state exceeds size limit")
}
if text != "" {
if target.text == nil {
target.text = &strings.Builder{}
target.text.WriteString(previous)
target.scalar = ""
}
target.text.WriteString(text)
}
target.bytes += addedBytes
}
for _, ancestor := range ancestors {
ancestor.bytes += addedBytes
ancestor.nodes += addedNodes
}
}
return nil
}
// Value returns the final ordinary JSON tree and verifies the incremental byte
// accounting against the configured codec before callers can persist it.
func (s *JSONState) Value() (any, error) {
if s.failed != nil {
return nil, s.failed
}
value := s.root.value()
encoded, err := common.Marshal(value)
if err != nil {
return nil, err
}
if len(encoded) != s.root.bytes || len(encoded) > s.limit {
return nil, fmt.Errorf("JSON state encoded size does not match its bounded representation")
}
return value, nil
}
func (n *jsonStateNode) value() any {
if n.object != nil {
result := make(map[string]any, len(n.object))
for key, child := range n.object {
result[key] = child.value()
}
return result
}
if n.array != nil {
result := make([]any, len(n.array))
for index, child := range n.array {
result[index] = child.value()
}
return result
}
if n.text != nil {
return n.text.String()
}
return n.scalar
}
// Native JS containers are essential here: a lazily converted Go map would
// read as a clone, but the result must be ordinary objects a plugin owns.
func (n *jsonStateNode) jsValue(realm *engine.Realm) (engine.Value, error) {
if n.object != nil {
result := realm.NewObject()
for key, child := range n.object {
value, err := child.jsValue(realm)
if err != nil {
return engine.Undefined(), err
}
if err = result.CreateDataPropertyOrThrow(realm, realm.KeyFromGoString(key), value); err != nil {
return engine.Undefined(), err
}
}
return engine.ObjectValue(result), nil
}
if n.array != nil {
items := make([]engine.Value, len(n.array))
for index, child := range n.array {
value, err := child.jsValue(realm)
if err != nil {
return engine.Undefined(), err
}
items[index] = value
}
return engine.ObjectValue(realm.NewArrayFromSlice(items)), nil
}
return realm.FromGo(n.value())
}