cryptohome: Cleanup DeriveVaultKeyset for AuthBlock

BUG=chromium:728125
TEST=FEATURES=test emerge-nami cryptohome

Change-Id: I8a46f169fd9fcd3f59ba75b6cc6430cea5cee4f5
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform2/+/2500634
Tested-by: Hardik Goyal <hardikgoyal@chromium.org>
Tested-by: Greg Kerr <kerrnel@chromium.org>
Reviewed-by: Greg Kerr <kerrnel@chromium.org>
Commit-Queue: Hardik Goyal <hardikgoyal@chromium.org>
Auto-Submit: Hardik Goyal <hardikgoyal@chromium.org>
diff --git a/cryptohome/crypto.cc b/cryptohome/crypto.cc
index b44c3cf..df79605 100644
--- a/cryptohome/crypto.cc
+++ b/cryptohome/crypto.cc
@@ -441,90 +441,28 @@
   PopulateError(error, CryptoError::CE_NONE);
 
   unsigned int flags = serialized.flags();
+  std::unique_ptr<AuthBlock> auth_block = DeriveAuthBlock(flags);
+  if (!auth_block) {
+    LOG(ERROR) << "Keyset wrapped with unknown method.";
+    return false;
+  }
+
+  AuthInput auth_input = {vault_key, locked_to_single_user};
+  AuthBlockState auth_state = {serialized};
+  KeyBlobs vkk_data;
+  if (!auth_block->Derive(auth_input, auth_state, &vkk_data, error)) {
+    return false;
+  }
 
   if (flags & SerializedVaultKeyset::LE_CREDENTIAL) {
-    PinWeaverAuthBlock pin_weaver_auth(le_manager_.get(), tpm_init_);
-
-    AuthInput auth_input = {vault_key};
-    AuthBlockState auth_state = {serialized};
-    KeyBlobs vkk_data;
-    if (!pin_weaver_auth.Derive(auth_input, auth_state, &vkk_data, error)) {
-      return false;
-    }
-
     // This is possible to be empty if an old version of CR50 is running.
     if (vkk_data.reset_secret.has_value() &&
         !vkk_data.reset_secret.value().empty()) {
       vault_keyset->set_reset_secret(vkk_data.reset_secret.value());
     }
-
-    return UnwrapVaultKeyset(serialized, vkk_data, vault_keyset, error);
   }
 
-  if (flags & SerializedVaultKeyset::SIGNATURE_CHALLENGE_PROTECTED) {
-    AuthInput user_input = {vault_key};
-    AuthBlockState auth_state = {serialized};
-    KeyBlobs vkk_data;
-    ChallengeCredentialAuthBlock auth_block;
-    if (!auth_block.Derive(user_input, auth_state, &vkk_data, error)) {
-      return false;
-    }
-    return UnwrapVaultKeyset(serialized, vkk_data, vault_keyset, error);
-  }
-
-  if (flags & SerializedVaultKeyset::SCRYPT_WRAPPED &&
-      flags & SerializedVaultKeyset::TPM_WRAPPED) {
-    LOG(ERROR) << "Keyset wrapped with both TPM and Scrypt?";
-    ReportCryptohomeError(cryptohome::kBothTpmAndScryptWrappedKeyset);
-
-    AuthInput auth_input = {vault_key};
-    AuthBlockState auth_state = {serialized};
-    KeyBlobs vkk_data;
-    DoubleWrappedCompatAuthBlock auth_block(tpm_, tpm_init_);
-    if (!auth_block.Derive(auth_input, auth_state, &vkk_data, error)) {
-      return false;
-    }
-    return UnwrapVaultKeyset(serialized, vkk_data, vault_keyset, error);
-  }
-
-  if (flags & SerializedVaultKeyset::TPM_WRAPPED) {
-    std::unique_ptr<AuthBlock> tpm_auth;
-    if (flags & SerializedVaultKeyset::PCR_BOUND) {
-      tpm_auth = std::make_unique<TpmBoundToPcrAuthBlock>(tpm_, tpm_init_);
-    } else {
-      tpm_auth = std::make_unique<TpmNotBoundToPcrAuthBlock>(tpm_, tpm_init_);
-    }
-
-    KeyBlobs vkk_data;
-    AuthInput auth_input;
-    auth_input.user_input = vault_key;
-    auth_input.locked_to_single_user = locked_to_single_user;
-
-    AuthBlockState auth_state = {serialized};
-    if (!tpm_auth->Derive(auth_input, auth_state, &vkk_data, error)) {
-      return false;
-    }
-
-    return UnwrapVaultKeyset(serialized, vkk_data, vault_keyset, error);
-  }
-
-  if (flags & SerializedVaultKeyset::SCRYPT_WRAPPED) {
-    KeyBlobs vkk_data;
-    AuthInput auth_input;
-    auth_input.user_input = vault_key;
-    auth_input.locked_to_single_user = locked_to_single_user;
-
-    AuthBlockState auth_state = {serialized};
-    LibScryptCompatAuthBlock auth_block;
-    if (!auth_block.Derive(auth_input, auth_state, &vkk_data, error)) {
-      return false;
-    }
-
-    return UnwrapVaultKeyset(serialized, vkk_data, vault_keyset, error);
-  }
-
-  LOG(ERROR) << "Keyset wrapped with unknown method.";
-  return false;
+  return UnwrapVaultKeyset(serialized, vkk_data, vault_keyset, error);
 }
 
 bool Crypto::GenerateEncryptedRawKeyset(const VaultKeyset& vault_keyset,
@@ -1125,4 +1063,25 @@
 #endif
 }
 
