Integrate authorizations with the operations.

Bug: 171503362, 171503128
Test: TBD
Change-Id: If12104eec4f9f32a9af4f4da8e620543ce26548d
diff --git a/keystore2/src/operation.rs b/keystore2/src/operation.rs
index 9a35154..f306df4 100644
--- a/keystore2/src/operation.rs
+++ b/keystore2/src/operation.rs
@@ -125,6 +125,23 @@
 //! or it transitions to its end-of-life, which means we may get a free slot.
 //! Either way, we have to revaluate the pruning scores.
 
+use crate::auth_token_handler::AuthTokenHandler;
+use crate::error::{map_km_error, map_or_log_err, Error, ErrorCode, ResponseCode};
+use crate::globals::ENFORCEMENTS;
+use crate::key_parameter::KeyParameter;
+use crate::utils::Asp;
+use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+    ByteArray::ByteArray, HardwareAuthToken::HardwareAuthToken,
+    IKeyMintOperation::IKeyMintOperation, KeyParameter::KeyParameter as KmParam,
+    KeyParameterArray::KeyParameterArray, KeyParameterValue::KeyParameterValue as KmParamValue,
+    Tag::Tag, VerificationToken::VerificationToken,
+};
+use android_system_keystore2::aidl::android::system::keystore2::{
+    IKeystoreOperation::BnKeystoreOperation, IKeystoreOperation::IKeystoreOperation,
+    OperationChallenge::OperationChallenge,
+};
+use anyhow::{anyhow, Context, Result};
+use binder::{IBinder, Interface};
 use std::{
     collections::HashMap,
     sync::{Arc, Mutex, MutexGuard, Weak},
@@ -132,19 +149,6 @@
     time::Instant,
 };
 
-use crate::error::{map_km_error, map_or_log_err, Error, ErrorCode, ResponseCode};
-use crate::utils::Asp;
-use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    ByteArray::ByteArray, IKeyMintOperation::IKeyMintOperation,
-    KeyParameter::KeyParameter as KmParam, KeyParameterArray::KeyParameterArray,
-    KeyParameterValue::KeyParameterValue as KmParamValue, Tag::Tag,
-};
-use android_system_keystore2::aidl::android::system::keystore2::{
-    IKeystoreOperation::BnKeystoreOperation, IKeystoreOperation::IKeystoreOperation,
-};
-use anyhow::{anyhow, Context, Result};
-use binder::{IBinder, Interface};
-
 /// Operations have `Outcome::Unknown` as long as they are active. They transition
 /// to one of the other variants exactly once. The distinction in outcome is mainly
 /// for the statistic.
@@ -168,6 +172,11 @@
     last_usage: Mutex<Instant>,
     outcome: Mutex<Outcome>,
     owner: u32, // Uid of the operation's owner.
