Integrate authorizations with the operations.

Bug: 171503362, 171503128
Test: TBD
Change-Id: If12104eec4f9f32a9af4f4da8e620543ce26548d
diff --git a/keystore2/src/auth_token_handler.rs b/keystore2/src/auth_token_handler.rs
index 8c10442..a1f9399 100644
--- a/keystore2/src/auth_token_handler.rs
+++ b/keystore2/src/auth_token_handler.rs
@@ -15,7 +15,6 @@
 //! This module defines the AuthTokenHandler enum and its methods. AuthTokenHandler enum represents
 //! the different states an auth token and an associated verification token can be expressed during
 //! the operation life cycle.
-use crate::error::Error as KeystoreError;
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     HardwareAuthToken::HardwareAuthToken, VerificationToken::VerificationToken,
 };
@@ -23,7 +22,8 @@
 use std::sync::mpsc::Receiver;
 
 /// AuthTokenHandler enum has five different variants which are described by the comments above
-// each variant.
+// each variant, as follows.
+#[derive(Debug)]
 pub enum AuthTokenHandler {
     /// Used when an operation does not require an auth token for authorization.
     NoAuthRequired,
@@ -45,39 +45,37 @@
 }
 
 impl AuthTokenHandler {
-    /// Retrieve auth token and verification token from the Token variant of an AuthTokenHandler
-    /// instance
-    pub fn get_auth_and_verification_tokens(
-        &self,
-    ) -> Option<(&HardwareAuthToken, &VerificationToken)> {
-        if let AuthTokenHandler::Token(auth_token, Some(verification_token)) = self {
-            Some((auth_token, verification_token))
-        } else {
-            None
-        }
-    }
-
-    /// Retrieve auth token from the Token variant of an AuthTokenHandler instance
-    pub fn get_auth_token(&self) -> Option<&HardwareAuthToken> {
-        if let AuthTokenHandler::Token(auth_token, _) = self {
-            Some(auth_token)
-        } else {
-            None
-        }
-    }
-
     /// If Channel variant, block on it until the verification token is sent by the
-    /// keystore2 worker thread which obtains verification tokens from TEE Keymint
-    pub fn receive_verification_token(&mut self) -> Result<()> {
+    /// keystore2 worker thread which obtains verification tokens from TEE Keymint and converts the
+    /// object from Channel variant to Token variant.
+    /// Retrieve auth token and verification token from the Token variant of an AuthTokenHandler
+    /// instance.
+    pub fn retrieve_auth_and_verification_tokens(
+        &mut self,
+    ) -> Result<(Option<&HardwareAuthToken>, Option<&VerificationToken>)> {
+        // Converts to Token variant if Channel variant found, after retrieving the
+        // VerificationToken
         if let AuthTokenHandler::Channel(recv) = self {
             let (auth_token, verification_token) =
                 recv.recv().context("In receive_verification_token: sender disconnected.")?;
             *self = AuthTokenHandler::Token(auth_token, Some(verification_token));
-            Ok(())
+        }
+        // get the tokens from the Token variant
+        if let AuthTokenHandler::Token(auth_token, optional_verification_token) = self {
+            Ok((Some(auth_token), optional_verification_token.as_ref()))
         } else {
-            Err(KeystoreError::sys()).context(
-                "In receive_verification_token: Wrong variant found in the authorization object.",
-            )
+            Ok((None, None))
+        }
+    }
+
+    /// Retrieve auth token from VerificationRequired and Token variants of an
+    /// AuthTokenHandler instance. This method is useful when we only expect an auth token and
+    /// do not expect a verification token.
+    pub fn get_auth_token(&self) -> Option<&HardwareAuthToken> {
+        match self {
+            AuthTokenHandler::VerificationRequired(auth_token) => Some(auth_token),
+            AuthTokenHandler::Token(auth_token, _) => Some(auth_token),
+            _ => None,
         }
     }
 }
diff --git a/keystore2/src/enforcements.rs b/keystore2/src/enforcements.rs
index 1f82e57..41a4e48 100644
--- a/keystore2/src/enforcements.rs
+++ b/keystore2/src/enforcements.rs
@@ -53,10 +53,13 @@
     /// try to receive the auth token from the op_auth_map. We assume that by the time update/finish
     /// is called, the auth token has been delivered to keystore. Therefore, we do not wait for it
     /// and if the auth token is not found in the map, an error is returned.
+    /// This method is called only during the first call to update or if finish is called right
+    /// after create operation, because the operation caches the authorization decisions and tokens
+    /// from previous calls to enforcement module.
     pub fn authorize_update_or_finish(
         &self,
         key_params: &[KeyParameter],
-        op_challenge: Option<OperationChallenge>,
+        op_challenge: Option<&OperationChallenge>,
     ) -> Result<AuthTokenHandler> {
         let mut user_auth_type: Option<HardwareAuthenticatorType> = None;
         let mut user_secure_ids = Vec::<i64>::new();
@@ -403,6 +406,15 @@
             .context("In add_auth_token.")?;
         Ok(())
     }
+
+    /// This allows adding an entry to the op_auth_map, indexed by the operation challenge.
+    /// This is to be called by create_operation, once it has received the operation challenge
+    /// from keymint for an operation whose authorization decision is OpAuthRequired, as signalled
+    /// by the AuthTokenHandler.
+    pub fn insert_to_op_auth_map(&self, op_challenge: i64) {
+        let mut op_auth_map_guard = self.op_auth_map.lock().unwrap();
+        op_auth_map_guard.insert(op_challenge, None);
+    }
 }
 
 impl Default for Enforcements {
diff --git a/keystore2/src/globals.rs b/keystore2/src/globals.rs
index 18bed9a..075a801 100644
--- a/keystore2/src/globals.rs
+++ b/keystore2/src/globals.rs
@@ -17,6 +17,7 @@
 //! to talk to.
 
 use crate::async_task::AsyncTask;
+use crate::enforcements::Enforcements;
 use crate::gc::Gc;
 use crate::super_key::SuperKeyManager;
 use crate::utils::Asp;
@@ -81,6 +82,10 @@
     /// A single on-demand worker thread that handles deferred tasks with two different
     /// priorities.
     pub static ref ASYNC_TASK: AsyncTask = Default::default();
+    /// Singeleton for enforcements.
+    /// It is safe for this enforcements object to be called by multiple threads because the two
+    /// data structures which maintain its state are protected by mutexes.
+    pub static ref ENFORCEMENTS: Enforcements = Enforcements::new();
 }
 
 static KEYMINT_SERVICE_NAME: &str = "android.hardware.security.keymint.IKeyMintDevice";
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
             }
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 1d0608e..c5b5da0 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -17,19 +17,23 @@
 //! This crate implements the IKeystoreSecurityLevel interface.
 
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    Algorithm::Algorithm, HardwareAuthenticatorType::HardwareAuthenticatorType,
-    IKeyMintDevice::IKeyMintDevice, KeyCreationResult::KeyCreationResult, KeyFormat::KeyFormat,
-    KeyParameter::KeyParameter, KeyParameterValue::KeyParameterValue, SecurityLevel::SecurityLevel,
-    Tag::Tag,
+    Algorithm::Algorithm, HardwareAuthToken::HardwareAuthToken,
+    HardwareAuthenticatorType::HardwareAuthenticatorType, IKeyMintDevice::IKeyMintDevice,
+    KeyCreationResult::KeyCreationResult, KeyFormat::KeyFormat, KeyParameter::KeyParameter,
+    KeyParameterValue::KeyParameterValue, SecurityLevel::SecurityLevel, Tag::Tag,
+    VerificationToken::VerificationToken,
 };
 use android_system_keystore2::aidl::android::system::keystore2::{
     AuthenticatorSpec::AuthenticatorSpec, CreateOperationResponse::CreateOperationResponse,
     Domain::Domain, IKeystoreOperation::IKeystoreOperation,
     IKeystoreSecurityLevel::BnKeystoreSecurityLevel,
     IKeystoreSecurityLevel::IKeystoreSecurityLevel, KeyDescriptor::KeyDescriptor,
-    KeyMetadata::KeyMetadata, KeyParameters::KeyParameters,
+    KeyMetadata::KeyMetadata, KeyParameters::KeyParameters, OperationChallenge::OperationChallenge,
 };
 
