From cc2fc3f84bdcaea5ee623098e2aadbe7a4d2add9 Mon Sep 17 00:00:00 2001
From: IsaacShelton <isaac.a.shelton@gmail.com>
Date: Wed, 28 Aug 2024 13:58:51 -0500
Subject: [PATCH] Restricted pragma directives to only allow basic constructs
 for now

---
 src/cli.rs                          |  2 +
 src/pragma_section/parse.rs         | 69 ++++++++++++++++++++++++++---
 src/pragma_section/run.rs           |  1 +
 src/workspace/compile/module/mod.rs |  3 +-
 4 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/src/cli.rs b/src/cli.rs
index c2f40a02..d99059ca 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -112,6 +112,7 @@ pub struct BuildOptions {
     pub interpret: bool,
     pub coerce_main_signature: bool,
     pub excute_result: bool,
+    pub allow_experimental_pragma_features: bool,
     pub use_pic: Option<bool>,
     pub target: Target,
     pub infrastructure: Option<PathBuf>,
@@ -133,6 +134,7 @@ impl Default for BuildOptions {
             interpret: false,
             coerce_main_signature: true,
             excute_result: false,
+            allow_experimental_pragma_features: false,
             use_pic: None,
             target: Target::HOST,
             infrastructure: Some(infrastructure),
diff --git a/src/pragma_section/parse.rs b/src/pragma_section/parse.rs
index e6fa6c29..82a1b669 100644
--- a/src/pragma_section/parse.rs
+++ b/src/pragma_section/parse.rs
@@ -1,14 +1,17 @@
 use super::PragmaSection;
 use crate::{
-    ast::{AstFile, Function, Parameters, StmtKind, TypeKind},
+    ast::{AstFile, Expr, ExprKind, Function, Parameters, Stmt, StmtKind, TypeKind},
+    diagnostics::ErrorDiagnostic,
     inflow::Inflow,
     parser::{self, error::ParseError, Input},
     show::{into_show, Show},
+    source_files::Source,
     token::{Token, TokenKind},
 };
 
 impl PragmaSection {
     pub fn parse<'a, I: Inflow<Token>>(
+        allow_experimental_pragma_features: bool,
         mut input: Input<'a, I>,
     ) -> Result<(PragmaSection, Input<'a, I>), Box<dyn Show>> {
         // pragma ...
@@ -33,11 +36,18 @@ impl PragmaSection {
             // "Whole-file" mode
 
             // Parse top-level contructs until we hit a '}'
-            while !parser.input.peek_is(TokenKind::CloseCurly) {
-                parser
-                    .parse_top_level(&mut ast_file, vec![])
-                    .map_err(into_show)?;
-                parser.input.ignore_newlines();
+            if allow_experimental_pragma_features {
+                while !parser.input.peek_is_or_eof(TokenKind::CloseCurly) {
+                    parser
+                        .parse_top_level(&mut ast_file, vec![])
+                        .map_err(into_show)?;
+                    parser.input.ignore_newlines();
+                }
+            } else {
+                return Err(Box::new(ErrorDiagnostic::new(
+                    "Whole-file pragma directives are an experimental feature and may be removed",
+                    pragma_source,
+                )));
             }
 
             // Eat the final '}'
@@ -62,6 +72,12 @@ impl PragmaSection {
                 vec![StmtKind::Expr(expr).at(expr_source)]
             };
 
+            if !allow_experimental_pragma_features {
+                for stmt in stmts.iter() {
+                    restrict_allowed_stmt(stmt)?;
+                }
+            }
+
             ast_file.functions.push(Function {
                 name: "main".into(),
                 parameters: Parameters {
@@ -92,3 +108,44 @@ impl PragmaSection {
         ))
     }
 }
+
+fn restrict_allowed_stmt(stmt: &Stmt) -> Result<(), Box<dyn Show>> {
+    match &stmt.kind {
+        StmtKind::Expr(inner) => restrict_allowed_expr(inner)?,
+        _ => return Err(found_forbidden_stmt(stmt.source)),
+    }
+
+    Ok(())
+}
+
+fn restrict_allowed_expr(expr: &Expr) -> Result<(), Box<dyn Show>> {
+    match &expr.kind {
+        ExprKind::NullTerminatedString(_) => (),
+        ExprKind::Call(call) => {
+            if call.expected_to_return.is_some() {
+                return Err(found_forbidden_expr(expr.source));
+            }
+
+            for argument in call.arguments.iter() {
+                restrict_allowed_expr(argument)?;
+            }
+        }
+        _ => return Err(found_forbidden_expr(expr.source)),
+    }
+
+    Ok(())
+}
+
+fn found_forbidden_stmt(source: Source) -> Box<dyn Show> {
+    Box::new(ErrorDiagnostic::new(
+        "Support for this statement inside pragma directives is an experimental feature that may be removed",
+        source,
+    ))
+}
+
+fn found_forbidden_expr(source: Source) -> Box<dyn Show> {
+    Box::new(ErrorDiagnostic::new(
+        "Support for this expression inside pragma directives is an experimental feature that may be removed",
+        source,
+    ))
+}
diff --git a/src/pragma_section/run.rs b/src/pragma_section/run.rs
index 36ec6d5d..6d79c2dd 100644
--- a/src/pragma_section/run.rs
+++ b/src/pragma_section/run.rs
@@ -29,6 +29,7 @@ impl PragmaSection {
                 coerce_main_signature: false,
                 excute_result: false,
                 use_pic: None,
+                allow_experimental_pragma_features: false,
                 target: Target::HOST,
                 infrastructure: None,
             },
diff --git a/src/workspace/compile/module/mod.rs b/src/workspace/compile/module/mod.rs
index 8bfe1bb1..52d1cbd0 100644
--- a/src/workspace/compile/module/mod.rs
+++ b/src/workspace/compile/module/mod.rs
@@ -36,7 +36,8 @@ pub fn compile_module_file<'a>(
     let mut settings = None;
 
     while input.peek_is(TokenKind::PragmaKeyword) {
-        let (section, rest_input) = PragmaSection::parse(input)?;
+        let (section, rest_input) =
+            PragmaSection::parse(compiler.options.allow_experimental_pragma_features, input)?;
         input = rest_input;
         settings = Some(section.run(compiler, path, settings)?);
         input.ignore_newlines();