Resolve public documentation contract questions

internal/validate/standard_validator.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io/fs"
+	"net/url"
 	"os"
 	"path"
 	"path/filepath"
@@ -16,6 +17,8 @@ import (
 	"github.com/santhosh-tekuri/jsonschema/v6"
 )
 
+const jsonSchemaDraft2020 = "https://json-schema.org/draft/2020-12/schema"
+
 // StandardValidator provides basic, JSON, and JSON Schema output validation.
 type StandardValidator struct {
 	schemaBaseDir string
@@ -121,7 +124,11 @@ func (v *StandardValidator) validateJSONSchema(instance any, schemaPath string)
 		return nil, err
 	}
 
-	compiler := jsonschema.NewCompiler()
+	schemaRoot, err := v.schemaRoot()
+	if err != nil {
+		return nil, err
+	}
+	compiler := newSchemaCompiler(standardSchemaLoader{root: schemaRoot})
 	schema, err := compiler.Compile(resolvedSchemaPath)
 	if err != nil {
 		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", resolvedSchemaPath, err)
@@ -140,7 +147,10 @@ func (v *FSValidator) validateJSONSchema(instance any, schemaPath string) ([]str
 	}
 
 	resourceURL := fsSchemaResourceURL(schemaName)
-	compiler := jsonschema.NewCompiler()
+	if err := validateSchemaDialect(schemaDoc); err != nil {
+		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", schemaName, err)
+	}
+	compiler := newSchemaCompiler(fsSchemaLoader{fsys: v.fsys, root: filecatalog.CleanFSRoot(v.root)})
 	if err := compiler.AddResource(resourceURL, schemaDoc); err != nil {
 		return nil, fmt.Errorf("failed to register JSON schema %q: %w", schemaName, err)
 	}
@@ -184,6 +194,9 @@ func (v *StandardValidator) LoadSchemaDocument(ctx context.Context, schemaPath s
 	if err := json.Unmarshal(raw, &doc); err != nil {
 		return nil, fmt.Errorf("failed to decode JSON schema %q: %w", resolved, err)
 	}
+	if err := validateSchemaDialect(doc); err != nil {
+		return nil, fmt.Errorf("failed to decode JSON schema %q: %w", resolved, err)
+	}
 	return doc, nil
 }
 
@@ -206,12 +219,14 @@ func (v *StandardValidator) resolveSchemaPath(schemaPath string) (string, error)
 		return "", errors.New("schema path is required for json_schema validation")
 	}
 
-	resolved := schemaPath
-	if !filepath.IsAbs(schemaPath) {
-		resolved = filepath.Join(v.schemaBaseDir, schemaPath)
+	root, err := v.schemaRoot()
+	if err != nil {
+		return "", err
+	}
+	resolved, err := containedFilesystemPath(root, schemaPath)
+	if err != nil {
+		return "", err
 	}
-
-	resolved = filepath.Clean(resolved)
 	if _, err := os.Stat(resolved); err != nil {
 		return "", fmt.Errorf("failed to access schema file %q: %w", resolved, err)
 	}
@@ -219,6 +234,22 @@ func (v *StandardValidator) resolveSchemaPath(schemaPath string) (string, error)
 	return resolved, nil
 }
 
