internal/mcp/jsonschema: check for structure separately

Redo iteration. It was too subtle.

Write a separate recursive check for whether the schemas form a tree and
have any invalid nils.

While doing that, assign paths to schemas for better error messages,
even as we check structure. This also lets us detect re-occurring
schemas without a separate map.

Everything else, including the iterators, can now assume valid structure
and populated paths.

Change-Id: I90ec93d79d1d95c92258ad7afc94ef068fdf123f
Reviewed-on: https://go-review.googlesource.com/c/tools/+/676376
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
Auto-Submit: Jonathan Amsterdam <jba@google.com>

jsonschema/resolve.go

@@ -11,6 +11,7 @@ import (
 	"errors"
 	"fmt"
 	"net/url"
+	"reflect"
 	"regexp"
 	"strings"
 )
@@ -111,32 +112,83 @@ func (r *resolver) resolve(s *Schema, baseURI *url.URL) (*Resolved, error) {
 }
 
 func (root *Schema) check() error {
-	if root == nil {
+	// Check for structural validity. Do this first and fail fast:
-		return errors.New("nil schema")
+	// bad structure will cause other code to panic.
+	if err := root.checkStructure(); err != nil {
+		return err
 	}
+
 	var errs []error
 	report := func(err error) { errs = append(errs, err) }
 
-	seen := map[*Schema]bool{}
+	for ss := range root.all() {
-	for ss, path := range root.all() {
-		if seen[ss] {
-			// The schema graph rooted at s is not a tree, but it needs to
-			// be because we assume a unique parent when we store a schema's base
-			// in the Schema. A cycle would also put Schema.all into an infinite
-			// recursion.
-			return fmt.Errorf("schemas rooted at %s do not form a tree (saw %s twice)", root, ss)
-		}
-		seen[ss] = true
-		if len(path) == 0 {
-			ss.path = "root"
-		} else {
-			ss.path = "/" + strings.Join(path, "/")
-		}
 		ss.checkLocal(report)
 	}
 	return errors.Join(errs...)
 }
 
+// checkStructure verifies that root and its subschemas form a tree.
+// It also assigns each schema a unique path, to improve error messages.
+func (root *Schema) checkStructure() error {
+	var check func(reflect.Value, []byte) error
+	check = func(v reflect.Value, path []byte) error {
+		// For the purpose of error messages, the root schema has path "root"
+		// and other schemas' paths are their JSON Pointer from the root.
+		p := "root"
+		if len(path) > 0 {
+			p = string(path)
+		}
+		s := v.Interface().(*Schema)
+		if s == nil {
+			return fmt.Errorf("jsonschema: schema at %s is nil", p)
+		}
+		if s.path != "" {
+			// We've seen s before.
+			// The schema graph at root is not a tree, but it needs to
+			// be because we assume a unique parent when we store a schema's base
+			// in the Schema. A cycle would also put Schema.all into an infinite
+			// recursion.
+			return fmt.Errorf("jsonschema: schemas at %s do not form a tree; %s appears more than once (also at %s)",
+				root, s.path, p)
+		}
+		s.path = p
+
+		for _, info := range schemaFieldInfos {
+			fv := v.Elem().FieldByIndex(info.sf.Index)
+			switch info.sf.Type {
+			case schemaType:
+				// A field that contains an individual schema.
+				// A nil is valid: it just means the field isn't present.
+				if !fv.IsNil() {
+					if err := check(fv, fmt.Appendf(path, "/%s", info.jsonName)); err != nil {
+						return err
+					}
+				}
+
+			case schemaSliceType:
+				for i := range fv.Len() {
+					if err := check(fv.Index(i), fmt.Appendf(path, "/%s/%d", info.jsonName, i)); err != nil {
+						return err
+					}
+				}
+
+			case schemaMapType:
+				iter := fv.MapRange()
+				for iter.Next() {
+					key := escapeJSONPointerSegment(iter.Key().String())
+					if err := check(iter.Value(), fmt.Appendf(path, "/%s/%s", info.jsonName, key)); err != nil {
+						return err
+					}
+				}
+			}
+
+		}
+		return nil
+	}
+
+	return check(reflect.ValueOf(root), make([]byte, 0, 256))
+}
+
 // checkLocal checks s for validity, independently of other schemas it may refer to.
 // Since checking a regexp involves compiling it, checkLocal saves those compiled regexps
 // in the schema for later use.

jsonschema/resolve_test.go

@@ -14,18 +14,41 @@ import (
 	"testing"
 )
 
+func TestSchemaStructure(t *testing.T) {
+	check := func(s *Schema, want string) {
+		t.Helper()
+		err := s.checkStructure()
+		if err == nil || !strings.Contains(err.Error(), want) {
+			t.Errorf("checkStructure returned error %q, want %q", err, want)
+		}
+	}
+
+	dag := &Schema{Type: "number"}
+	dag = &Schema{Items: dag, Contains: dag}
+	check(dag, "do not form a tree")
+
+	tree := &Schema{Type: "number"}
+	tree.Items = tree
+	check(tree, "do not form a tree")
+
+	sliceNil := &Schema{PrefixItems: []*Schema{nil}}
+	check(sliceNil, "is nil")
+
+	sliceMap := &Schema{Properties: map[string]*Schema{"a": nil}}
+	check(sliceMap, "is nil")
+}
+
 func TestCheckLocal(t *testing.T) {
 	for _, tt := range []struct {
 		s    *Schema
 		want string // error must be non-nil and match this regexp
 	}{
-		{nil, "nil"},
 		{
 			&Schema{Pattern: "]["},
 			"regexp",
 		},
 		{
-			&Schema{PatternProperties: map[string]*Schema{"*": nil}},
+			&Schema{PatternProperties: map[string]*Schema{"*": {}}},
 			"regexp",
 		},
 	} {
@@ -35,26 +58,48 @@ func TestCheckLocal(t *testing.T) {
 			continue
 		}
 		if !regexp.MustCompile(tt.want).MatchString(err.Error()) {
-			t.Errorf("%s: did not match\nerror: %s\nregexp: %s",
+			t.Errorf("checkLocal returned error\n%q\nwanted it to match\n%s\nregexp: %s",
 				tt.s.json(), err, tt.want)
 		}
 	}
 }
 
-func TestSchemaNonTree(t *testing.T) {
+func TestPaths(t *testing.T) {
-	run := func(s *Schema, kind string) {
+	// CheckStructure should assign paths to schemas.
-		err := s.check()
+	// This test also verifies that Schema.all visits maps in sorted order.
-		if err == nil || !strings.Contains(err.Error(), "tree") {
+	root := &Schema{
-			t.Fatalf("did not detect %s", kind)
+		Type:        "string",
-		}
+		PrefixItems: []*Schema{{Type: "int"}, {Items: &Schema{Type: "null"}}},
+		Contains: &Schema{Properties: map[string]*Schema{
+			"~1": {Type: "boolean"},
+			"p":  {},
+		}},
 	}
 
-	s := &Schema{Type: "number"}
+	type item struct {
-	run(&Schema{Items: s, Contains: s}, "DAG")
+		s *Schema
+		p string
+	}
+	want := []item{
+		{root, "root"},
+		{root.Contains, "/contains"},
+		{root.Contains.Properties["p"], "/contains/properties/p"},
+		{root.Contains.Properties["~1"], "/contains/properties/~01"},
+		{root.PrefixItems[0], "/prefixItems/0"},
+		{root.PrefixItems[1], "/prefixItems/1"},
+		{root.PrefixItems[1].Items, "/prefixItems/1/items"},
+	}
+	if err := root.checkStructure(); err != nil {
+		t.Fatal(err)
+	}
 
-	root := &Schema{Items: s}
+	var got []item
-	s.Items = root
+	for s := range root.all() {
-	run(root, "cycle")
+		got = append(got, item{s, s.path})
+	}
+	if !slices.Equal(got, want) {
+		t.Errorf("\ngot  %v\nwant %v", got, want)
+	}
 }
 
 func TestResolveURIs(t *testing.T) {

jsonschema/schema.go

@@ -17,7 +17,6 @@ import (
 	"reflect"
 	"regexp"
 	"slices"
-	"strconv"
 )
 
 // A Schema is a JSON schema object.
@@ -361,27 +360,12 @@ func Ptr[T any](x T) *T { return &x }
 // every applies f preorder to every schema under s including s.
 // The second argument to f is the path to the schema appended to the argument path.
 // It stops when f returns false.
-func (s *Schema) every(f func(*Schema, []string) bool, path []string) bool {
+func (s *Schema) every(f func(*Schema) bool) bool {
-	return s == nil ||
+	return f(s) && s.everyChild(func(s *Schema) bool { return s.every(f) })
-		f(s, path) && s.everyChild(func(s *Schema, p []string) bool { return s.every(f, p) }, path)
 }
 
 // everyChild reports whether f is true for every immediate child schema of s.
-// The second argument to f is the path to the schema appended to the argument path.
+func (s *Schema) everyChild(f func(*Schema) bool) bool {
-//
-// It does not call f on nil-valued fields holding individual schemas, like Contains,
-// because a nil value indicates that the field is absent.
-// It does call f on nils when they occur in slices and maps, so those invalid values
-// can be detected when the schema is validated.
-func (s *Schema) everyChild(f func(*Schema, []string) bool, path []string) bool {
-	if s == nil {
-		return false
-	}
-	var (
-		schemaType      = reflect.TypeFor[*Schema]()
-		schemaSliceType = reflect.TypeFor[[]*Schema]()
-		schemaMapType   = reflect.TypeFor[map[string]*Schema]()
-	)
 	v := reflect.ValueOf(s)
 	for _, info := range schemaFieldInfos {
 		fv := v.Elem().FieldByIndex(info.sf.Index)
@@ -389,25 +373,23 @@ func (s *Schema) everyChild(f func(*Schema, []string) bool, path []string) bool
 		case schemaType:
 			// A field that contains an individual schema. A nil is valid: it just means the field isn't present.
 			c := fv.Interface().(*Schema)
-			if c != nil && !f(c, append(path, info.jsonName)) {
+			if c != nil && !f(c) {
 				return false
 			}
 
 		case schemaSliceType:
 			slice := fv.Interface().([]*Schema)
-			// A field that contains a slice of schemas. Yield nils so we can check for their presence.
+			for _, c := range slice {
-			for i, c := range slice {
+				if !f(c) {
-				if !f(c, append(path, info.jsonName, strconv.Itoa(i))) {
 					return false
 				}
 			}
 
 		case schemaMapType:
-			// A field that is a map of schemas. Ditto about nils.
 			// Sort keys for determinism.
 			m := fv.Interface().(map[string]*Schema)
 			for _, k := range slices.Sorted(maps.Keys(m)) {
-				if !f(m[k], append(path, info.jsonName, escapeJSONPointerSegment(k))) {
+				if !f(m[k]) {
 					return false
 				}
 			}
@@ -417,16 +399,21 @@ func (s *Schema) everyChild(f func(*Schema, []string) bool, path []string) bool
 }
 
 // all wraps every in an iterator.
-func (s *Schema) all() iter.Seq2[*Schema, []string] {
+func (s *Schema) all() iter.Seq[*Schema] {
-	return func(yield func(*Schema, []string) bool) { s.every(yield, nil) }
+	return func(yield func(*Schema) bool) { s.every(yield) }
 }
 
 // children wraps everyChild in an iterator.
-func (s *Schema) children() iter.Seq2[*Schema, []string] {
+func (s *Schema) children() iter.Seq[*Schema] {
-	var pathBuffer [4]string
+	return func(yield func(*Schema) bool) { s.everyChild(yield) }
-	return func(yield func(*Schema, []string) bool) { s.everyChild(yield, pathBuffer[:0]) }
 }
 
+var (
+	schemaType      = reflect.TypeFor[*Schema]()
+	schemaSliceType = reflect.TypeFor[[]*Schema]()
+	schemaMapType   = reflect.TypeFor[map[string]*Schema]()
+)
+
 type structFieldInfo struct {
 	sf       reflect.StructField
 	jsonName string

jsonschema/schema_test.go

@@ -9,8 +9,6 @@ import (
 	"fmt"
 	"math"
 	"regexp"
-	"slices"
-	"strings"
 	"testing"
 )
 
@@ -111,37 +109,6 @@ func TestUnmarshalErrors(t *testing.T) {
 	}
 }
 
-func TestEvery(t *testing.T) {
-	// Schema.every should visit all descendants of a schema, not just the immediate ones.
-	s := &Schema{
-		Type:        "string",
-		PrefixItems: []*Schema{{Type: "int"}, {Items: &Schema{Type: "null"}}},
-		Contains: &Schema{Properties: map[string]*Schema{
-			"~1": {Type: "boolean"},
-			"p":  {},
-		}},
-	}
-
-	type item struct {
-		s *Schema
-		p string
-	}
-	want := []item{
-		{s, ""},
-		{s.Contains, "contains"},
-		{s.Contains.Properties["p"], "contains/properties/p"},
-		{s.Contains.Properties["~1"], "contains/properties/~01"},
-		{s.PrefixItems[0], "prefixItems/0"},
-		{s.PrefixItems[1], "prefixItems/1"},
-		{s.PrefixItems[1].Items, "prefixItems/1/items"},
-	}
-	var got []item
-	s.every(func(s *Schema, p []string) bool { got = append(got, item{s, strings.Join(p, "/")}); return true }, nil)
-	if !slices.Equal(got, want) {
-		t.Errorf("\n got  %v\nwant %v", got, want)
-	}
-}
-
 func mustUnmarshal(t *testing.T, data []byte, ptr any) {
 	t.Helper()
 	if err := json.Unmarshal(data, ptr); err != nil {

jsonschema/validate.go

@@ -32,8 +32,7 @@ func (rs *Resolved) Validate(instance any) error {
 
 // state is the state of single call to ResolvedSchema.Validate.
 type state struct {
-	rs    *Resolved
+	rs *Resolved
-	depth int
 	// stack holds the schemas from recursive calls to validate.
 	// These are the "dynamic scopes" used to resolve dynamic references.
 	// https://json-schema.org/draft/2020-12/json-schema-core#scopes
@@ -48,9 +47,6 @@ func (st *state) validate(instance reflect.Value, schema *Schema, callerAnns *an
 	defer func() {
 		st.stack = st.stack[:len(st.stack)-1] // pop
 	}()
-	if depth := len(st.stack); depth >= 100 {
-		return fmt.Errorf("max recursion depth of %d reached", depth)
-	}
 
 	// We checked for nil schemas in [Schema.Resolve].
 	assert(schema != nil, "nil schema")