Fix deep-review findings for prepare/closed-handle safety

Agent-Logs-Url: https://github.com/microsoft/regorus/sessions/c2c0f3ef-f641-4d83-94f8-a3cd3ede6db8

Co-authored-by: anakrish <35780660+anakrish@users.noreply.github.com>

bindings/java/src/lib.rs

@@ -27,13 +27,17 @@ pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeNewEngine(
 
 #[no_mangle]
 pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeClone(
-    _env: EnvUnowned,
+    env: EnvUnowned,
     _class: JClass,
     engine_ptr: jlong,
 ) -> jlong {
-    let engine = unsafe { &mut *(engine_ptr as *mut Engine) };
-    let c = engine.clone();
-    Box::into_raw(Box::new(c)) as jlong
+    let res = throw_err(env, |_env| {
+        let engine = unsafe { &mut *get_engine_ptr(engine_ptr)? };
+        let c = engine.clone();
+        Ok(Box::into_raw(Box::new(c)) as jlong)
+    });
+
+    res.unwrap_or_default()
 }
 
 #[no_mangle]
@@ -43,7 +47,7 @@ pub extern "system" fn Java_com_microsoft_regorus_Engine_nativePrepare(
     engine_ptr: jlong,
 ) {
     let _ = throw_err(env, |_env| {
-        let engine = unsafe { &mut *(engine_ptr as *mut Engine) };
+        let engine = unsafe { &mut *get_engine_ptr(engine_ptr)? };
         engine.prepare()?;
         Ok(())
     });
@@ -450,6 +454,9 @@ pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeDestroyEngine(
     _class: JClass,
     engine_ptr: jlong,
 ) {
+    if engine_ptr == 0 {
+        return;
+    }
     unsafe {
         let _engine = Box::from_raw(engine_ptr as *mut Engine);
     }
@@ -829,6 +836,13 @@ fn throw_err<T>(mut env: EnvUnowned, f: impl FnOnce(&mut Env) -> Result<T>) -> R
     }
 }
 
+fn get_engine_ptr(engine_ptr: jlong) -> Result<*mut Engine> {
+    if engine_ptr == 0 {
+        return Err(anyhow::anyhow!("Engine is closed"));
+    }
+    Ok(engine_ptr as *mut Engine)
+}
+
 fn get_string_array(env: &mut Env, array: jobjectArray) -> Result<Vec<String>> {
     if array.is_null() {
         return Ok(Vec::new());

bindings/java/src/main/java/com/microsoft/regorus/Engine.java

@@ -46,7 +46,7 @@ public class Engine implements AutoCloseable, Cloneable {
 
     // Pointer to Engine allocated on Rust's heap, all native methods works on
     // engine expects this pointer. It is free'd in `close` method.
-    private final long enginePtr;
+    private long enginePtr;
 
     /**
      * Creates a new Regorus Engine.
@@ -64,7 +64,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * Efficiently clones an Engine.
      */
     public Engine clone() {
-	return new Engine(nativeClone(enginePtr));
+	return new Engine(nativeClone(requireOpen()));
     }
 
     /**
@@ -72,7 +72,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * Optional: if skipped, first evaluation performs the same setup.
      */
     public void prepare() {
-        nativePrepare(enginePtr);
+        nativePrepare(requireOpen());
     }
     
      /**
@@ -82,7 +82,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * 
      */
     public void setRegoV0(boolean enable) {
-        nativeSetRegoV0(enginePtr, enable);
+        nativeSetRegoV0(requireOpen(), enable);
     }
 
     /**
@@ -94,7 +94,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * @return Rego package defined in the policy.
      */
     public String addPolicy(String filename, String rego) {
-        return nativeAddPolicy(enginePtr, filename, rego);
+        return nativeAddPolicy(requireOpen(), filename, rego);
     }
 
     /**
@@ -105,7 +105,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * @return Rego package defined in the policy.
      */
     public String addPolicyFromFile(String path) {
-        return nativeAddPolicyFromFile(enginePtr, path);
+        return nativeAddPolicyFromFile(requireOpen(), path);
     }
 
     /**
@@ -114,7 +114,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * @return List of Rego packages as a JSON array of strings.
      */
     public String getPackages() {
-        return nativeGetPackages(enginePtr);
+        return nativeGetPackages(requireOpen());
     }
     
     /**
@@ -123,14 +123,14 @@ public class Engine implements AutoCloseable, Cloneable {
      * @return List of Rego policies as a JSON array of sources.
      */
     public String getPolicies() {
-        return nativeGetPolicies(enginePtr);
+        return nativeGetPolicies(requireOpen());
     }
     
     /**
      * Clears the data  document.
      */
     public void clearData() {
-        nativeClearData(enginePtr);
+        nativeClearData(requireOpen());
     }
 
     /**
@@ -152,7 +152,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * @param data Inline data document.
      */
     public void addDataJson(String data) throws RuntimeException {
-        nativeAddDataJson(enginePtr, data);
+        nativeAddDataJson(requireOpen(), data);
     }
 
     /**
@@ -169,7 +169,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * @param path Path to JSON data document.
      */
     public void addDataJsonFromFile(String path) throws RuntimeException {
-        nativeAddDataJsonFromFile(enginePtr, path);
+        nativeAddDataJsonFromFile(requireOpen(), path);
     }
 
     /**
@@ -178,7 +178,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * @param input inline JSON input.
      */
     public void setInputJson(String input) {
-        nativeSetInputJson(enginePtr, input);
+        nativeSetInputJson(requireOpen(), input);
     }
 
     /**
@@ -187,7 +187,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * @param path Path to JSON input.
      */
     public void setInputJsonFromFile(String path) {
-        nativeSetInputJsonFromFile(enginePtr, path);
+        nativeSetInputJsonFromFile(requireOpen(), path);
     }
 
     /**
@@ -198,7 +198,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * @return Query results as a JSON string.
      */
     public String evalQuery(String query) {
-        return nativeEvalQuery(enginePtr, query);
+        return nativeEvalQuery(requireOpen(), query);
     }
     
     /**
@@ -209,7 +209,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * @return Value of the rule as a JSON string.
      */
     public String evalRule(String rule) {
-        return nativeEvalRule(enginePtr, rule);
+        return nativeEvalRule(requireOpen(), rule);
     }
     
     /**
@@ -219,7 +219,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * 
      */
     public void setEnableCoverage(boolean enable) {
-        nativeSetEnableCoverage(enginePtr, enable);
+        nativeSetEnableCoverage(requireOpen(), enable);
     }
     
     /**
@@ -227,7 +227,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * 
      */
     public void clearCoverageData() {
-        nativeClearCoverageData(enginePtr);
+        nativeClearCoverageData(requireOpen());
     }
     
     /**
@@ -237,7 +237,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * 
      */
     public String getCoverageReport() {
-        return nativeGetCoverageReport(enginePtr);
+        return nativeGetCoverageReport(requireOpen());
     }
     
     /**
@@ -247,7 +247,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * 
      */
     public String getCoverageReportPretty() {
-        return nativeGetCoverageReportPretty(enginePtr);
+        return nativeGetCoverageReportPretty(requireOpen());
     }
     
     /**
@@ -257,7 +257,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * 
      */
     public void setGatherPrints(boolean b) {
-        nativeSetGatherPrints(enginePtr, b);
+        nativeSetGatherPrints(requireOpen(), b);
     }
     
     /**
@@ -267,7 +267,7 @@ public class Engine implements AutoCloseable, Cloneable {
      * 
      */
     public String takePrints() {
-        return nativeTakePrints(enginePtr);
+        return nativeTakePrints(requireOpen());
     }
 
     /**
@@ -276,24 +276,34 @@ public class Engine implements AutoCloseable, Cloneable {
      * @param config Policy length configuration.
      */
     public void setPolicyLengthConfig(PolicyLengthConfig config) {
-        nativeSetPolicyLengthConfig(enginePtr, config.maxCol, config.maxFileBytes, config.maxLines);
+        nativeSetPolicyLengthConfig(requireOpen(), config.maxCol, config.maxFileBytes, config.maxLines);
     }
 
     /**
      * Clear the policy length configuration, reverting to defaults.
      */
     public void clearPolicyLengthConfig() {
-        nativeClearPolicyLengthConfig(enginePtr);
+        nativeClearPolicyLengthConfig(requireOpen());
     }
 
     long getPtr() {
+        return requireOpen();
+    }
+
+    private long requireOpen() {
+        if (enginePtr == 0) {
+            throw new IllegalStateException("Engine is closed");
+        }
         return enginePtr;
     }
 
     
     @Override
     public void close() {
-        nativeDestroyEngine(enginePtr);
+        if (enginePtr != 0) {
+            nativeDestroyEngine(enginePtr);
+            enginePtr = 0;
+        }
     }
 
     // Loading native library from JAR is adapted from:

bindings/java/src/test/java/com/microsoft/regorus/EngineTest.java

@@ -44,4 +44,28 @@ public class EngineTest extends TestCase
         Map expression = (Map) expressions.get(0);
         Assert.assertEquals("Hello, World!", expression.get("value"));
     }
+
+    public void test_closed_engine_operations_throw()
+    {
+        Engine engine = new Engine();
+        engine.close();
+
+        try {
+            engine.prepare();
+            fail("prepare should fail on closed engine");
+        } catch (IllegalStateException expected) {
+        }
+
+        try {
+            engine.clone();
+            fail("clone should fail on closed engine");
+        } catch (IllegalStateException expected) {
+        }
+
+        try {
+            engine.evalQuery("data");
+            fail("evalQuery should fail on closed engine");
+        } catch (IllegalStateException expected) {
+        }
+    }
 }

src/engine.rs

@@ -1125,11 +1125,10 @@ impl Engine {
         limits::enforce_memory_limit().map_err(|err| anyhow!(err))?;
 
         self.interpreter.set_traces(enable_tracing);
-        #[cfg(feature = "azure_policy")]
-        let was_prepared = self.prepared;
+        let newly_prepared = !self.prepared;
 
         // if the data/policies have changed or the interpreter has never been prepared
-        if !self.prepared {
+        if newly_prepared {
             // Analyze the modules and determine how statements must be scheduled.
             let analyzer = Analyzer::new();
             let schedule = Rc::new(analyzer.analyze(&self.modules)?);
@@ -1159,8 +1158,6 @@ impl Engine {
 
             // Set schedule after hoisting completes
             self.interpreter.set_schedule(Some(schedule));
-
-            self.prepared = true;
         }
 
         #[cfg(feature = "azure_policy")]
@@ -1174,7 +1171,7 @@ impl Engine {
                 )?;
                 // Infer resource types
                 crate::interpreter::target::infer::infer_resource_type(&mut self.interpreter)?;
-            } else if !was_prepared {
+            } else if newly_prepared {
                 // Check if any module specifies a target and warn if so.
                 self.warn_if_targets_present();
             }
@@ -1182,6 +1179,10 @@ impl Engine {
         #[cfg(not(feature = "azure_policy"))]
         let _ = for_target;
 
+        if newly_prepared {
+            self.prepared = true;
+        }
+
         Ok(())
     }
 

tests/engine/mod.rs

@@ -172,7 +172,56 @@ fn prepare_then_compile_for_target() -> Result<()> {
     engine.prepare()?;
     let compiled = engine.compile_for_target()?;
     let info = compiled.get_policy_info()?;
-    assert_eq!(info.target_name.as_deref(), Some("target.tests.sample_test_target"));
+    assert_eq!(
+        info.target_name.as_deref(),
+        Some("target.tests.sample_test_target")
+    );
+
+    let result = compiled.eval_with_input(Value::from_json_str(
+        r#"{"name":"resource-1","type":"test_resource"}"#,
+    )?)?;
+    assert_eq!(result, Value::from(true));
+
+    Ok(())
+}
+
+#[test]
+#[cfg(feature = "azure_policy")]
+#[cfg_attr(docsrs, doc(cfg(feature = "azure_policy")))]
+fn prepare_then_compile_for_target_error_recovery() -> Result<()> {
+    let target_name = "target.tests.prepare_recovery_test_target";
+
+    let mut engine = Engine::new();
+    engine.add_policy(
+        "test.rego".to_string(),
+        format!(
+            r#"package test
+           import rego.v1
+           __target__ := "{target_name}"
+
+           default allow := false
+
+           allow if {{
+               input.type == "test_resource"
+           }}
+        "#
+        ),
+    )?;
+
+    engine.prepare()?;
+    assert!(engine.compile_for_target().is_err());
+
+    if !registry::targets::contains(target_name) {
+        let target_json =
+            include_str!("../interpreter/cases/target/definitions/sample_target.json")
+                .replace("target.tests.sample_test_target", target_name);
+        let target = Target::from_json_str(&target_json)?;
+        registry::targets::register(Rc::new(target))?;
+    }
+
+    let compiled = engine.compile_for_target()?;
+    let info = compiled.get_policy_info()?;
+    assert_eq!(info.target_name.as_deref(), Some(target_name));
 
     let result = compiled.eval_with_input(Value::from_json_str(
         r#"{"name":"resource-1","type":"test_resource"}"#,