Restrict the user from using the wrapper name

ballerina-platform · Oct 27, 2023 · 3ddab8a · 3ddab8a
1 parent 9747eee
commit 3ddab8a
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 9 deletions.
diff --git a/...ina-shell/modules/shell-core/src/main/java/io/ballerina/shell/parser/ParserConstants.java b/...ina-shell/modules/shell-core/src/main/java/io/ballerina/shell/parser/ParserConstants.java
@@ -26,6 +26,7 @@
  * @since 2.0.0
  */
 public class ParserConstants {
+    public static final String WRAPPER_PREFIX = "__shell_wrapper__";
 
     public static final Set<String> RESTRICTED_FUNCTION_NAMES = Set.of("main", "init", "__java_recall",
             "__java_memorize", "__recall_any", "__recall_any_error", "__memorize", "__stmts", "__run");

diff --git a/...na-shell/modules/shell-core/src/main/java/io/ballerina/shell/parser/SerialTreeParser.java b/...na-shell/modules/shell-core/src/main/java/io/ballerina/shell/parser/SerialTreeParser.java
@@ -123,7 +123,8 @@ public Collection<Node> parseDeclarations(String source) throws TreeParserExcept
     private boolean isModuleDeclarationAllowed(ModuleMemberDeclarationNode declarationNode) {
         if (declarationNode instanceof FunctionDefinitionNode) {
             String functionName = ((FunctionDefinitionNode) declarationNode).functionName().text();
-            if (RESTRICTED_FUNCTION_NAMES.contains(functionName)) {
+            if (RESTRICTED_FUNCTION_NAMES.contains(functionName) ||
+                    functionName.startsWith(ParserConstants.WRAPPER_PREFIX)) {
                 addWarnDiagnostic("Found '" + functionName + "' function in the declarations.\n" +
                         "Discarded '" + functionName + "' function without loading.");
                 return false;

diff --git a/.../modules/shell-core/src/main/java/io/ballerina/shell/parser/trials/ModuleMemberTrial.java b/.../modules/shell-core/src/main/java/io/ballerina/shell/parser/trials/ModuleMemberTrial.java
@@ -80,7 +80,8 @@ public Collection<Node> parse(String source) throws ParserTrialFailedException {
     private void validateModuleDeclaration(ModuleMemberDeclarationNode declarationNode) {
         if (declarationNode instanceof FunctionDefinitionNode) {
             String functionName = ((FunctionDefinitionNode) declarationNode).functionName().text();
-            if (RESTRICTED_FUNCTION_NAMES.contains(functionName)) {
+            if (RESTRICTED_FUNCTION_NAMES.contains(functionName) ||
+                    functionName.startsWith(ParserConstants.WRAPPER_PREFIX)) {
                 String message = "Function name " + "'" + functionName + "'" + " not allowed in Ballerina Shell.\n";
                 throw new InvalidMethodException(message);
             }

diff --git a/...ules/shell-core/src/main/java/io/ballerina/shell/snippet/factory/BasicSnippetFactory.java b/...ules/shell-core/src/main/java/io/ballerina/shell/snippet/factory/BasicSnippetFactory.java
@@ -65,6 +65,7 @@
 import io.ballerina.compiler.syntax.tree.WhileStatementNode;
 import io.ballerina.compiler.syntax.tree.XMLNamespaceDeclarationNode;
 import io.ballerina.shell.exceptions.SnippetException;
+import io.ballerina.shell.parser.ParserConstants;
 import io.ballerina.shell.snippet.SnippetSubKind;
 import io.ballerina.shell.snippet.types.ExpressionSnippet;
 import io.ballerina.shell.snippet.types.ImportDeclarationSnippet;
@@ -168,14 +169,14 @@ public List<VariableDeclarationSnippet> createVariableDeclarationSnippets(Node n
                 String functionBody = varNode.initializer().get().toString();
                 String functionDefinition = (hasCheck ? "function() returns error|" :
                         "function() returns ") + typeDescriptorNode;
-                String functionString =
-                        functionDefinition + "f_" + varFunctionCount + " = " + functionDefinition + "{" +
-                                typeDescriptorNode + bindingPatternNode.toString() + " = " + functionBody + ";" +
-                                "return " + bindingPatternNode + ";};";
+                String functionName = ParserConstants.WRAPPER_PREFIX + varFunctionCount;
+                String functionString = String.format("%s %s = %s {%s %s = %s; return %s;};", functionDefinition,
+                        functionName, functionDefinition, typeDescriptorNode, bindingPatternNode, functionBody,
+                        bindingPatternNode);
                 varNode = (VariableDeclarationNode) NodeParser.parseStatement(functionString);
-                newNode = (VariableDeclarationNode) NodeParser
-                        .parseStatement(typeDescriptorNode + " " + bindingPatternNode + " = " +
-                                (hasCheck ? "check " : "") + " f_" + varFunctionCount + "();");
+                newNode = (VariableDeclarationNode) NodeParser.parseStatement(
+                        String.format("%s %s = %s %s();", typeDescriptorNode, bindingPatternNode,
+                                (hasCheck ? "check " : ""), functionName));
             }
 
             varFunctionCount += 1;