+use crate::auth_token_handler::AuthTokenHandler;
+use crate::globals::ENFORCEMENTS;
+use crate::key_parameter::KeyParameter as KsKeyParam;
 use crate::utils::{check_key_permission, Asp};
 use crate::{database::KeyIdGuard, globals::DB};
 use crate::{
@@ -47,6 +51,7 @@
 };
 use anyhow::{Context, Result};
 use binder::{IBinder, Interface, ThreadState};
+use std::sync::mpsc::channel;
 
 /// Implementation of the IKeystoreSecurityLevel Interface.
 pub struct KeystoreSecurityLevel {
@@ -159,7 +164,7 @@
         // so that we can use it by reference like the blob provided by the key descriptor.
         // Otherwise, we would have to clone the blob from the key descriptor.
         let scoping_blob: Vec<u8>;
-        let (km_blob, key_id_guard) = match key.domain {
+        let (km_blob, key_id_guard, key_parameters) = match key.domain {
             Domain::BLOB => {
                 check_key_permission(KeyPerm::use_(), key, &None)
                     .context("In create_operation: checking use permission for Domain::BLOB.")?;
@@ -174,6 +179,7 @@
                         }
                     },
                     None,
+                    None,
                 )
             }
             _ => {
@@ -197,14 +203,10 @@
                         ))
                     }
                 };
