449 lines
12 KiB
Go
449 lines
12 KiB
Go
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
|
|
}
|
|
|
|
// 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())
|
|
}
|