Merge changes I2ed07420,If12104ee

* changes:
  Implement and integrate the verification token handler.
  Integrate authorizations with the operations.
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/background_task_handler.rs b/keystore2/src/background_task_handler.rs
new file mode 100644
index 0000000..fbb6778
--- /dev/null
+++ b/keystore2/src/background_task_handler.rs
@@ -0,0 +1,144 @@
+// Copyright 2020, 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.
+
+//! This module implements the handling of background tasks such as obtaining timestamp tokens from
+//! the timestamp service (or TEE KeyMint in legacy devices), via a separate thread.
+
+use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+    HardwareAuthToken::HardwareAuthToken, IKeyMintDevice::IKeyMintDevice,
+    SecurityLevel::SecurityLevel, VerificationToken::VerificationToken,
+};
+use android_system_keystore2::aidl::android::system::keystore2::OperationChallenge::OperationChallenge;
+use anyhow::Result;
+use log::error;
+use std::sync::mpsc::{Receiver, Sender};
+use std::sync::Mutex;
+use std::thread::{spawn, JoinHandle};
+/// This is the struct encapsulating the thread which handles background tasks such as
+/// obtaining verification tokens.
+pub struct BackgroundTaskHandler {
+    task_handler: Mutex<Option<JoinHandle<()>>>,
+}
+/// This enum defines the two variants of a message that can be passed down to the
+/// BackgroundTaskHandler via the channel.
+pub enum Message {
+    ///This variant represents a message sent down the channel when requesting a timestamp token.
+    Inputs((HardwareAuthToken, OperationChallenge, Sender<(HardwareAuthToken, VerificationToken)>)),
+    ///This variant represents a message sent down the channel when signalling the thread to stop.
+    Shutdown,
+}
+
+impl BackgroundTaskHandler {
+    /// Initialize the BackgroundTaskHandler with the task_handler field set to None.
+    /// The thread is not started during initialization, as it needs the receiver end of a channel
+    /// to function.
+    pub fn new() -> Self {
+        BackgroundTaskHandler { task_handler: Mutex::new(None) }
+    }
+
+    /// Start the background task handler (bth) by passing in the receiver end of a channel, through
+    /// which the enforcement module can send messages to the bth thread.
+    pub fn start_bth(&self, receiver: Receiver<Message>) -> Result<()> {
+        let task_handler = Self::start_thread(receiver)?;
+        // it is ok to unwrap here because there is no way that this lock can get poisoned.
+        let mut thread_guard = self.task_handler.lock().unwrap();
+        *thread_guard = Some(task_handler);
+        Ok(())
+    }
+
+    fn start_thread(receiver: Receiver<Message>) -> Result<JoinHandle<()>> {
+        // TODO: initialize timestamp service/keymint instances.
+        // First lookup timestamp token service, if this is not a legacy device.
+        // Otherwise, lookup keymaster 4.1 or 4.0 (previous ones are not relevant, because strongbox
+        // was introduced with keymaster 4.0).
+        // If either a timestamp service or a keymint instance is expected to be found and neither
+        // is found, an error is returned.
+        // If neither is expected to be found, make timestamp_service field None, and in the thread,
+        // send a default verification token down the channel to the operation.
+        // Until timestamp service is available and proper probing of legacy keymaster devices are
+        // done, the keymint service is initialized here as it is done in security_level module.
+        Ok(spawn(move || {
+            while let Message::Inputs((auth_token, op_challenge, op_sender)) = receiver
+                .recv()
+                .expect(
+                "In background task handler thread. Failed to receive message over the channel.",
+            ) {
+                // TODO: call the timestamp service/old TEE keymaster to get
+                // timestamp/verification tokens and pass it down the sender that is
+                // coupled with a particular operation's receiver.
+                // If none of the services are available, pass the authtoken and a default
+                // verification token down the channel.
+                let km_dev: Box<dyn IKeyMintDevice> =
+                    crate::globals::get_keymint_device(SecurityLevel::TRUSTED_ENVIRONMENT)
+                        .expect("A TEE Keymint must be present.")
+                        .get_interface()
+                        .expect("Fatal: The keymint device does not implement IKeyMintDevice.");
+                let result = km_dev.verifyAuthorization(op_challenge.challenge, &auth_token);
+                match result {
+                    Ok(verification_token) => {
+                        // this can fail if the operation is dropped and hence the channel
+                        // is hung up.
+                        op_sender.send((auth_token, verification_token)).unwrap_or_else(|e| {
+                            error!(
+                                "In background task handler thread. Failed to send
+                                         verification token to operation {} due to error {:?}.",
+                                op_challenge.challenge, e
+                            )
+                        });
+                    }
+                    Err(e) => {
+                        // log error
+                        error!(
+                            "In background task handler thread. Failed to receive
+                                     verification token for operation {} due to error {:?}.",
+                            op_challenge.challenge, e
+                        );
+                        // send default verification token
+                        // this can fail if the operation is dropped and the channel is
+                        // hung up.
+                        op_sender.send((auth_token, VerificationToken::default())).unwrap_or_else(
+                            |e| {
+                                error!(
+                                    "In background task handler thread. Failed to send default
+                                         verification token to operation {} due to error {:?}.",
+                                    op_challenge.challenge, e
+                                )
+                            },
+                        );
+                    }
+                }
+            }
+        }))
+    }
+}
+
+impl Default for BackgroundTaskHandler {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+// TODO: Verify if we want the thread to finish the requests they are working on, during drop.
+impl Drop for BackgroundTaskHandler {
+    fn drop(&mut self) {
+        // it is ok to unwrap here as there is no way this lock can get poisoned.
+        let mut thread_guard = self.task_handler.lock().unwrap();
+        if let Some(thread) = (*thread_guard).take() {
+            // TODO: Verify how best to handle the error in this case.
+            thread.join().unwrap_or_else(|e| {
+                panic!("Failed to join the background task handling thread because of {:?}.", e);
+            });
+        }
+    }
+}
diff --git a/keystore2/src/enforcements.rs b/keystore2/src/enforcements.rs
index 1f82e57..ae41432 100644
--- a/keystore2/src/enforcements.rs
+++ b/keystore2/src/enforcements.rs
@@ -15,6 +15,7 @@
 //! This is the Keystore 2.0 Enforcements module.
 // TODO: more description to follow.
 use crate::auth_token_handler::AuthTokenHandler;
+use crate::background_task_handler::Message;
 use crate::database::AuthTokenEntry;
 use crate::error::Error as KeystoreError;
 use crate::globals::DB;
@@ -23,10 +24,12 @@
     Algorithm::Algorithm, ErrorCode::ErrorCode as Ec, HardwareAuthToken::HardwareAuthToken,
     HardwareAuthenticatorType::HardwareAuthenticatorType, KeyPurpose::KeyPurpose,
     SecurityLevel::SecurityLevel, Tag::Tag, Timestamp::Timestamp,
+    VerificationToken::VerificationToken,
 };
 use android_system_keystore2::aidl::android::system::keystore2::OperationChallenge::OperationChallenge;
 use anyhow::{Context, Result};
 use std::collections::{HashMap, HashSet};
