From bd0663393175215ddba96ecf21ad087dd7a33c86 Mon Sep 17 00:00:00 2001
From: crschnick <crschnick@xpipe.io>
Date: Sun, 13 Oct 2024 05:40:45 +0000
Subject: [PATCH] Fixes for potential stack overflow

---
 .../io/xpipe/app/storage/DataStorage.java     | 15 +++++++++++----
 .../ext/base/script/SimpleScriptStore.java    | 19 ++++++++++---------
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/app/src/main/java/io/xpipe/app/storage/DataStorage.java b/app/src/main/java/io/xpipe/app/storage/DataStorage.java
index 33eeb41f7..4ff1c28c0 100644
--- a/app/src/main/java/io/xpipe/app/storage/DataStorage.java
+++ b/app/src/main/java/io/xpipe/app/storage/DataStorage.java
@@ -788,13 +788,20 @@ public abstract class DataStorage {
 
     public Set<DataStoreEntry> getDeepStoreChildren(DataStoreEntry entry) {
         var set = new HashSet<DataStoreEntry>();
-        getStoreChildren(entry).forEach(c -> {
-            set.add(c);
-            set.addAll(getDeepStoreChildren(c));
-        });
+        getDeepStoreChildren(entry, set);
         return set;
     }
 
+    private void getDeepStoreChildren(DataStoreEntry entry, Set<DataStoreEntry> current) {
+        getStoreChildren(entry).forEach(c -> {
+            var added = current.add(c);
+            // Guard against loop
+            if (added) {
+                getDeepStoreChildren(c, current);
+            }
+        });
+    }
+
     public Set<DataStoreEntry> getStoreChildren(DataStoreEntry entry) {
         if (entry.getValidity() == DataStoreEntry.Validity.LOAD_FAILED) {
             return Set.of();
diff --git a/ext/base/src/main/java/io/xpipe/ext/base/script/SimpleScriptStore.java b/ext/base/src/main/java/io/xpipe/ext/base/script/SimpleScriptStore.java
index d1e251b73..197b1880c 100644
--- a/ext/base/src/main/java/io/xpipe/ext/base/script/SimpleScriptStore.java
+++ b/ext/base/src/main/java/io/xpipe/ext/base/script/SimpleScriptStore.java
@@ -87,16 +87,17 @@ public class SimpleScriptStore extends ScriptStore implements ShellInitCommand.T
     }
 
     public void queryFlattenedScripts(LinkedHashSet<DataStoreEntryRef<SimpleScriptStore>> all) {
-        // Prevent loop
         DataStoreEntryRef<SimpleScriptStore> ref = getSelfEntry().ref();
-        all.add(ref);
-        getEffectiveScripts().stream()
-                .filter(scriptStoreDataStoreEntryRef -> !all.contains(scriptStoreDataStoreEntryRef))
-                .forEach(scriptStoreDataStoreEntryRef -> {
-                    scriptStoreDataStoreEntryRef.getStore().queryFlattenedScripts(all);
-                });
-        all.remove(ref);
-        all.add(ref);
+        var added = all.add(ref);
+        // Prevent loop
+        if (added) {
+            getEffectiveScripts().stream().filter(scriptStoreDataStoreEntryRef -> !all.contains(scriptStoreDataStoreEntryRef)).forEach(
+                    scriptStoreDataStoreEntryRef -> {
+                        scriptStoreDataStoreEntryRef.getStore().queryFlattenedScripts(all);
+                    });
+            all.remove(ref);
+            all.add(ref);
+        }
     }
 
     @Override