1 files changed, 53 insertions, 0 deletions
diff --git a/security/dice/aidl/vts/functional/utils.rs b/security/dice/aidl/vts/functional/utils.rs
new file mode 100644
index 0000000000..4e6708ec5e
--- /dev/null
+++ b/security/dice/aidl/vts/functional/utils.rs
@@ -0,0 +1,53 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use android_hardware_security_dice::aidl::android::hardware::security::dice::IDiceDevice::IDiceDevice;
+use anyhow::Result;
+use binder::Strong;
+use keystore2_vintf::get_aidl_instances;
+use std::sync::Arc;
+
+static DICE_DEVICE_SERVICE_NAME: &str = &"android.hardware.security.dice";
+static DICE_DEVICE_INTERFACE_NAME: &str = &"IDiceDevice";
+
+/// This function iterates through all announced IDiceDevice services and runs the given test
+/// closure against connections to each of them. It also modifies the panic hook to indicate
+/// on which instance the test failed in case the test closure panics.
+pub fn with_connection<R, F>(test: F)
+where
+    F: Fn(&Strong<dyn IDiceDevice>) -> Result<R>,
+{
+    let instances = get_aidl_instances(DICE_DEVICE_SERVICE_NAME, 1, DICE_DEVICE_INTERFACE_NAME);
+    let panic_hook = Arc::new(std::panic::take_hook());
+    for i in instances.into_iter() {
+        let panic_hook_clone = panic_hook.clone();
+        let instance_clone = i.clone();
+        std::panic::set_hook(Box::new(move |v| {
+            println!("While testing instance: \"{}\"", instance_clone);
+            panic_hook_clone(v)
+        }));
+        let connection: Strong<dyn IDiceDevice> = binder::get_interface(&format!(
+            "{}.{}/{}",
+            DICE_DEVICE_SERVICE_NAME, DICE_DEVICE_INTERFACE_NAME, i
+        ))
+        .unwrap();
+        test(&connection).unwrap();
+        drop(std::panic::take_hook());
+    }
+    // Cannot call unwrap here because the panic hook is not Debug.
+    std::panic::set_hook(match Arc::try_unwrap(panic_hook) {
+        Ok(hook) => hook,
+        _ => panic!("Failed to unwrap and reset previous panic hook."),
+    })
+}