+func (v *StandardValidator) schemaRoot() (string, error) {
+	root := v.schemaBaseDir
+	if strings.TrimSpace(root) == "" {
+		root = "."
+	}
+	absolute, err := filepath.Abs(root)
+	if err != nil {
+		return "", fmt.Errorf("failed to resolve schema source %q: %w", root, err)
+	}
+	resolved, err := filepath.EvalSymlinks(absolute)
+	if err != nil {
+		return "", fmt.Errorf("failed to access schema source %q: %w", root, err)
+	}
+	return resolved, nil
+}
+
 func (v *FSValidator) loadSchemaDocument(schemaPath string) (string, any, error) {
 	resolved, err := v.resolveSchemaPath(schemaPath)
 	if err != nil {
@@ -234,6 +265,9 @@ func (v *FSValidator) loadSchemaDocument(schemaPath string) (string, any, error)
 	if err := json.Unmarshal(raw, &doc); err != nil {
 		return "", nil, fmt.Errorf("failed to decode JSON schema %q: %w", resolved, err)
 	}
+	if err := validateSchemaDialect(doc); err != nil {
+		return "", nil, fmt.Errorf("failed to decode JSON schema %q: %w", resolved, err)
+	}
 	return resolved, doc, nil
 }
 
@@ -290,3 +324,131 @@ func cleanSchemaFSPath(schemaPath string) (string, error) {
 func fsSchemaResourceURL(schemaName string) string {
 	return "promptkit-schema:///" + strings.TrimPrefix(path.Clean(schemaName), "/")
 }
+
+func newSchemaCompiler(loader jsonschema.URLLoader) *jsonschema.Compiler {
+	compiler := jsonschema.NewCompiler()
+	compiler.DefaultDraft(jsonschema.Draft2020)
+	compiler.UseLoader(loader)
+	return compiler
+}
+
+func validateSchemaDialect(doc any) error {
+	object, ok := doc.(map[string]any)
+	if !ok {
+		return nil
+	}
+	value, ok := object["$schema"]
+	if !ok {
+		return nil
+	}
+	dialect, ok := value.(string)
+	if !ok {
+		return errors.New("$schema must be a string")
+	}
+	if dialect != jsonSchemaDraft2020 && dialect != jsonSchemaDraft2020+"#" {
+		return fmt.Errorf("unsupported JSON Schema dialect %q; expected %q", dialect, jsonSchemaDraft2020)
+	}
+	return nil
+}
+
+type standardSchemaLoader struct {
+	root string
+}
+
+func (l standardSchemaLoader) Load(resourceURL string) (any, error) {
+	fileName, err := (jsonschema.FileLoader{}).ToFile(resourceURL)
+	if err != nil {
+		return nil, fmt.Errorf("schema reference %q is not a contained file reference: %w", resourceURL, err)
+	}
+	resolved, err := containedFilesystemPath(l.root, fileName)
+	if err != nil {
+		return nil, err
+	}
+	return loadJSONSchemaFile(resolved)
+}
+
+func containedFilesystemPath(root, name string) (string, error) {
+	candidate := name
+	if !filepath.IsAbs(candidate) {
+		candidate = filepath.Join(root, candidate)
+	}
+	candidate, err := filepath.Abs(candidate)
+	if err != nil {
+		return "", fmt.Errorf("failed to resolve schema path %q: %w", name, err)
+	}
+	candidate, err = filepath.EvalSymlinks(candidate)
+	if err != nil {
+		return "", fmt.Errorf("failed to access schema file %q: %w", candidate, err)
+	}
+	relative, err := filepath.Rel(root, candidate)
+	if err != nil {
+		return "", fmt.Errorf("failed to compare schema path %q with source root: %w", candidate, err)
+	}
+	if relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
+		return "", fmt.Errorf("schema path %q escapes source root", name)
+	}
+	return candidate, nil
+}
+
+func loadJSONSchemaFile(name string) (any, error) {
+	raw, err := os.ReadFile(name)
+	if err != nil {
+		return nil, err
+	}
+	var doc any
+	if err := json.Unmarshal(raw, &doc); err != nil {
+		return nil, err
+	}
+	if err := validateSchemaDialect(doc); err != nil {
+		return nil, err
+	}
+	return doc, nil
+}
+
+type fsSchemaLoader struct {
+	fsys fs.FS
+	root string
+}
+
+func (l fsSchemaLoader) Load(resourceURL string) (any, error) {
+	parsed, err := url.Parse(resourceURL)
+	if err != nil {
+		return nil, fmt.Errorf("invalid schema reference %q: %w", resourceURL, err)
+	}
+	if parsed.Scheme != "promptkit-schema" || parsed.Host != "" {
+		return nil, fmt.Errorf("schema reference %q is not allowed", resourceURL)
+	}
+	name, err := url.PathUnescape(strings.TrimPrefix(parsed.Path, "/"))
+	if err != nil {
+		return nil, fmt.Errorf("invalid schema reference %q: %w", resourceURL, err)
+	}
+	name = path.Clean(name)
+	if l.root == "." {
+		if strings.HasPrefix(name, "../") || name == ".." {
+			return nil, fmt.Errorf("schema reference %q escapes source root", resourceURL)
+		}
+	} else if name != l.root && !strings.HasPrefix(name, l.root+"/") {
+		return nil, fmt.Errorf("schema reference %q escapes source root", resourceURL)
+	}
+
+	rootInfo, err := fs.Stat(l.fsys, l.root)
+	if err != nil {
+		return nil, err
+	}
+	if !rootInfo.IsDir() && name != l.root {
+		return nil, fmt.Errorf("schema reference %q is outside the configured schema file", resourceURL)
+	}
+
+	raw, err := fs.ReadFile(l.fsys, name)
+	if err != nil {
+		return nil, err
+	}
+	var doc any
+	if err := json.Unmarshal(raw, &doc); err != nil {
+		return nil, err
+	}
+	if err := validateSchemaDialect(doc); err != nil {
+		return nil, err
+	}
+	return doc, nil
+}