From f823b8891067c5f0844e3901b48a70a4ca873688 Mon Sep 17 00:00:00 2001
From: Timon Back <timonback@users.noreply.github.com>
Date: Fri, 8 Nov 2024 15:41:17 +0100
Subject: [PATCH] fix(amqp): avoid StackOverflowError when used with protobuf
 (#1062)

* fix(amqp): avoid StackOverflowError when used with protobuf

Additionally, add a ProtobufSchemaPostProcessor to remove fields generated by protobuf java generator

* chore(core): remove unnecessary ProtobufSchemaPostProcessor.java
---
 .../AvroSchemaPostProcessor.java              | 36 ++++++++++++++-----
 .../AvroSchemaPostProcessorTest.java          | 22 ++++++++++++
 2 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/springwolf-core/src/main/java/io/github/springwolf/core/asyncapi/components/postprocessors/AvroSchemaPostProcessor.java b/springwolf-core/src/main/java/io/github/springwolf/core/asyncapi/components/postprocessors/AvroSchemaPostProcessor.java
index f71278148..7fc14fab5 100644
--- a/springwolf-core/src/main/java/io/github/springwolf/core/asyncapi/components/postprocessors/AvroSchemaPostProcessor.java
+++ b/springwolf-core/src/main/java/io/github/springwolf/core/asyncapi/components/postprocessors/AvroSchemaPostProcessor.java
@@ -5,6 +5,10 @@
 import io.swagger.v3.oas.models.media.Schema;
 import org.springframework.util.StringUtils;
 
+import java.util.Deque;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -15,6 +19,7 @@
  * This may change in the future.
  */
 public class AvroSchemaPostProcessor implements SchemasPostProcessor {
+    private static final String SCHEMA_AVRO_PREFIX = "org.apache.avro.";
     private static final String SCHEMA_PROPERTY = "schema";
     private static final String SPECIFIC_DATA_PROPERTY = "specificData";
     private static final String SCHEMA_REF = "org.apache.avro.Schema";
@@ -22,18 +27,33 @@ public class AvroSchemaPostProcessor implements SchemasPostProcessor {
 
     @Override
     public void process(Schema schema, Map<String, Schema> definitions, String contentType) {
-        removeAvroSchemas(definitions);
-        removeAvroProperties(schema, definitions);
+        removeSchemas(definitions);
+
+        Deque<Schema> queue = new LinkedList<>(List.of(schema));
+        HashSet<Schema> visited = new HashSet<>();
+        while (!queue.isEmpty()) {
+            Schema currentSchema = queue.pop();
+            if (visited.contains(currentSchema)) {
+                continue;
+            }
+            visited.add(currentSchema);
+
+            processRefSchema(currentSchema, queue, definitions);
+            processProperties(currentSchema, queue);
+        }
     }
 
-    private void removeAvroProperties(Schema schema, Map<String, Schema> definitions) {
+    private void processRefSchema(Schema schema, Deque<Schema> queue, Map<String, Schema> definitions) {
         if (schema.get$ref() != null) {
             String schemaName = MessageReference.extractRefName(schema.get$ref());
-            if (definitions.containsKey(schemaName)) {
-                removeAvroProperties(definitions.get(schemaName), definitions);
+            Schema refedSchema = definitions.get(schemaName);
+            if (refedSchema != null) {
+                queue.add(refedSchema);
             }
         }
+    }
 
+    private void processProperties(Schema schema, Deque<Schema> queue) {
         Map<String, Schema> properties = schema.getProperties();
         if (properties != null) {
             Schema schemaPropertySchema = properties.getOrDefault(SCHEMA_PROPERTY, null);
@@ -46,11 +66,11 @@ private void removeAvroProperties(Schema schema, Map<String, Schema> definitions
                 }
             }
 
-            properties.forEach((key, value) -> removeAvroProperties(value, definitions));
+            properties.forEach((key, value) -> queue.add(value));
         }
     }
 
-    private void removeAvroSchemas(Map<String, Schema> definitions) {
-        definitions.entrySet().removeIf(entry -> StringUtils.startsWithIgnoreCase(entry.getKey(), "org.apache.avro"));
+    private void removeSchemas(Map<String, Schema> definitions) {
+        definitions.entrySet().removeIf(entry -> StringUtils.startsWithIgnoreCase(entry.getKey(), SCHEMA_AVRO_PREFIX));
     }
 }
diff --git a/springwolf-core/src/test/java/io/github/springwolf/core/asyncapi/components/postprocessors/AvroSchemaPostProcessorTest.java b/springwolf-core/src/test/java/io/github/springwolf/core/asyncapi/components/postprocessors/AvroSchemaPostProcessorTest.java
index c294109a5..ca09832b0 100644
--- a/springwolf-core/src/test/java/io/github/springwolf/core/asyncapi/components/postprocessors/AvroSchemaPostProcessorTest.java
+++ b/springwolf-core/src/test/java/io/github/springwolf/core/asyncapi/components/postprocessors/AvroSchemaPostProcessorTest.java
@@ -8,6 +8,7 @@
 import java.util.Map;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Fail.fail;
 
 class AvroSchemaPostProcessorTest {
     SchemasPostProcessor processor = new AvroSchemaPostProcessor();
@@ -80,4 +81,25 @@ void avroSchemasAreRemovedInRefsTest() {
                         "customClassRefUnusedInThisTest",
                         new StringSchema()));
     }
+
+    @Test
+    void handleRecursiveSchemasTest() {
+        var schema = new io.swagger.v3.oas.models.media.Schema();
+        schema.set$ref("#/components/schemas/intermediateSchema");
+
+        var intermediateSchema = new io.swagger.v3.oas.models.media.Schema();
+        intermediateSchema.set$ref("#/components/schemas/schema");
+
+        var definitions = new HashMap<String, io.swagger.v3.oas.models.media.Schema>();
+        definitions.put("schema", schema);
+        definitions.put("intermediateSchema", new StringSchema());
+
+        // when
+        try {
+            processor.process(schema, definitions, "content-type-ignored");
+        } catch (StackOverflowError ex) {
+            // then, no StackOverflowException is thrown
+            fail();
+        }
+    }
 }