+use std::sync::mpsc::{channel, Sender};
 use std::sync::Mutex;
 use std::time::SystemTime;
 
@@ -38,25 +41,40 @@
     // This maps the operation challenge to an optional auth token, to maintain op-auth tokens
     // in-memory, until they are picked up and given to the operation by authorise_update_finish().
     op_auth_map: Mutex<HashMap<i64, Option<HardwareAuthToken>>>,
+    // sender end of the channel via which the enforcement module communicates with the
+    // background task handler (bth). This is of type Mutex in an Option because it is initialized
+    // after the global enforcement object is created.
+    sender_to_bth: Mutex<Option<Sender<Message>>>,
 }
 
 impl Enforcements {
-    /// Creates an enforcement object with the two data structures it holds.
+    /// Creates an enforcement object with the two data structures it holds and the sender as None.
     pub fn new() -> Self {
         Enforcements {
             device_unlocked_set: Mutex::new(HashSet::new()),
             op_auth_map: Mutex::new(HashMap::new()),
+            sender_to_bth: Mutex::new(None),
         }
     }
 
+    /// Initialize the sender_to_bth field, using the given sender end of a channel.
+    pub fn set_sender_to_bth(&self, sender: Sender<Message>) {
+        // It is ok to unwrap here because there is no chance of poisoning this mutex.
+        let mut sender_guard = self.sender_to_bth.lock().unwrap();
+        *sender_guard = Some(sender);
+    }
+
     /// Checks if update or finish calls are authorized. If the operation is based on per-op key,
     /// 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 +421,43 @@
             .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);
+    }
+
+    /// Requests a verification token from the background task handler which will retrieve it from
+    /// Timestamp Service or TEE KeyMint.
+    /// Once the create_operation receives an operation challenge from KeyMint, if it has
+    /// previously received a VerificationRequired variant of AuthTokenHandler during
+    /// authorize_create_operation, it calls this method to obtain a VerificationToken.
+    pub fn request_verification_token(
+        &self,
+        auth_token: HardwareAuthToken,
+        op_challenge: OperationChallenge,
+    ) -> Result<AuthTokenHandler> {
+        // create a channel for this particular operation
+        let (op_sender, op_receiver) = channel::<(HardwareAuthToken, VerificationToken)>();
+        // it is ok to unwrap here because there is no way this mutex gets poisoned.
+        let sender_guard = self.sender_to_bth.lock().unwrap();
+        if let Some(sender) = &*sender_guard {
+            let sender_cloned = sender.clone();
+            drop(sender_guard);
+            sender_cloned
+                .send(Message::Inputs((auth_token, op_challenge, op_sender)))
+                .map_err(|_| KeystoreError::sys())
+                .context(
+                    "In request_verification_token. Sending a request for a verification token
+             failed.",
+                )?;
+        }
+        Ok(AuthTokenHandler::Channel(op_receiver))
+    }
 }
 
 impl Default for Enforcements {
@@ -411,4 +466,21 @@
     }
 }
 
