feat(tree): account for resolveInlineRef behavior when $ref has siblings (#25)

* fix(tree): account for resolveInlineRef behavior when $ref has siblings

* feat(walker): adds a max depth option for refs

---------

Co-authored-by: Daniel A. White <[email protected]>

Author: P0lip

src/__tests__/__fixtures__/recursive-schema.json

@@ -0,0 +1,24 @@
+{
+  "title": "Thing",
+  "allOf": [
+    {
+      "$ref": "#/definitions/User"
+    }
+  ],
+  "description": "baz",
+  "definitions": {
+    "User": {
+      "type": "object",
+      "description": "user",
+      "properties": {
+        "manager": {
+          "$ref": "#/definitions/Boss"
+        }
+      }
+    },
+    "Boss": {
+      "$ref": "#/definitions/User",
+      "description": "xyz"
+    }
+  }
+}

src/__tests__/__fixtures__/references/with-overrides.json

@@ -0,0 +1,23 @@
+{
+  "oneOf": [
+    {
+      "$ref": "#/definitions/User"
+    }
+  ],
+  "description": "User Model",
+  "definitions": {
+    "User": {
+      "type": "object",
+      "description": "Plain User",
+      "properties": {
+        "manager": {
+          "$ref": "#/definitions/Admin"
+        }
+      }
+    },
+    "Admin": {
+      "$ref": "#/definitions/User",
+      "description": "Admin User"
+    }
+  }
+}

src/__tests__/__snapshots__/tree.spec.ts.snap

@@ -1293,6 +1293,29 @@ exports[`SchemaTree output should generate valid tree for references/nullish.jso
 "
 `;
 
+exports[`SchemaTree output should generate valid tree for references/with-overrides.json 1`] = `
+"└─ #
+   ├─ combiners
+   │  └─ 0: oneOf
+   └─ children
+      └─ 0
+         └─ #/oneOf/0
+            ├─ types
+            │  └─ 0: object
+            ├─ primaryType: object
+            └─ children
+               └─ 0
+                  └─ #/oneOf/0/properties/manager
+                     ├─ types
+                     │  └─ 0: object
+                     ├─ primaryType: object
+                     └─ children
+                        └─ 0
+                           └─ #/oneOf/0/properties/manager/properties/manager
+                              └─ mirrors: #/oneOf/0/properties/manager
+"
+`;
+
 exports[`SchemaTree output should generate valid tree for tickets.schema.json 1`] = `
 "└─ #
    ├─ types

src/__tests__/tree.spec.ts

@@ -13,7 +13,7 @@ describe('SchemaTree', () => {
     it.each(
       fastGlob.sync('**/*.json', {
         cwd: path.join(__dirname, '__fixtures__'),
-        ignore: ['stress-schema.json'],
+        ignore: ['stress-schema.json', 'recursive-schema.json'],
       }),
     )('should generate valid tree for %s', async filename => {
       const schema = JSON.parse(await fs.promises.readFile(path.resolve(__dirname, '__fixtures__', filename), 'utf8'));
@@ -985,4 +985,17 @@ describe('SchemaTree', () => {
       });
     });
   });
+
+  describe('recursive walking', () => {
+    it('should load with a max depth', async () => {
+      const schema = JSON.parse(
+        await fs.promises.readFile(path.resolve(__dirname, '__fixtures__', 'recursive-schema.json'), 'utf8'),
+      );
+
+      const w = new SchemaTree(schema, {
+        maxRefDepth: 1000,
+      });
+      w.populate();
+    });
+  });
 });

src/tree/tree.ts

@@ -11,18 +11,22 @@ import type { SchemaTreeOptions } from './types';
 export class SchemaTree {
   public walker: Walker;
   public root: RootNode;
+  private readonly resolvedRefs = new Map();
 
   constructor(public schema: SchemaFragment, protected readonly opts?: Partial<SchemaTreeOptions>) {
     this.root = new RootNode(schema);
+    this.resolvedRefs = new Map();
     this.walker = new Walker(this.root, {
       mergeAllOf: this.opts?.mergeAllOf !== false,
       resolveRef: opts?.refResolver === null ? null : this.resolveRef,
+      maxRefDepth: opts?.maxRefDepth,
     });
   }
 
   public destroy() {
     this.root.children.length = 0;
     this.walker.destroy();
+    this.resolvedRefs.clear();
   }
 
   public populate() {
@@ -34,6 +38,10 @@ export class SchemaTree {
   }
 
   protected resolveRef: WalkerRefResolver = (path, $ref) => {
+    if (this.resolvedRefs.has($ref)) {
+      return this.resolvedRefs.get($ref);
+    }
+
     const seenRefs: string[] = [];
     let cur$ref: unknown = $ref;
     let resolvedValue!: SchemaFragment;
@@ -48,6 +56,7 @@ export class SchemaTree {
       cur$ref = resolvedValue.$ref;
     }
 
+    this.resolvedRefs.set($ref, resolvedValue);
     return resolvedValue;
   };
 

src/tree/types.ts

@@ -2,7 +2,10 @@ import type { SchemaFragment } from '../types';
 
 export type SchemaTreeOptions = {
   mergeAllOf: boolean;
+  /** Resolves references to the schemas. If providing a custom implementation, it must return the same object reference for the same reference string. */
   refResolver: SchemaTreeRefDereferenceFn | null;
+  /** Controls the level of recursion of refs. Prevents overly complex trees and running out of stack depth. */
+  maxRefDepth?: number | null;
 };
 
 export type SchemaTreeRefInfo = {

src/walker/types.ts

@@ -6,7 +6,10 @@ export type WalkerRefResolver = (path: string[] | null, $ref: string) => SchemaF
 
 export type WalkingOptions = {
   mergeAllOf: boolean;
+  /** Resolves references to the schemas. If providing a custom implementation, it must return the same object reference for the same reference string. */
   resolveRef: WalkerRefResolver | null;
+  /** Controls the level of recursion of refs. Prevents overly complex trees and running out of stack depth. */
+  maxRefDepth?: number | null;
 };
 
 export type WalkerSnapshot = {

src/walker/walker.ts

@@ -36,11 +36,22 @@ export class Walker extends EventEmitter<WalkerEmitter> {
   constructor(protected readonly root: RootNode, protected readonly walkingOptions: WalkingOptions) {
     super();
 
+    let maxRefDepth = walkingOptions.maxRefDepth ?? null;
+    if (typeof maxRefDepth === 'number') {
+      if (maxRefDepth < 1) {
+        maxRefDepth = null;
+      } else if (maxRefDepth > 1000) {
+        // experimented with 1500 and the recursion limit is still lower than that
+        maxRefDepth = 1000;
+      }
+    }
+    walkingOptions.maxRefDepth = maxRefDepth;
+
     this.path = [];
     this.depth = -1;
     this.fragment = root.fragment;
     this.schemaNode = root;
-    this.processedFragments = new WeakMap<SchemaFragment, SchemaNode>();
+    this.processedFragments = new WeakMap();
     this.mergedAllOfs = new WeakMap();
 
     this.hooks = {};
@@ -51,7 +62,7 @@ export class Walker extends EventEmitter<WalkerEmitter> {
     this.depth = -1;
     this.fragment = this.root.fragment;
     this.schemaNode = this.root;
-    this.processedFragments = new WeakMap<SchemaFragment, RegularNode | ReferenceNode>();
+    this.processedFragments = new WeakMap();
     this.mergedAllOfs = new WeakMap();
   }
 
@@ -265,7 +276,7 @@ export class Walker extends EventEmitter<WalkerEmitter> {
   }
 
   protected processFragment(): [SchemaNode, ProcessedFragment] {
-    const { walkingOptions, path, fragment: originalFragment } = this;
+    const { walkingOptions, path, fragment: originalFragment, depth } = this;
     let { fragment } = this;
 
     let retrieved = isNonNullable(fragment) ? this.retrieveFromFragment(fragment, originalFragment) : null;
@@ -275,7 +286,9 @@ export class Walker extends EventEmitter<WalkerEmitter> {
     }
 
     if ('$ref' in fragment) {
-      if (typeof fragment.$ref !== 'string') {
+      if (typeof walkingOptions.maxRefDepth === 'number' && walkingOptions.maxRefDepth < depth) {
+        return [new ReferenceNode(fragment, `max $ref depth limit reached`), fragment];
+      } else if (typeof fragment.$ref !== 'string') {
         return [new ReferenceNode(fragment, '$ref is not a string'), fragment];
       } else if (walkingOptions.resolveRef !== null) {
         try {