Update the VTS module according to the refactoring in aosp/2826571

Test: atest VtsAidlAuthGraphRoleTest, atest VtsAidlAuthGraphSessionTest
Change-Id: I6cc3bd17952f602b58668d35e09c6a5385c7de61
diff --git a/security/authgraph/aidl/vts/functional/source.rs b/security/authgraph/aidl/vts/functional/source.rs
index 9aaaaee..4178a99 100644
--- a/security/authgraph/aidl/vts/functional/source.rs
+++ b/security/authgraph/aidl/vts/functional/source.rs
@@ -16,19 +16,22 @@
 
 //! VTS tests for sources
 use super::*;
-use authgraph_core::traits;
+use authgraph_core::{key, keyexchange as ke};
 
 /// Run AuthGraph tests against the provided source, using a local test sink implementation.
-pub fn test(impls: &mut traits::TraitImpl, source: binder::Strong<dyn IAuthGraphKeyExchange>) {
-    test_mainline(impls, source.clone());
-    test_corrupt_sig(impls, source.clone());
-    test_corrupt_key(impls, source);
+pub fn test(
+    local_sink: &mut ke::AuthGraphParticipant,
+    source: binder::Strong<dyn IAuthGraphKeyExchange>,
+) {
+    test_mainline(local_sink, source.clone());
+    test_corrupt_sig(local_sink, source.clone());
+    test_corrupt_key(local_sink, source);
 }
 
 /// Perform mainline AuthGraph key exchange with the provided source.
 /// Return the agreed AES keys in plaintext.
 pub fn test_mainline(
-    impls: &mut traits::TraitImpl,
+    local_sink: &mut ke::AuthGraphParticipant,
     source: binder::Strong<dyn IAuthGraphKeyExchange>,
 ) -> [key::AesKey; 2] {
     // Step 1: create an ephemeral ECDH key at the (remote) source.
@@ -40,14 +43,14 @@
     let source_pub_key = extract_plain_pub_key(&source_init_info.key.pubKey);
 
     // Step 2: pass the source's ECDH public key and other session info to the (local) sink.
-    let init_result = ke::init(
-        impls,
-        &source_pub_key.plainPubKey,
-        &source_init_info.identity.identity,
-        &source_init_info.nonce,
-        source_init_info.version,
-    )
-    .expect("failed to init() with local impl");
+    let init_result = local_sink
+        .init(
+            &source_pub_key.plainPubKey,
+            &source_init_info.identity.identity,
+            &source_init_info.nonce,
+            source_init_info.version,
+        )
+        .expect("failed to init() with local impl");
     let sink_init_info = init_result.session_init_info;
     let sink_pub_key = sink_init_info
         .ke_key
@@ -58,14 +61,17 @@
     assert!(!sink_info.session_id.is_empty());
 
     // The AuthGraph core library will verify the session ID signature, but do it here too.
-    let sink_verification_key = verification_key_from_identity(&impls, &sink_init_info.identity);
-    ke::verify_signature_on_session_id(
-        &sink_verification_key,
-        &sink_info.session_id,
-        &sink_info.session_id_signature,
-        &*impls.ecdsa,
-    )
-    .expect("failed verification of signed session ID");
+    let sink_verification_key = key::Identity::from_slice(&sink_init_info.identity)
+        .expect("invalid identity CBOR")
+        .cert_chain
+        .root_key;
+    local_sink
+        .verify_signature_on_session_id(
+            &sink_verification_key,
+            &sink_info.session_id,
+            &sink_info.session_id_signature,
+        )
+        .expect("failed verification of signed session ID");
 
     // Step 3: pass the sink's ECDH public key and other session info to the (remote) source, so it
     // can calculate the same pair of symmetric keys.
@@ -86,36 +92,41 @@
     assert!(!source_info.sessionId.is_empty());
 
     // The AuthGraph core library will verify the session ID signature, but do it here too.
-    let source_verification_key =
-        verification_key_from_identity(&impls, &source_init_info.identity.identity);
-    ke::verify_signature_on_session_id(
-        &source_verification_key,
-        &source_info.sessionId,
-        &source_info.signature.signature,
-        &*impls.ecdsa,
-    )
-    .expect("failed verification of signed session ID");
+    let source_verification_key = local_sink
+        .peer_verification_key_from_identity(&source_init_info.identity.identity)
+        .expect("failed to get peer verification from identity");
+    local_sink
+        .verify_signature_on_session_id(
+            &source_verification_key,
+            &source_info.sessionId,
+            &source_info.signature.signature,
+        )
+        .expect("failed verification of signed session ID");
 
     // Both ends should agree on the session ID.
     assert_eq!(source_info.sessionId, sink_info.session_id);
 
     // Step 4: pass the (remote) source's session ID signature back to the sink, so it can check it
     // and update the symmetric keys so they're marked as authentication complete.
-    let sink_arcs = ke::authentication_complete(
-        impls,
-        &source_info.signature.signature,
-        sink_info.shared_keys,
-    )
-    .expect("failed to authenticationComplete() with local sink");
-
+    let sink_arcs = local_sink
+        .authentication_complete(&source_info.signature.signature, sink_info.shared_keys)
+        .expect("failed to authenticationComplete() with local sink");
     // Decrypt and return the session keys.
-    decipher_aes_keys(&impls, &sink_arcs)
+    let decrypted_shared_keys = local_sink
+        .decipher_shared_keys_from_arcs(&sink_arcs)
+        .expect("failed to decrypt shared key arcs")
+        .try_into();
+    let decrypted_shared_keys_array = match decrypted_shared_keys {
+        Ok(array) => array,
+        Err(_) => panic!("wrong number of decrypted shared key arcs"),
+    };
+    decrypted_shared_keys_array
 }
 
 /// Perform mainline AuthGraph key exchange with the provided source, but provide an invalid session
 /// ID signature.
 pub fn test_corrupt_sig(
-    impls: &mut traits::TraitImpl,
+    local_sink: &mut ke::AuthGraphParticipant,
     source: binder::Strong<dyn IAuthGraphKeyExchange>,
 ) {
     // Step 1: create an ephemeral ECDH key at the (remote) source.
@@ -127,14 +138,14 @@
     let source_pub_key = extract_plain_pub_key(&source_init_info.key.pubKey);
 
     // Step 2: pass the source's ECDH public key and other session info to the (local) sink.
-    let init_result = ke::init(
-        impls,
-        &source_pub_key.plainPubKey,
-        &source_init_info.identity.identity,
-        &source_init_info.nonce,
-        source_init_info.version,
-    )
-    .expect("failed to init() with local impl");
+    let init_result = local_sink
+        .init(
+            &source_pub_key.plainPubKey,
+            &source_init_info.identity.identity,
+            &source_init_info.nonce,
+            source_init_info.version,
+        )
+        .expect("failed to init() with local impl");
     let sink_init_info = init_result.session_init_info;
     let sink_pub_key = sink_init_info
         .ke_key
@@ -172,7 +183,7 @@
 /// Perform mainline AuthGraph key exchange with the provided source, but give it back
 /// a corrupted key.
 pub fn test_corrupt_key(
-    impls: &mut traits::TraitImpl,
+    local_sink: &mut ke::AuthGraphParticipant,
     source: binder::Strong<dyn IAuthGraphKeyExchange>,
 ) {
     // Step 1: create an ephemeral ECDH key at the (remote) source.
@@ -184,14 +195,14 @@
     let source_pub_key = extract_plain_pub_key(&source_init_info.key.pubKey);
 
     // Step 2: pass the source's ECDH public key and other session info to the (local) sink.
-    let init_result = ke::init(
-        impls,
-        &source_pub_key.plainPubKey,
-        &source_init_info.identity.identity,
-        &source_init_info.nonce,
-        source_init_info.version,
-    )
-    .expect("failed to init() with local impl");
+    let init_result = local_sink
+        .init(
+            &source_pub_key.plainPubKey,
+            &source_init_info.identity.identity,
+            &source_init_info.nonce,
+            source_init_info.version,
+        )
+        .expect("failed to init() with local impl");
     let sink_init_info = init_result.session_init_info;
     let sink_pub_key = sink_init_info
         .ke_key
@@ -202,14 +213,17 @@
     assert!(!sink_info.session_id.is_empty());
 
     // The AuthGraph core library will verify the session ID signature, but do it here too.
-    let sink_verification_key = verification_key_from_identity(&impls, &sink_init_info.identity);
-    ke::verify_signature_on_session_id(
-        &sink_verification_key,
-        &sink_info.session_id,
-        &sink_info.session_id_signature,
-        &*impls.ecdsa,
-    )
-    .expect("failed verification of signed session ID");
+    let sink_verification_key = key::Identity::from_slice(&sink_init_info.identity)
+        .expect("invalid identity CBOR")
+        .cert_chain
+        .root_key;
+    local_sink
+        .verify_signature_on_session_id(
+            &sink_verification_key,
+            &sink_info.session_id,
+            &sink_info.session_id_signature,
+        )
+        .expect("failed verification of signed session ID");
 
     // Deliberately corrupt the source's encrypted key.
     let mut corrupt_key = source_init_info.key.clone();