From d3840edd73c0fb52a8e67812392d2efffd1a2c16 Mon Sep 17 00:00:00 2001
From: ValKmjolnir <lhk101lhk101@qq.com>
Date: Mon, 18 Dec 2023 23:43:37 +0800
Subject: [PATCH] :zap: avoid infinite recursion when loading module

---
 src/main.cpp         |  2 +-
 src/nasal_import.cpp | 35 +++++++++++++++++++++++++----------
 src/nasal_import.h   |  3 ++-
 src/repl.cpp         |  2 +-
 4 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/src/main.cpp b/src/main.cpp
index d9e7588..c92a703 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -132,7 +132,7 @@ void execute(
     }
 
     // linker gets parser's ast and load import files to this ast
-    ld.link(parse, file, cmd&VM_DETAIL).chkerr();
+    ld.link(parse, cmd&VM_DETAIL).chkerr();
     if (cmd&VM_REF_FILE) {
         if (ld.get_file_list().size()) {
             std::cout << "referenced file(s):\n";
diff --git a/src/nasal_import.cpp b/src/nasal_import.cpp
index 0d06ea7..b7fd3af 100644
--- a/src/nasal_import.cpp
+++ b/src/nasal_import.cpp
@@ -179,23 +179,38 @@ code_block* linker::import_regular_file(
     // check self import, avoid infinite loading loop
     if (check_self_import(filename)) {
         err.err("link",
-            "self-referenced module <" + filename + ">:\n" +
-            "    reference path: " + generate_self_import_path(filename)
+            node->get_location(),
+            "self-referenced module <" + filename + ">, " +
+            "reference path: " + generate_self_import_path(filename)
         );
         return new code_block({0, 0, 0, 0, filename});
     }
     check_exist_or_record_file(filename);
     
     module_load_stack.push_back(filename);
+    if (module_load_stack.size()>MAX_RECURSION_DEPTH) {
+        err.err("link",
+            node->get_location(),
+            "too deep module import stack (>" +
+            std::to_string(MAX_RECURSION_DEPTH) + ")."
+        );
+        return new code_block({0, 0, 0, 0, filename});
+    }
     // start importing...
     lexer nasal_lexer;
     parse nasal_parser;
     if (nasal_lexer.scan(filename).geterr()) {
-        err.err("link", "error occurred when analysing <" + filename + ">");
+        err.err("link",
+            node->get_location(),
+            "error occurred when analysing <" + filename + ">"
+        );
         return new code_block({0, 0, 0, 0, filename});
     }
     if (nasal_parser.compile(nasal_lexer).geterr())  {
-        err.err("link", "error occurred when analysing <" + filename + ">");
+        err.err("link",
+            node->get_location(),
+            "error occurred when analysing <" + filename + ">"
+        );
         return new code_block({0, 0, 0, 0, filename});
     }
     // swap result out
@@ -378,19 +393,19 @@ code_block* linker::load(code_block* program_root, const std::string& filename)
     return tree;
 }
 
-const error& linker::link(
-    parse& parse, const std::string& self, bool spath = false) {
+const error& linker::link(parse& parse, bool spath = false) {
     // switch for showing path when errors occur
     show_path_flag = spath;
 
     // initializing file map
-    this_file = self;
-    imported_files = {self};
-    module_load_stack = {self};
+    this_file = parse.tree()->get_location().file;
+    imported_files = {this_file};
+    module_load_stack = {this_file};
 
     // scan root and import files
     // then generate a new ast and return to import_ast
-    auto new_tree_root = load(parse.tree(), self);
+    // dfs load file
+    auto new_tree_root = load(parse.tree(), this_file);
     auto old_tree_root = parse.swap(new_tree_root);
     delete old_tree_root;
     return err;
diff --git a/src/nasal_import.h b/src/nasal_import.h
index 692c922..efac02f 100644
--- a/src/nasal_import.h
+++ b/src/nasal_import.h
@@ -23,6 +23,7 @@ namespace nasal {
 
 class linker {
 private:
+    const u32 MAX_RECURSION_DEPTH = 256;
     bool show_path_flag;
     bool library_loaded;
     std::string this_file;
@@ -48,7 +49,7 @@ private:
 
 public:
     linker();
-    const error& link(parse&, const std::string&, bool);
+    const error& link(parse&, bool);
     const auto& get_file_list() const {return imported_files;}
 };
 
diff --git a/src/repl.cpp b/src/repl.cpp
index 2773700..fe6c9f5 100644
--- a/src/repl.cpp
+++ b/src/repl.cpp
@@ -92,7 +92,7 @@ bool repl::run() {
         return false;
     }
 
-    if (nasal_linker->link(*nasal_parser, "<nasal-repl>", true).geterr()) {
+    if (nasal_linker->link(*nasal_parser, true).geterr()) {
         return false;
     }