From 7e9164e7b9616598ef15e5df89d0b5a446544b00 Mon Sep 17 00:00:00 2001
From: Nicholas Bollweg <nick.bollweg@gmail.com>
Date: Fri, 17 May 2024 09:29:28 -0500
Subject: [PATCH] add basic support for declaring schema with inline `$schema`

---
 src/languageservice/services/dollarUtils.ts   | 26 +++++++++
 .../services/yamlSchemaService.ts             | 58 ++++++++++++-------
 2 files changed, 64 insertions(+), 20 deletions(-)
 create mode 100644 src/languageservice/services/dollarUtils.ts

diff --git a/src/languageservice/services/dollarUtils.ts b/src/languageservice/services/dollarUtils.ts
new file mode 100644
index 00000000..3459e87d
--- /dev/null
+++ b/src/languageservice/services/dollarUtils.ts
@@ -0,0 +1,26 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Red Hat, Inc. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { SingleYAMLDocument } from '../parser/yamlParser07';
+import { JSONDocument } from '../parser/jsonParser07';
+
+/**
+ * Retrieve schema if declared by `$schema`.
+ * Public for testing purpose, not part of the API.
+ * @param doc
+ */
+export function getDollarSchema(doc: SingleYAMLDocument | JSONDocument): string | undefined {
+  if (doc instanceof SingleYAMLDocument && doc.root.type === 'object') {
+    let dollarSchema = doc.root.properties['$schema'];
+    dollarSchema = typeof dollarSchema === 'string' ? dollarSchema.trim() : undefined;
+    if (typeof dollarSchema === 'string') {
+      return dollarSchema.trim();
+    }
+    if (dollarSchema) {
+      console.log('The $schema attribute is not a string, and will be ignored');
+    }
+  }
+  return undefined;
+}
diff --git a/src/languageservice/services/yamlSchemaService.ts b/src/languageservice/services/yamlSchemaService.ts
index 453b74eb..8873be61 100644
--- a/src/languageservice/services/yamlSchemaService.ts
+++ b/src/languageservice/services/yamlSchemaService.ts
@@ -29,6 +29,7 @@ import { SchemaVersions } from '../yamlTypes';
 
 import Ajv, { DefinedError } from 'ajv';
 import { getSchemaTitle } from '../utils/schemaUtils';
+import { getDollarSchema } from './dollarUitls';
 
 const localize = nls.loadMessageBundle();
 
@@ -343,33 +344,46 @@ export class YAMLSchemaService extends JSONSchemaService {
   }
 
   public getSchemaForResource(resource: string, doc: JSONDocument): Promise<ResolvedSchema> {
+    const normalizeSchemaRef = (schemaRef: string): string | undefined => {
+      if (!schemaRef.startsWith('file:') && !schemaRef.startsWith('http')) {
+        // If path contains a fragment and it is left intact, "#" will be
+        // considered part of the filename and converted to "%23" by
+        // path.resolve() -> take it out and add back after path.resolve
+        let appendix = '';
+        if (schemaRef.indexOf('#') > 0) {
+          const segments = schemaRef.split('#', 2);
+          schemaRef = segments[0];
+          appendix = segments[1];
+        }
+        if (!path.isAbsolute(schemaRef)) {
+          const resUri = URI.parse(resource);
+          schemaRef = URI.file(path.resolve(path.parse(resUri.fsPath).dir, schemaRef)).toString();
+        } else {
+          schemaRef = URI.file(schemaRef).toString();
+        }
+        if (appendix.length > 0) {
+          schemaRef += '#' + appendix;
+        }
+      }
+      return schemaRef;
+    };
+
     const resolveModelineSchema = (): string | undefined => {
       let schemaFromModeline = getSchemaFromModeline(doc);
       if (schemaFromModeline !== undefined) {
-        if (!schemaFromModeline.startsWith('file:') && !schemaFromModeline.startsWith('http')) {
-          // If path contains a fragment and it is left intact, "#" will be
-          // considered part of the filename and converted to "%23" by
-          // path.resolve() -> take it out and add back after path.resolve
-          let appendix = '';
-          if (schemaFromModeline.indexOf('#') > 0) {
-            const segments = schemaFromModeline.split('#', 2);
-            schemaFromModeline = segments[0];
-            appendix = segments[1];
-          }
-          if (!path.isAbsolute(schemaFromModeline)) {
-            const resUri = URI.parse(resource);
-            schemaFromModeline = URI.file(path.resolve(path.parse(resUri.fsPath).dir, schemaFromModeline)).toString();
-          } else {
-            schemaFromModeline = URI.file(schemaFromModeline).toString();
-          }
-          if (appendix.length > 0) {
-            schemaFromModeline += '#' + appendix;
-          }
-        }
+        schemaFromModeline = normalizeSchemaRef(schemaFromModeline);
         return schemaFromModeline;
       }
     };
 
+    const resolveDollarSchema = (): string | undefined => {
+      let dollarSchema = getDollarSchema(doc);
+      if (dollarSchema !== undefined) {
+        dollarSchema = normalizeSchemaRef(dollarSchema);
+        return dollarSchema;
+      }
+    };
+
     const resolveSchemaForResource = (schemas: string[]): Promise<ResolvedSchema> => {
       const schemaHandle = super.createCombinedSchema(resource, schemas);
       return schemaHandle.getResolvedSchema().then((schema) => {
@@ -416,6 +430,10 @@ export class YAMLSchemaService extends JSONSchemaService {
     if (modelineSchema) {
       return resolveSchemaForResource([modelineSchema]);
     }
+    const dollarSchema = resolveDollarSchema();
+    if (dollarSchema) {
+      return resolveSchemaForResource([dollarSchema]);
+    }
     if (this.customSchemaProvider) {
       return this.customSchemaProvider(resource)
         .then((schemaUri) => {