From 915059de2a761da307caecc2ddd6b4997ba37e5d Mon Sep 17 00:00:00 2001
From: haoqixu <hq.xu0o0@gmail.com>
Date: Fri, 9 Aug 2024 12:47:18 +0800
Subject: [PATCH] encoding/jsonschema: fix decoding of `$id`

When decoding "$id", jsonschema decoder mistakenly assumed that
the schema is an object and added an object constraint with the id tag.

This change postpones the attachment of id tag until an object
constraint is added or the schema is finalizing.

Fixes #3361

Change-Id: I7829af165af726524d48022fc87ed538b7e451cd
Signed-off-by: haoqixu <hq.xu0o0@gmail.com>
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/1199139
Unity-Result: CUE porcuepine <cue.porcuepine@gmail.com>
Reviewed-by: Roger Peppe <rogpeppe@gmail.com>
TryBot-Result: CUEcueckoo <cueckoo@cuelang.org>
---
 encoding/jsonschema/constraints.go            |  9 +-----
 encoding/jsonschema/decode.go                 | 25 ++++++++++++++++-
 .../jsonschema/testdata/id_in_oneOf.txtar     | 28 +++++++++++++++++++
 encoding/jsonschema/testdata/id_scalar.txtar  | 12 ++++++++
 4 files changed, 65 insertions(+), 9 deletions(-)
 create mode 100644 encoding/jsonschema/testdata/id_in_oneOf.txtar
 create mode 100644 encoding/jsonschema/testdata/id_scalar.txtar

diff --git a/encoding/jsonschema/constraints.go b/encoding/jsonschema/constraints.go
index 8296efff477..1abd67f90b3 100644
--- a/encoding/jsonschema/constraints.go
+++ b/encoding/jsonschema/constraints.go
@@ -15,7 +15,6 @@
 package jsonschema
 
 import (
-	"fmt"
 	"math/big"
 	"path"
 	"regexp"
@@ -143,13 +142,7 @@ var constraints = []*constraint{
 			return
 		}
 		s.id = u
-
-		obj := s.object(n)
-
-		// TODO: handle the case where this is always defined and we don't want
-		// to include the default value.
-		obj.Elts = append(obj.Elts, &ast.Attribute{
-			Text: fmt.Sprintf("@jsonschema(id=%q)", u)})
+		s.idPos = n.Pos()
 	}),
 
 	// Generic constraint
diff --git a/encoding/jsonschema/decode.go b/encoding/jsonschema/decode.go
index 82cba4f11d2..0582e18a2e9 100644
--- a/encoding/jsonschema/decode.go
+++ b/encoding/jsonschema/decode.go
@@ -336,6 +336,7 @@ type state struct {
 	exclusiveMax bool // For OpenAPI and legacy support.
 	jsonschema   string
 	id           *url.URL // base URI for $ref
+	idPos        token.Pos
 
 	definitions []ast.Decl
 
@@ -362,9 +363,19 @@ type refs struct {
 	refs  []*ast.Ident
 }
 
+func (s *state) idTag() *ast.Attribute {
+	return &ast.Attribute{
+		At:   s.idPos,
+		Text: fmt.Sprintf("@jsonschema(id=%q)", s.id)}
+}
+
 func (s *state) object(n cue.Value) *ast.StructLit {
 	if s.obj == nil {
 		s.obj = &ast.StructLit{}
+
+		if s.id != nil {
+			s.obj.Elts = append(s.obj.Elts, s.idTag())
+		}
 		s.add(n, objectType, s.obj)
 	}
 	return s.obj
@@ -382,7 +393,8 @@ func (s *state) hasConstraints() bool {
 	return len(s.patterns) > 0 ||
 		s.title != "" ||
 		s.description != "" ||
-		s.obj != nil
+		s.obj != nil ||
+		s.id != nil
 }
 
 const allTypes = cue.NullKind | cue.BoolKind | cue.NumberKind | cue.IntKind |
@@ -499,6 +511,17 @@ outer:
 
 	s.linkReferences()
 
+	// If an "$id" exists and has not been included in any object constraints
+	if s.id != nil && s.obj == nil {
+		if st, ok := e.(*ast.StructLit); ok {
+			st.Elts = append([]ast.Decl{s.idTag()}, st.Elts...)
+		} else {
+			st = &ast.StructLit{Elts: []ast.Decl{s.idTag()}}
+			st.Elts = append(st.Elts, &ast.EmbedDecl{Expr: e})
+			e = st
+		}
+	}
+
 	return e
 }
 
diff --git a/encoding/jsonschema/testdata/id_in_oneOf.txtar b/encoding/jsonschema/testdata/id_in_oneOf.txtar
new file mode 100644
index 00000000000..62fd52636e3
--- /dev/null
+++ b/encoding/jsonschema/testdata/id_in_oneOf.txtar
@@ -0,0 +1,28 @@
+-- schema.json --
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "$id": "https://test.example/foo",
+  "oneOf": [
+      {
+          "$schema": "http://json-schema.org/draft-07/schema#",
+          "$id": "https://1.test.example/string",
+          "type": "string"
+      },
+      {
+          "$schema": "http://json-schema.org/draft-07/schema#",
+          "$id": "https://2.test.example/object",
+          "type": "object"
+      }
+  ]
+}
+
+-- out/decode/cue --
+@jsonschema(schema="http://json-schema.org/draft-07/schema#")
+@jsonschema(id="https://test.example/foo")
+{
+	@jsonschema(id="https://1.test.example/string")
+	string
+} | {
+	@jsonschema(id="https://2.test.example/object")
+	...
+}
diff --git a/encoding/jsonschema/testdata/id_scalar.txtar b/encoding/jsonschema/testdata/id_scalar.txtar
new file mode 100644
index 00000000000..ff86e853a2e
--- /dev/null
+++ b/encoding/jsonschema/testdata/id_scalar.txtar
@@ -0,0 +1,12 @@
+-- schema.json --
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "$id": "http://cuelang.org/go/encoding/openapi/testdata/id_scalar.json",
+
+  "type": "string"
+}
+
+-- out/decode/cue --
+@jsonschema(schema="http://json-schema.org/draft-07/schema#")
+@jsonschema(id="http://cuelang.org/go/encoding/openapi/testdata/id_scalar.json")
+string