Restrict subdirectory imports to main.kcl (#7094)

Signed-off-by: Nick Cameron <nrc@ncameron.org>

rust/kcl-lib/src/parsing/ast/types/mod.rs

@@ -1811,14 +1811,25 @@ impl ImportStatement {
 
         match &self.path {
             ImportPath::Kcl { filename: s } | ImportPath::Foreign { path: s } => {
-                let name = s.file_name().map(|f| f.to_string());
+                let name = s.to_string_lossy();
+                if name.ends_with("/main.kcl") || name.ends_with("\\main.kcl") {
+                    let name = &name[..name.len() - 9];
+                    let start = name.rfind(['/', '\\']).map(|s| s + 1).unwrap_or(0);
+                    return Some(name[start..].to_owned());
+                }
+
+                let name = s.file_name().map(|f| f.to_string())?;
+                if name.contains('\\') || name.contains('/') {
+                    return None;
+                }
+
                 // Remove the extension if it exists.
                 let extension = s.extension();
-                if let Some(extension) = extension {
-                    name.map(|n| n.trim_end_matches(extension).trim_end_matches('.').to_string())
+                Some(if let Some(extension) = extension {
+                    name.trim_end_matches(extension).trim_end_matches('.').to_string()
                 } else {
                     name
-                }
+                })
             }
             ImportPath::Std { path } => path.last().cloned(),
         }
@@ -4330,4 +4341,20 @@ startSketchOn(XY)
 "#
         );
     }
+
+    #[test]
+    fn module_name() {
+        #[track_caller]
+        fn assert_mod_name(stmt: &str, name: &str) {
+            let tokens = crate::parsing::token::lex(stmt, ModuleId::default()).unwrap();
+            let stmt = crate::parsing::parser::import_stmt(&mut tokens.as_slice()).unwrap();
+            assert_eq!(stmt.module_name().unwrap(), name);
+        }
+
+        assert_mod_name("import 'foo.kcl'", "foo");
+        assert_mod_name("import 'foo.kcl' as bar", "bar");
+        assert_mod_name("import 'main.kcl'", "main");
+        assert_mod_name("import 'foo/main.kcl'", "foo");
+        assert_mod_name("import 'foo\\bar\\main.kcl'", "bar");
+    }
 }

rust/kcl-lib/src/parsing/parser.rs

@@ -1729,7 +1729,7 @@ fn glob(i: &mut TokenSlice) -> PResult<Token> {
         .parse_next(i)
 }
 
-fn import_stmt(i: &mut TokenSlice) -> PResult<BoxNode<ImportStatement>> {
+pub(super) fn import_stmt(i: &mut TokenSlice) -> PResult<BoxNode<ImportStatement>> {
     let (visibility, visibility_token) = opt(terminated(item_visibility, whitespace))
         .parse_next(i)?
         .map_or((ItemVisibility::Default, None), |pair| (pair.0, Some(pair.1)));
@@ -1867,7 +1867,7 @@ fn validate_path_string(path_string: String, var_name: bool, path_range: SourceR
             return Err(ErrMode::Cut(
                 CompilationError::fatal(
                     path_range,
-                    "import path may only contain alphanumeric characters, underscore, hyphen, and period. KCL files in other directories are not yet supported.",
+                    "import path may only contain alphanumeric characters, `_`, `-`, `.`, `/`, and `\\`.",
                 )
                 .into(),
             ));
@@ -1894,6 +1894,15 @@ fn validate_path_string(path_string: String, var_name: bool, path_range: SourceR
             ));
         }
 
+        if (path_string.contains('/') || path_string.contains('\\'))
+            && !(path_string.ends_with("/main.kcl") || path_string.ends_with("\\main.kcl"))
+        {
+            return Err(ErrMode::Cut(
+                CompilationError::fatal(path_range, "import path to a subdirectory must only refer to main.kcl.")
+                    .into(),
+            ));
+        }
+
         ImportPath::Kcl {
             filename: TypedPath::new(&path_string),
         }
@@ -4569,9 +4578,14 @@ e
         );
         assert_err(
             r#"import cube from "C:\cube.kcl""#,
-            "import path may only contain alphanumeric characters, underscore, hyphen, and period. KCL files in other directories are not yet supported.",
+            "import path may only contain alphanumeric characters, `_`, `-`, `.`, `/`, and `\\`.",
             [17, 30],
         );
+        assert_err(
+            r#"import cube from "cube/cube.kcl""#,
+            "import path to a subdirectory must only refer to main.kcl.",
+            [17, 32],
+        );
         assert_err(
             r#"import * as foo from "dsfs""#,
             "as is not the 'from' keyword",