test(attestation): feature-gate the pre-launch-script rejection

The fixture's app-compose carries the pre-launch script that exported its
signer key, which production verification rejects as arbitrary root code, so
tests could never verify it. The new allow-pre-launch-script feature relaxes
that single field: init_script and bash_script stay rejected, and the
production wasm never enables it. The check takes the policy as an argument so
the strict behavior keeps its unit test in test builds too.

Author: pbeza

crates/attestation/Cargo.toml

@@ -7,7 +7,11 @@ edition = { workspace = true }
 [features]
 borsh-schema = ["borsh/unstable__schema", "tee-verifier-interface/borsh-schema"]
 dstack-conversions = ["dep:dstack-sdk-types"]
-test-utils = []
+test-utils = ["allow-pre-launch-script"]
+# Accepts an app-compose carrying a `pre_launch_script`, which production rejects as
+# arbitrary root code. The test fixture needs one: it is how the node's in-enclave signer
+# key is exported during collection. `init_script` and `bash_script` stay rejected.
+allow-pre-launch-script = []
 # Pulls in `dcap-qvl` for full local DCAP + post-DCAP verification. Meant for
 # off-chain callers; `mpc-contract` enables it today.
 # TODO(#3264): contract drops this once DCAP moves to the verifier contract.

crates/attestation/src/attestation.rs

@@ -409,11 +409,19 @@ impl DstackAttestation {
             && app_compose.local_key_provider_enabled
             && app_compose.allowed_envs.is_empty()
             && app_compose.no_instance_id
-            // Reject all three arbitrary-root-code fields. `pre_launch_script` and `init_script` run
-            // unconditionally; `bash_script` only runs when `runner == "bash"` (so the runner pin
-            // above already neutralizes it), but we reject it explicitly so the guarantee does not
-            // silently depend on that pin.
-            && app_compose.pre_launch_script.is_none()
+            && Self::scripts_absent(app_compose)
+    }
+
+    /// Rejects the arbitrary-root-code fields. `bash_script` only runs when
+    /// `runner == "bash"`, but is rejected explicitly so the guarantee does not depend on
+    /// the runner pin above.
+    fn scripts_absent(app_compose: &AppCompose) -> bool {
+        Self::scripts_absent_with(app_compose, cfg!(feature = "allow-pre-launch-script"))
+    }
+
+    /// Takes the policy as an argument so tests can assert both, whatever features are on.
+    fn scripts_absent_with(app_compose: &AppCompose, allow_pre_launch_script: bool) -> bool {
+        (allow_pre_launch_script || app_compose.pre_launch_script.is_none())
             && app_compose.init_script.is_none()
             && app_compose.bash_script.is_none()
     }
@@ -620,19 +628,58 @@ mod tests {
     }
 
     #[test]
-    fn validate_app_compose_config__rejects_present_pre_launch_script() {
+    fn scripts_absent_with__rejects_pre_launch_script_when_disallowed() {
+        // Asserts the production policy, which test builds relax for the fixture.
+
         // Given
         let app_compose = AppCompose {
             pre_launch_script: Some("echo pwn".to_string()),
             ..valid_app_compose()
         };
         // When
-        let result = DstackAttestation::validate_app_compose_config(&app_compose);
+        let result = DstackAttestation::scripts_absent_with(&app_compose, false);
 
         // Then
         assert!(!result)
     }
 
+    #[test]
+    fn scripts_absent_with__accepts_pre_launch_script_when_allowed() {
+        // Given
+        let app_compose = AppCompose {
+            pre_launch_script: Some("echo collecting fixtures".to_string()),
+            ..valid_app_compose()
+        };
+        // When
+        let result = DstackAttestation::scripts_absent_with(&app_compose, true);
+
+        // Then
+        assert!(result)
+    }
+
+    #[test]
+    fn scripts_absent_with__rejects_other_scripts_even_when_pre_launch_is_allowed() {
+        // The relaxation must stay scoped to `pre_launch_script`.
+
+        // Given
+        let with_init = AppCompose {
+            init_script: Some("echo pwn".to_string()),
+            ..valid_app_compose()
+        };
+        let with_bash = AppCompose {
+            bash_script: Some("echo pwn".to_string()),
+            ..valid_app_compose()
+        };
+
+        // When
+        let init_result = DstackAttestation::scripts_absent_with(&with_init, true);
+        let bash_result = DstackAttestation::scripts_absent_with(&with_bash, true);
+
+        // Then
+        assert!(!init_result);
+        assert!(!bash_result);
+    }
+
     #[test]
     fn validate_app_compose_config__rejects_present_init_script() {
         // `init_script` is arbitrary root code run before dockerd. It is

crates/contract/Cargo.toml

@@ -59,6 +59,9 @@ bench-contract-methods = []
 # length of a fan-out queue). Distinct from `bench-contract-methods` because these are
 # behavioral hooks, not gas-measurement hooks.
 sandbox-test-methods = []
+# For the sandbox wasm of tests that submit the Dstack fixture and expect it to verify;
+# see `attestation/allow-pre-launch-script`.
+sandbox-test-attestation = ["mpc-attestation/allow-pre-launch-script"]
 dev-utils = ["rand", "threshold-signatures", "near-mpc-contract-interface/blstrs"]
 abi = [
     "borsh/unstable__schema",

crates/mpc-attestation/Cargo.toml

@@ -8,6 +8,8 @@ edition = { workspace = true }
 abi = ["borsh/unstable__schema", "mpc-primitives/abi", "attestation/borsh-schema"]
 dstack-conversions = ["attestation/dstack-conversions"]
 test-utils = ["attestation/test-utils"]
+# Forwards `attestation/allow-pre-launch-script`; see that crate for the rationale.
+allow-pre-launch-script = ["attestation/allow-pre-launch-script"]
 # Enables full local DCAP + post-DCAP verification, forwarding to
 # `attestation/local-verify` which pulls in `dcap-qvl`. Used off-chain (node,
 # tee-authority, attestation-cli) and, for now, by the contract's synchronous