Merge "Refactor Dice policy test" into main
diff --git a/secretkeeper/dice_policy/src/lib.rs b/secretkeeper/dice_policy/src/lib.rs
index f5d117c..327b8a4 100644
--- a/secretkeeper/dice_policy/src/lib.rs
+++ b/secretkeeper/dice_policy/src/lib.rs
@@ -174,7 +174,7 @@
     payload: Value,
     constraint_spec: &[ConstraintSpec],
 ) -> Result<NodeConstraints> {
-    let mut node_constraints: Vec<Constraint> = Vec::new();
+    let mut node_constraints: Vec<Constraint> = Vec::with_capacity(constraint_spec.len());
     for constraint_item in constraint_spec {
         let constraint_path = constraint_item.path.to_vec();
         if constraint_path.is_empty() {
@@ -258,15 +258,98 @@
     const COMPONENT_NAME: i64 = -70002;
     const KEY_MODE: i64 = -4670551;
 
-    // This is the number of certs in compos bcc (including the first ROT)
-    // To analyze a bcc use hwtrust tool from /tools/security/remote_provisioning/hwtrust
-    // `hwtrust --verbose dice-chain [path]/composbcc`
-    const COMPOS_DICE_CHAIN_SIZE: usize = 5;
-    const EXAMPLE_STRING: &str = "testing_dice_policy";
-    const EXAMPLE_NUM: i64 = 59765;
+    // Helper struct to encapsulate artifacts that are useful for unit tests.
+    struct TestArtifacts {
+        // A dice chain.
+        input_dice: Vec<u8>,
+        // A list of ConstraintSpec that can be applied on the input_dice to get a dice policy.
+        constraint_spec: Vec<ConstraintSpec>,
+        // The expected dice policy if above constraint_spec is applied to input_dice.
+        expected_dice_policy: DicePolicy,
+    }
+
+    impl TestArtifacts {
+        // Get an example instance of TestArtifacts. This uses a hard coded, hypothetical
+        // chain of certificates & a list of constraint_spec on this.
+        fn get_example() -> Self {
+            const EXAMPLE_NUM: i64 = 59765;
+            const EXAMPLE_STRING: &str = "testing_dice_policy";
+
+            let rot_key = CoseKey::default().to_cbor_value().unwrap();
+            let nested_payload = cbor!({
+                100 => EXAMPLE_NUM
+            })
+            .unwrap();
+            let payload = cbor!({
+                1 => EXAMPLE_STRING,
+                2 => "some_other_example_string",
+                3 => Value::Bytes(value_to_bytes(&nested_payload).unwrap()),
+            })
+            .unwrap();
+            let payload = value_to_bytes(&payload).unwrap();
+            let dice_node = CoseSign1 {
+                protected: ProtectedHeader::default(),
+                unprotected: Header::default(),
+                payload: Some(payload),
+                signature: b"ddef".to_vec(),
+            }
+            .to_cbor_value()
+            .unwrap();
+            let input_dice = Value::Array([rot_key.clone(), dice_node].to_vec());
+
+            let input_dice = value_to_bytes(&input_dice).unwrap();
+
+            // Now construct constraint_spec on the input dice, note this will use the keys
+            // which are also hardcoded within the get_dice_chain_helper.
+
+            let constraint_spec = vec![
+                ConstraintSpec::new(ConstraintType::ExactMatch, vec![1]).unwrap(),
+                // Notice how key "2" is (deliberately) absent in ConstraintSpec
+                // so policy should not constraint it.
+                ConstraintSpec::new(ConstraintType::GreaterOrEqual, vec![3, 100]).unwrap(),
+            ];
+            let expected_dice_policy = DicePolicy {
+                version: 1,
+                node_constraints_list: Box::new([
+                    NodeConstraints(Box::new([Constraint(
+                        ConstraintType::ExactMatch as u16,
+                        vec![],
+                        rot_key.clone(),
+                    )])),
+                    NodeConstraints(Box::new([
+                        Constraint(
+                            ConstraintType::ExactMatch as u16,
+                            vec![1],
+                            Value::Text(EXAMPLE_STRING.to_string()),
+                        ),
+                        Constraint(
+                            ConstraintType::GreaterOrEqual as u16,
+                            vec![3, 100],
+                            Value::from(EXAMPLE_NUM),
+                        ),
+                    ])),
+                ]),
+            };
+            Self { input_dice, constraint_spec, expected_dice_policy }
+        }
+    }
+
+    test!(policy_structure_check);
+    fn policy_structure_check() {
+        let example = TestArtifacts::get_example();
+        let policy =
+            DicePolicy::from_dice_chain(&example.input_dice, &example.constraint_spec).unwrap();
+
+        // Assert policy is exactly as expected!
+        assert_eq!(policy, example.expected_dice_policy);
+    }
 
     test!(policy_dice_size_is_same);
     fn policy_dice_size_is_same() {
+        // This is the number of certs in compos bcc (including the first ROT)
+        // To analyze a bcc use hwtrust tool from /tools/security/remote_provisioning/hwtrust
+        // `hwtrust --verbose dice-chain [path]/composbcc`
+        let compos_dice_chain_size: usize = 5;
         let input_dice = include_bytes!("../testdata/composbcc");
         let constraint_spec = [
             ConstraintSpec::new(ConstraintType::ExactMatch, vec![AUTHORITY_HASH]).unwrap(),
@@ -275,66 +358,7 @@
                 .unwrap(),
         ];
         let policy = DicePolicy::from_dice_chain(input_dice, &constraint_spec).unwrap();
-        assert_eq!(policy.node_constraints_list.len(), COMPOS_DICE_CHAIN_SIZE);
-    }
-
-    test!(policy_structure_check);
-    fn policy_structure_check() {
-        let rot_key = CoseKey::default().to_cbor_value().unwrap();
-        let nested_payload = cbor!({
-            100 => EXAMPLE_NUM
-        })
-        .unwrap();
-        let payload = cbor!({
-            1 => EXAMPLE_STRING,
-            2 => "some_other_example_string",
-            3 => Value::Bytes(value_to_bytes(&nested_payload).unwrap()),
-        })
-        .unwrap();
-        let payload = value_to_bytes(&payload).unwrap();
-        let dice_node = CoseSign1 {
-            protected: ProtectedHeader::default(),
-            unprotected: Header::default(),
-            payload: Some(payload),
-            signature: b"ddef".to_vec(),
-        }
-        .to_cbor_value()
-        .unwrap();
-        let input_dice = Value::Array([rot_key.clone(), dice_node].to_vec());
-
-        let input_dice = value_to_bytes(&input_dice).unwrap();
-        let constraint_spec = [
-            ConstraintSpec::new(ConstraintType::ExactMatch, vec![1]).unwrap(),
-            ConstraintSpec::new(ConstraintType::GreaterOrEqual, vec![3, 100]).unwrap(),
-        ];
-        let policy = DicePolicy::from_dice_chain(&input_dice, &constraint_spec).unwrap();
-
-        // Assert policy is exactly as expected!
-        assert_eq!(
-            policy,
-            DicePolicy {
-                version: 1,
-                node_constraints_list: Box::new([
-                    NodeConstraints(Box::new([Constraint(
-                        ConstraintType::ExactMatch as u16,
-                        vec![],
-                        rot_key
-                    )])),
-                    NodeConstraints(Box::new([
-                        Constraint(
-                            ConstraintType::ExactMatch as u16,
-                            vec![1],
-                            Value::Text(EXAMPLE_STRING.to_string())
-                        ),
-                        Constraint(
-                            ConstraintType::GreaterOrEqual as u16,
-                            vec![3, 100],
-                            Value::from(EXAMPLE_NUM)
-                        )
-                    ])),
-                ])
-            }
-        );
+        assert_eq!(policy.node_constraints_list.len(), compos_dice_chain_size);
     }
 
     /// Encodes a ciborium::Value into bytes.