+    auth_token_handler: Mutex<AuthTokenHandler>,
+    // optional because in create_operation, there is a case in which we might not load
+    // key parameters
+    key_params: Option<Vec<KeyParameter>>,
+    op_challenge: Option<OperationChallenge>,
 }
 
 struct PruningInfo {
@@ -181,13 +190,23 @@
 
 impl Operation {
     /// Constructor
-    pub fn new(index: usize, km_op: Box<dyn IKeyMintOperation>, owner: u32) -> Self {
+    pub fn new(
+        index: usize,
+        km_op: Box<dyn IKeyMintOperation>,
+        owner: u32,
+        auth_token_handler: AuthTokenHandler,
+        key_params: Option<Vec<KeyParameter>>,
+        op_challenge: Option<OperationChallenge>,
+    ) -> Self {
         Self {
             index,
             km_op: Asp::new(km_op.as_binder()),
             last_usage: Mutex::new(Instant::now()),
             outcome: Mutex::new(Outcome::Unknown),
             owner,
+            auth_token_handler: Mutex::new(auth_token_handler),
+            key_params,
+            op_challenge,
         }
     }
 
@@ -348,6 +367,43 @@
         Ok(())
     }
 
+    /// Based on the authorization information stored in the operation during create_operation(),
+    /// and any previous calls to update(), this function returns appropriate auth token and
+    /// verification token to be passed to keymint.
+    /// Note that the call to the global enforcement object happens only during the first call to
+    /// update or if finish() is called right after create_opertation.
+    fn handle_authorization<'a>(
+        auth_token_handler: &'a mut AuthTokenHandler,
+        key_params: Option<&Vec<KeyParameter>>,
+        op_challenge: Option<&OperationChallenge>,
+    ) -> Result<(Option<&'a HardwareAuthToken>, Option<&'a VerificationToken>)> {
+        // keystore performs authorization only if key parameters have been loaded during
+        // create_operation()
+        if let Some(key_parameters) = key_params {
+            match *auth_token_handler {
+                // this variant is found only in a first call to update or if finish is called
+                // right after create_operation.
+                AuthTokenHandler::OpAuthRequired => {
+                    *auth_token_handler = ENFORCEMENTS
+                        .authorize_update_or_finish(key_parameters.as_slice(), op_challenge)
+                        .context("In handle_authorization.")?;
+                    Ok((auth_token_handler.get_auth_token(), None))
+                }
+                // this variant is found only in a first call to update or if finish is called
+                // right after create_operation.
+                AuthTokenHandler::Channel(_)|
+                // this variant is found in every subsequent call to update/finish,
+                // unless the authorization is not required for the key
+                AuthTokenHandler::Token(_, _) => {
+                    auth_token_handler.retrieve_auth_and_verification_tokens()
+                }
+                _ => Ok((None, None))
+            }
+        } else {
+            Ok((None, None))
+        }
+    }
+
     /// Implementation of `IKeystoreOperation::update`.
     /// Refer to the AIDL spec at system/hardware/interfaces/keystore2 for details.
     fn update(&self, input: &[u8]) -> Result<Option<Vec<u8>>> {
@@ -361,15 +417,21 @@
         let km_op: Box<dyn IKeyMintOperation> =
             self.km_op.get_interface().context("In update: Failed to get KeyMintOperation.")?;
 
+        let mut auth_handler = self.auth_token_handler.lock().unwrap();
+        let (auth_token_for_km, verification_token_for_km) = Self::handle_authorization(
+            &mut auth_handler,
+            self.key_params.as_ref(),
+            self.op_challenge.as_ref(),
+        )
+        .context("In update.")?;
+
         self.update_outcome(
             &mut *outcome,
             map_km_error(km_op.update(
                 None,
                 Some(input),
-                // TODO Get auth token from enforcement module if required.
-                None,
-                // TODO Get verification token from enforcement module if required.
-                None,
+                auth_token_for_km,
+                verification_token_for_km,
                 &mut out_params,
                 &mut output,
             )),
@@ -402,6 +464,14 @@
         let km_op: Box<dyn IKeyMintOperation> =
             self.km_op.get_interface().context("In finish: Failed to get KeyMintOperation.")?;
 
+        let mut auth_handler = self.auth_token_handler.lock().unwrap();
+        let (auth_token_for_km, verification_token_for_km) = Self::handle_authorization(
+            &mut auth_handler,
+            self.key_params.as_ref(),
+            self.op_challenge.as_ref(),
+        )
+        .context("In finish.")?;
+
         let output = self
             .update_outcome(
                 &mut *outcome,
@@ -409,10 +479,8 @@
                     None,
                     input,
                     signature,
-                    // TODO Get auth token from enforcement module if required.
-                    None,
-                    // TODO Get verification token from enforcement module if required.
-                    None,
+                    auth_token_for_km,
+                    verification_token_for_km,
                     &mut out_params,
                 )),
             )
@@ -475,6 +543,9 @@
         &self,
         km_op: Box<dyn IKeyMintOperation>,
         owner: u32,
+        auth_token_handler: AuthTokenHandler,
+        key_params: Option<Vec<KeyParameter>>,
+        op_challenge: Option<OperationChallenge>,
     ) -> Arc<Operation> {
         // We use unwrap because we don't allow code that can panic while locked.
         let mut operations = self.operations.lock().expect("In create_operation.");
@@ -487,12 +558,26 @@
             s.upgrade().is_none()
         }) {
             Some(free_slot) => {
-                let new_op = Arc::new(Operation::new(index - 1, km_op, owner));
+                let new_op = Arc::new(Operation::new(
+                    index - 1,
+                    km_op,
+                    owner,
+                    auth_token_handler,
+                    key_params,
+                    op_challenge,
+                ));
                 *free_slot = Arc::downgrade(&new_op);
                 new_op
             }
             None => {
-                let new_op = Arc::new(Operation::new(operations.len(), km_op, owner));
+                let new_op = Arc::new(Operation::new(
+                    operations.len(),
+                    km_op,
+                    owner,
+                    auth_token_handler,
+                    key_params,
+                    op_challenge,
+                ));
                 operations.push(Arc::downgrade(&new_op));
                 new_op
             }