-                (&scoping_blob, Some(key_id_guard))
+                (&scoping_blob, Some(key_id_guard), Some(key_entry.into_key_parameters()))
             }
         };
 
-        // TODO Authorize begin operation.
-        // Check if we need an authorization token.
-        // Lookup authorization token and request VerificationToken if required.
-
         let purpose = operation_parameters.iter().find(|p| p.tag == Tag::PURPOSE).map_or(
             Err(Error::Km(ErrorCode::INVALID_ARGUMENT))
                 .context("In create_operation: No operation purpose specified."),
@@ -215,6 +217,35 @@
             },
         )?;
 
+        let mut auth_token_for_km: &HardwareAuthToken = &Default::default();
+        let mut auth_token_handler = AuthTokenHandler::NoAuthRequired;
+
+        // keystore performs authorizations only if the key parameters are loaded above
+        if let Some(ref key_params) = key_parameters {
+            // Note: although currently only one operation parameter is checked in authorizing the
+            // operation, the whole operation_parameter vector is converted into the internal
+            // representation of key parameter because we might need to sanitize operation
+            // parameters (b/175792701)
+            let mut op_params: Vec<KsKeyParam> = Vec::new();
+            for op_param in operation_parameters.iter() {
+                op_params.push(KsKeyParam::new(op_param.into(), self.security_level));
+            }
+            // authorize the operation, and receive an AuthTokenHandler, if authorized, else
+            // propagate the error
+            auth_token_handler = ENFORCEMENTS
+                .authorize_create(
+                    purpose,
+                    key_params.as_slice(),
+                    op_params.as_slice(),
+                    self.security_level,
+                )
+                .context("In create_operation.")?;
+            // if an auth token was found, pass it to keymint
+            if let Some(auth_token) = auth_token_handler.get_auth_token() {
+                auth_token_for_km = auth_token;
+            }
+        }
+
         let km_dev: Box<dyn IKeyMintDevice> = self
             .keymint
             .get_interface()
@@ -231,7 +262,7 @@
                         purpose,
                         blob,
                         &operation_parameters,
-                        &Default::default(),
+                        auth_token_for_km,
                     )) {
                         Err(Error::Km(ErrorCode::TOO_MANY_OPERATIONS)) => {
                             self.operation_db.prune(caller_uid)?;
@@ -243,8 +274,32 @@
             )
             .context("In create_operation: Failed to begin operation.")?;
 
+        let mut operation_challenge: Option<OperationChallenge> = None;
+
+        // take actions based on the authorization decision (if any) received via auth token handler
+        match auth_token_handler {
+            AuthTokenHandler::OpAuthRequired => {
+                operation_challenge =
+                    Some(OperationChallenge { challenge: begin_result.challenge });
+                ENFORCEMENTS.insert_to_op_auth_map(begin_result.challenge);
+            }
+            AuthTokenHandler::VerificationRequired(auth_token) => {
+                let (_sender, receiver) = channel::<(HardwareAuthToken, VerificationToken)>();
+                //TODO: call the worker thread and hand over the sender, auth token and challenge
+                auth_token_handler = AuthTokenHandler::Channel(receiver);
+            }
+            _ => {}
+        }
+
         let operation = match begin_result.operation {
-            Some(km_op) => self.operation_db.create_operation(km_op, caller_uid),
+            Some(km_op) => {
+                let mut op_challenge_copy: Option<OperationChallenge> = None;
+                if let Some(ref op_challenge) = operation_challenge {
+                    op_challenge_copy = Some(OperationChallenge{challenge: op_challenge.challenge});
+                }
+                self.operation_db.create_operation(km_op, caller_uid,
+                 auth_token_handler, key_parameters, op_challenge_copy)
+            },
             None => return Err(Error::sys()).context("In create_operation: Begin operation returned successfully, but did not return a valid operation."),
         };
 
@@ -258,7 +313,7 @@
         // We return None for now because we don't support auth bound keys yet.
         Ok(CreateOperationResponse {
             iOperation: Some(op_binder),
-            operationChallenge: None,
+            operationChallenge: operation_challenge,
             parameters: match begin_result.params.len() {
                 0 => None,
                 _ => Some(KeyParameters { keyParameter: begin_result.params }),