From 8056e5a472613654200cd669962efa3043060967 Mon Sep 17 00:00:00 2001
From: Jonathan Amsterdam <jba@google.com>
Date: Wed, 7 May 2025 21:11:11 -0400
Subject: [PATCH] internal/mcp/jsonschema: check that schemas form a tree

When we check the root schema during resolution, first make sure
its descendants form a tree.

Change-Id: Icc430e7df463aa84b961c28616be567f98960430
Reviewed-on: https://go-review.googlesource.com/c/tools/+/670776
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
---
 jsonschema/resolve.go      | 15 ++++++++++++---
 jsonschema/resolve_test.go | 17 +++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/jsonschema/resolve.go b/jsonschema/resolve.go
index 5cc1cc6..9d3a064 100644
--- a/jsonschema/resolve.go
+++ b/jsonschema/resolve.go
@@ -161,14 +161,23 @@ func (r *resolver) resolveRefs(rs *Resolved) error {
 	return nil
 }
 
-func (s *Schema) check() error {
-	if s == nil {
+func (root *Schema) check() error {
+	if root == nil {
 		return errors.New("nil schema")
 	}
 	var errs []error
 	report := func(err error) { errs = append(errs, err) }
 
-	for ss := range s.all() {
+	seen := map[*Schema]bool{}
+	for ss := 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
 		ss.checkLocal(report)
 	}
 	return errors.Join(errs...)
diff --git a/jsonschema/resolve_test.go b/jsonschema/resolve_test.go
index 54f6074..610d8de 100644
--- a/jsonschema/resolve_test.go
+++ b/jsonschema/resolve_test.go
@@ -10,6 +10,7 @@ import (
 	"net/url"
 	"regexp"
 	"slices"
+	"strings"
 	"testing"
 )
 
@@ -40,6 +41,22 @@ func TestCheckLocal(t *testing.T) {
 	}
 }
 
+func TestSchemaNonTree(t *testing.T) {
+	run := func(s *Schema, kind string) {
+		err := s.check()
+		if err == nil || !strings.Contains(err.Error(), "tree") {
+			t.Fatalf("did not detect %s", kind)
+		}
+	}
+
+	s := &Schema{Type: "number"}
+	run(&Schema{Items: s, Contains: s}, "DAG")
+
+	root := &Schema{Items: s}
+	s.Items = root
+	run(root, "cycle")
+}
+
 func TestResolveURIs(t *testing.T) {
 	for _, baseURI := range []string{"", "http://a.com"} {
 		t.Run(baseURI, func(t *testing.T) {