+impl Drop for Enforcements {
+    fn drop(&mut self) {
+        let sender_guard = self.sender_to_bth.lock().unwrap();
+        if let Some(sender) = &*sender_guard {
+            let sender_cloned = sender.clone();
+            drop(sender_guard);
+            // TODO: Verify how best to handle the error in this case.
+            sender_cloned.send(Message::Shutdown).unwrap_or_else(|e| {
+                panic!(
+                    "Failed to send shutdown message to background task handler because of {:?}.",
+                    e
+                );
+            });
+        }
+    }
+}
+
 // TODO: Add tests to enforcement module (b/175578618).
diff --git a/keystore2/src/globals.rs b/keystore2/src/globals.rs
index 18bed9a..035dac1 100644
--- a/keystore2/src/globals.rs
+++ b/keystore2/src/globals.rs
@@ -17,6 +17,8 @@
 //! to talk to.
 
 use crate::async_task::AsyncTask;
+use crate::background_task_handler::BackgroundTaskHandler;
+use crate::enforcements::Enforcements;
 use crate::gc::Gc;
 use crate::super_key::SuperKeyManager;
 use crate::utils::Asp;
@@ -81,6 +83,14 @@
     /// 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();
+    /// Background task handler is initialized and exists globally.
+    /// The other modules (e.g. enforcements) communicate with it via a channel initialized during
+    /// keystore startup.
+    pub static ref BACKGROUND_TASK_HANDLER: BackgroundTaskHandler = BackgroundTaskHandler::new();
 }
 
 static KEYMINT_SERVICE_NAME: &str = "android.hardware.security.keymint.IKeyMintDevice";
diff --git a/keystore2/src/keystore2_main.rs b/keystore2/src/keystore2_main.rs
index 7391f37..8607eef 100644
--- a/keystore2/src/keystore2_main.rs
+++ b/keystore2/src/keystore2_main.rs
@@ -16,9 +16,12 @@
 
 use binder::Interface;
 use keystore2::apc::ApcManager;
+use keystore2::background_task_handler::Message;
+use keystore2::globals::{BACKGROUND_TASK_HANDLER, ENFORCEMENTS};
 use keystore2::service::KeystoreService;
 use log::{error, info};
 use std::panic;
+use std::sync::mpsc::channel;
 
 static KS2_SERVICE_NAME: &str = "android.system.keystore2";
 static APC_SERVICE_NAME: &str = "android.security.apc";
@@ -50,6 +53,14 @@
         panic!("Must specify a working directory.");
     }
 
+    // initialize the channel via which the enforcement module and background task handler module
+    // communicate, and hand over the sender and receiver ends to the respective objects.
+    let (sender, receiver) = channel::<Message>();
+    ENFORCEMENTS.set_sender_to_bth(sender);
+    BACKGROUND_TASK_HANDLER.start_bth(receiver).unwrap_or_else(|e| {
+        panic!("Failed to start background task handler because of {:?}.", e);
+    });
+
     info!("Starting thread pool now.");
     binder::ProcessState::start_thread_pool();
 
diff --git a/keystore2/src/lib.rs b/keystore2/src/lib.rs
index 29b3992..f73cd59 100644
--- a/keystore2/src/lib.rs
+++ b/keystore2/src/lib.rs
@@ -17,6 +17,7 @@
 
 pub mod apc;
 pub mod auth_token_handler;
+pub mod background_task_handler;
 pub mod database;
 pub mod enforcements;
 pub mod error;
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..079e92a 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -17,19 +17,22 @@
 //! 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,
 };
 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::{
@@ -159,7 +162,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 +177,7 @@
                         }
                     },
                     None,
+                    None,
                 )
             }
             _ => {
@@ -197,14 +201,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 +215,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 +260,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 +272,36 @@
             )
             .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) => {
+                //request a verification token, given the auth token and the challenge
+                auth_token_handler = ENFORCEMENTS
+                    .request_verification_token(
+                        auth_token,
+                        OperationChallenge { challenge: begin_result.challenge },
+                    )
+                    .context("In create_operation.")?;
+            }
+            _ => {}
+        }
+
         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 +315,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 }),