fix: disable text/json/files to be property of json schema

langgenius · Dec 24, 2024 · 1a694a9 · 1a694a9
1 parent 3100864
commit 1a694a9
Show file tree

Hide file tree

Showing 2 changed files with 121 additions and 0 deletions.
diff --git a/internal/types/entities/plugin_entities/tool_declaration.go b/internal/types/entities/plugin_entities/tool_declaration.go
@@ -109,7 +109,37 @@ type ToolDeclaration struct {
 }
 
 func isJSONSchema(fl validator.FieldLevel) bool {
+	// get schema from interface
+	schemaMap, ok := fl.Field().Interface().(map[string]any)
+	if !ok {
+		return false
+	}
+
+	// validate root schema must be object type
+	rootType, ok := schemaMap["type"].(string)
+	if !ok || rootType != "object" {
+		return false
+	}
+
+	// validate properties
+	properties, ok := schemaMap["properties"].(map[string]any)
+	if !ok {
+		return false
+	}
+
+	// disallow text, json, files as property names
+	disallowedProps := []string{"text", "json", "files"}
+	for _, prop := range disallowedProps {
+		if _, exists := properties[prop]; exists {
+			return false
+		}
+	}
+
 	_, err := gojsonschema.NewSchema(gojsonschema.NewGoLoader(fl.Field().Interface()))
+	if err != nil {
+		return false
+	}
+
 	return err == nil
 }
 

diff --git a/internal/types/entities/plugin_entities/tool_declaration_test.go b/internal/types/entities/plugin_entities/tool_declaration_test.go
@@ -1175,3 +1175,94 @@ func TestParameterScope_Validate(t *testing.T) {
 		t.Errorf("ParameterScope_Validate() error = %v", err)
 	}
 }
+
+func TestInvalidJSONSchemaToolProvider_Validate(t *testing.T) {
+	type Test struct {
+		Text map[string]any `json:"text" validate:"json_schema"`
+	}
+
+	data := parser.MarshalJsonBytes(Test{
+		Text: map[string]any{
+			"text": "text",
+		},
+	})
+
+	if _, err := parser.UnmarshalJsonBytes[Test](data); err == nil {
+		t.Errorf("TestInvalidJSONSchemaToolProvider_Validate() error = %v", err)
+	}
+
+	data = parser.MarshalJsonBytes(Test{
+		Text: map[string]any{
+			"type": "object",
+			"properties": map[string]any{
+				"text": map[string]any{
+					"type": "object",
+				},
+			},
+		},
+	})
+
+	if _, err := parser.UnmarshalJsonBytes[Test](data); err == nil {
+		t.Errorf("TestInvalidJSONSchemaToolProvider_Validate() error = %v", err)
+	}
+
+	data = parser.MarshalJsonBytes(Test{
+		Text: map[string]any{
+			"type": "array",
+			"properties": map[string]any{
+				"a": map[string]any{
+					"type": "object",
+				},
+			},
+		},
+	})
+
+	if _, err := parser.UnmarshalJsonBytes[Test](data); err == nil {
+		t.Errorf("TestInvalidJSONSchemaToolProvider_Validate() error = %v", err)
+	}
+
+	data = parser.MarshalJsonBytes(Test{
+		Text: map[string]any{
+			"type": "object",
+			"properties": map[string]any{
+				"json": map[string]any{
+					"type": "object",
+				},
+			},
+		},
+	})
+
+	if _, err := parser.UnmarshalJsonBytes[Test](data); err == nil {
+		t.Errorf("TestInvalidJSONSchemaToolProvider_Validate() error = %v", err)
+	}
+
+	data = parser.MarshalJsonBytes(Test{
+		Text: map[string]any{
+			"type": "object",
+			"properties": map[string]any{
+				"files": map[string]any{
+					"type": "object",
+				},
+			},
+		},
+	})
+
+	if _, err := parser.UnmarshalJsonBytes[Test](data); err == nil {
+		t.Errorf("TestInvalidJSONSchemaToolProvider_Validate() error = %v", err)
+	}
+
+	data = parser.MarshalJsonBytes(Test{
+		Text: map[string]any{
+			"type": "object",
+			"properties": map[string]any{
+				"aaa": map[string]any{
+					"type": "object",
+				},
+			},
+		},
+	})
+
+	if _, err := parser.UnmarshalJsonBytes[Test](data); err != nil {
+		t.Errorf("TestInvalidJSONSchemaToolProvider_Validate() error = %v", err)
+	}
+}