+std::unique_ptr<AuthBlock> Crypto::DeriveAuthBlock(int serialized_key_flags) {
+  if (serialized_key_flags & SerializedVaultKeyset::LE_CREDENTIAL) {
+    return std::make_unique<PinWeaverAuthBlock>(le_manager_.get(), tpm_init_);
+  } else if (serialized_key_flags &
+             SerializedVaultKeyset::SIGNATURE_CHALLENGE_PROTECTED) {
+    return std::make_unique<ChallengeCredentialAuthBlock>();
+  } else if (serialized_key_flags & SerializedVaultKeyset::SCRYPT_WRAPPED &&
+             serialized_key_flags & SerializedVaultKeyset::TPM_WRAPPED) {
+    return std::make_unique<DoubleWrappedCompatAuthBlock>(tpm_, tpm_init_);
+  } else if (serialized_key_flags & SerializedVaultKeyset::TPM_WRAPPED) {
+    if (serialized_key_flags & SerializedVaultKeyset::PCR_BOUND) {
+      return std::make_unique<TpmBoundToPcrAuthBlock>(tpm_, tpm_init_);
+    } else {
+      return std::make_unique<TpmNotBoundToPcrAuthBlock>(tpm_, tpm_init_);
+    }
+  } else if (serialized_key_flags & SerializedVaultKeyset::SCRYPT_WRAPPED) {
+    return std::make_unique<LibScryptCompatAuthBlock>();
+  }
+  return nullptr;
+}
+
 }  // namespace cryptohome
diff --git a/cryptohome/crypto.h b/cryptohome/crypto.h
index 75e2ef4..b16f022 100644
--- a/cryptohome/crypto.h
+++ b/cryptohome/crypto.h
@@ -34,6 +34,7 @@
 
 struct KeyBlobs;
 class VaultKeyset;
+class AuthBlock;
 
 extern const char kSystemSaltFile[];
 
@@ -275,6 +276,12 @@
                                 const brillo::SecureBlob& vkk_key,
                                 const brillo::SecureBlob& vkk_iv) const;
 
+  // This function serves as a factory method to return the authblock used in
+  // authentication.
+  // |serialized_key_flags| is the flag data stores in the vault keyset which
+  // helps determine the type of AuthBlock.
+  std::unique_ptr<AuthBlock> DeriveAuthBlock(int serialized_key_flags);
+
   // If set, the TPM will be used during the encryption of the vault keyset
   bool use_tpm_;