From bda18b44d9e4f6b87f81d47299f0e51cfcb187df Mon Sep 17 00:00:00 2001
From: "Daniel G. Taylor" <danielgtaylor@gmail.com>
Date: Sun, 22 Oct 2023 22:27:06 -0700
Subject: [PATCH] fix: add $schema field to input to allow round-trips

---
 huma_test.go  | 21 +++++++++++++++++++++
 transforms.go | 40 ++++++++++++++++++++++++++--------------
 2 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/huma_test.go b/huma_test.go
index 6bb6f3d5..79e9c0c1 100644
--- a/huma_test.go
+++ b/huma_test.go
@@ -480,6 +480,27 @@ func TestFeatures(t *testing.T) {
 				assert.Equal(t, 256, resp.Code)
 			},
 		},
+		{
+			// Simulate a request with a body that came from another call, which
+			// includes the `$schema` field. It should be allowed to be passed
+			// to the new operation as input without modification.
+			Name: "round-trip-schema-field",
+			Register: func(t *testing.T, api API) {
+				Register(api, Operation{
+					Method: http.MethodPut,
+					Path:   "/round-trip",
+				}, func(ctx context.Context, input *struct {
+					Body struct {
+						Name string `json:"name"`
+					}
+				}) (*struct{}, error) {
+					return nil, nil
+				})
+			},
+			Method: http.MethodPut,
+			URL:    "/round-trip",
+			Body:   `{"$schema": "...", "name": "foo"}`,
+		},
 		{
 			Name: "one-of input",
 			Register: func(t *testing.T, api API) {
diff --git a/transforms.go b/transforms.go
index e4d329da..44355345 100644
--- a/transforms.go
+++ b/transforms.go
@@ -47,32 +47,44 @@ func NewSchemaLinkTransformer(prefix, schemasPath string) *SchemaLinkTransformer
 	}
 }
 
+func (t *SchemaLinkTransformer) addSchemaField(oapi *OpenAPI, content *MediaType) bool {
+	if content == nil || content.Schema == nil || content.Schema.Ref == "" {
+		return true
+	}
+
+	schema := oapi.Components.Schemas.SchemaFromRef(content.Schema.Ref)
+	if schema.Type != TypeObject || (schema.Properties != nil && schema.Properties["$schema"] != nil) {
+		return true
+	}
+
+	schema.Properties["$schema"] = &Schema{
+		Type:        TypeString,
+		Format:      "uri",
+		Description: "A URL to the JSON Schema for this object.",
+		ReadOnly:    true,
+	}
+	return false
+}
+
 // OnAddOperation is triggered whenever a new operation is added to the API,
 // enabling this transformer to precompute & cache information about the
 // response and schema.
 func (t *SchemaLinkTransformer) OnAddOperation(oapi *OpenAPI, op *Operation) {
 	// Update registry to be able to get the type from a schema ref.
 	// Register the type in t.types with the generated ref
+	if op.RequestBody != nil && op.RequestBody.Content != nil {
+		for _, content := range op.RequestBody.Content {
+			t.addSchemaField(oapi, content)
+		}
+	}
+
 	registry := oapi.Components.Schemas
 	for _, resp := range op.Responses {
 		for _, content := range resp.Content {
-			if content == nil || content.Schema == nil || content.Schema.Ref == "" {
+			if t.addSchemaField(oapi, content) {
 				continue
 			}
 
-			schema := registry.SchemaFromRef(content.Schema.Ref)
-			if schema.Type != TypeObject || (schema.Properties != nil && schema.Properties["$schema"] != nil) {
-				continue
-			}
-
-			// First, modify the schema to have the $schema field.
-			schema.Properties["$schema"] = &Schema{
-				Type:        TypeString,
-				Format:      "uri",
-				Description: "A URL to the JSON Schema for this object.",
-				ReadOnly:    true,
-			}
-
 			// Then, create the wrapper Go type that has the $schema field.
 			typ := deref(registry.TypeFromRef(content.Schema.Ref))