ast/schema: Fix race condition in parsing with reused references (#5184)

Resolves #5183

Signed-off-by: Liam Galvin <liam.galvin@aquasec.com>

ast/compile.go

@@ -1068,6 +1068,10 @@ func newSchemaParser() *schemaParser {
 }
 
 func (parser *schemaParser) parseSchema(schema interface{}) (types.Type, error) {
+	return parser.parseSchemaWithPropertyKey(schema, "")
+}
+
+func (parser *schemaParser) parseSchemaWithPropertyKey(schema interface{}, propertyKey string) (types.Type, error) {
 	subSchema, ok := schema.(*gojsonschema.SubSchema)
 	if !ok {
 		return nil, fmt.Errorf("unexpected schema type %v", subSchema)
@@ -1075,10 +1079,10 @@ func (parser *schemaParser) parseSchema(schema interface{}) (types.Type, error)
 
 	// Handle referenced schemas, returns directly when a $ref is found
 	if subSchema.RefSchema != nil {
-		if existing, ok := parser.definitionCache[subSchema.RefSchema.ID.String()]; ok {
+		if existing, ok := parser.definitionCache[subSchema.Ref.String()]; ok {
 			return types.NewObject(existing.properties, nil), nil
 		}
-		return parser.parseSchema(subSchema.RefSchema)
+		return parser.parseSchemaWithPropertyKey(subSchema.RefSchema, subSchema.Ref.String())
 	}
 
 	// Handle anyOf
@@ -1154,8 +1158,8 @@ func (parser *schemaParser) parseSchema(schema interface{}) (types.Type, error)
 				for _, pSchema := range subSchema.PropertiesChildren {
 					def.properties = append(def.properties, types.NewStaticProperty(pSchema.Property, nil))
 				}
-				if subSchema.Parent != nil {
-					parser.definitionCache[subSchema.ID.String()] = def
+				if propertyKey != "" {
+					parser.definitionCache[propertyKey] = def
 				}
 				for _, pSchema := range subSchema.PropertiesChildren {
 					newtype, err := parser.parseSchema(pSchema)

ast/compile_test.go

@@ -7454,3 +7454,117 @@ deny {
 	}
 
 }
+
+func TestCompilerWithRecursiveSchemaAvoidRace(t *testing.T) {
+
+	jsonSchema := `{
+  "type": "object",
+  "properties": {
+    "aws": {
+      "type": "object",
+      "$ref": "#/$defs/example.pkg.providers.aws.AWS"
+    }
+  },
+  "$defs": {
+    "example.pkg.providers.aws.AWS": {
+      "type": "object",
+      "properties": {
+        "iam": {
+          "type": "object",
+          "$ref": "#/$defs/example.pkg.providers.aws.iam.IAM"
+        },
+        "sqs": {
+          "type": "object",
+          "$ref": "#/$defs/example.pkg.providers.aws.sqs.SQS"
+        }
+      }
+    },
+    "example.pkg.providers.aws.iam.Document": {
+      "type": "object"
+    },
+    "example.pkg.providers.aws.iam.IAM": {
+      "type": "object",
+      "properties": {
+        "policies": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "$ref": "#/$defs/example.pkg.providers.aws.iam.Policy"
+          }
+        }
+      }
+    },
+    "example.pkg.providers.aws.iam.Policy": {
+      "type": "object",
+      "properties": {
+        "builtin": {
+          "type": "object",
+          "properties": {
+            "value": {
+              "type": "boolean"
+            }
+          }
+        },
+        "document": {
+          "type": "object",
+          "$ref": "#/$defs/example.pkg.providers.aws.iam.Document"
+        }
+      }
+    },
+    "example.pkg.providers.aws.sqs.Queue": {
+      "type": "object",
+      "properties": {
+        "policies": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "$ref": "#/$defs/example.pkg.providers.aws.iam.Policy"
+          }
+        }
+      }
+    },
+    "example.pkg.providers.aws.sqs.SQS": {
+      "type": "object",
+      "properties": {
+        "queues": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "$ref": "#/$defs/example.pkg.providers.aws.sqs.Queue"
+          }
+        }
+      }
+    }
+  }
+}`
+
+	exampleModule := `# METADATA
+# schemas:
+#  - input: schema.input
+package race.condition
+
+deny {
+	queue := input.aws.sqs.queues[_]
+	policy := queue.policies[_]
+	doc := json.unmarshal(policy.document.value)
+	statement = doc.Statement[_]
+	action := statement.Action[_]
+	action == "*"
+}
+`
+
+	c := NewCompiler()
+	var schema interface{}
+	if err := json.Unmarshal([]byte(jsonSchema), &schema); err != nil {
+		t.Fatal(err)
+	}
+	schemaSet := NewSchemaSet()
+	schemaSet.Put(MustParseRef("schema.input"), schema)
+	c.WithSchemas(schemaSet)
+
+	m := MustParseModuleWithOpts(exampleModule, ParserOptions{ProcessAnnotation: true})
+	c.Compile(map[string]*Module{"testMod": m})
+	if c.Failed() {
+		t.Fatal(c.Errors)
+	}
+}