Keystore 2.0: Add safe crypto wrapper

 * Adds safe wrappers for AES_gcm_decrypt and AES_gcm_encrypt.
 * Adds AES256 key generation.
 * Adds ZVec, a simple fixed size owned vector type that locks
   the backing memory in place with mlock and zeroes the buffer
   before freeing it.

Test: keystore2_test
Bug: 173545997
Change-Id: Id7e30d50b024da1fa8aa58a07cd9bb7a861f81f0
diff --git a/keystore2/src/crypto/Android.bp b/keystore2/src/crypto/Android.bp
index 061cf9a..03c42b2 100644
--- a/keystore2/src/crypto/Android.bp
+++ b/keystore2/src/crypto/Android.bp
@@ -18,9 +18,14 @@
     srcs: ["lib.rs"],
     rustlibs: [
         "libkeystore2_crypto_bindgen",
+        "liblog_rust",
+        "libnix",
+        "libthiserror",
     ],
-    static_libs: ["libkeystore2_crypto"],
-    shared_libs: ["libcrypto"],
+    shared_libs: [
+        "libkeystore2_crypto",
+        "libcrypto",
+    ],
 }
 
 cc_library {
@@ -29,7 +34,7 @@
         "crypto.cpp",
         "certificate_utils.cpp",
     ],
-    export_include_dirs: ["include",],
+    export_include_dirs: ["include"],
     shared_libs: [
         "libcrypto",
         "liblog",
@@ -52,7 +57,9 @@
     auto_gen_config: true,
     rustlibs: [
         "libkeystore2_crypto_bindgen",
-        "libkeystore2_crypto_rust",
+        "liblog_rust",
+        "libnix",
+        "libthiserror",
     ],
     static_libs: [
         "libkeystore2_crypto",
@@ -60,11 +67,12 @@
     shared_libs: [
         "libc++",
         "libcrypto",
-	"liblog",
+        "liblog",
     ],
 }
 
 cc_test {
+    name: "keystore2_crypto_test",
     cflags: [
         "-Wall",
         "-Werror",
@@ -74,11 +82,11 @@
         "tests/certificate_utils_test.cpp",
         "tests/gtest_main.cpp",
     ],
+    test_suites: ["general-tests"],
     static_libs: [
         "libkeystore2_crypto",
     ],
     shared_libs: [
         "libcrypto",
     ],
-    name: "keystore2_crypto_test",
 }
diff --git a/keystore2/src/crypto/TEST_MAPPING b/keystore2/src/crypto/TEST_MAPPING
new file mode 100644
index 0000000..f6fb97a
--- /dev/null
+++ b/keystore2/src/crypto/TEST_MAPPING
@@ -0,0 +1,10 @@
+{
+  "presubmit": [
+    {
+      "name": "keystore2_crypto_test_rust"
+    },
+    {
+      "name": "keystore2_crypto_test"
+    }
+  ]
+}
diff --git a/keystore2/src/crypto/crypto.cpp b/keystore2/src/crypto/crypto.cpp
index 8c52e4c..173ed11 100644
--- a/keystore2/src/crypto/crypto.cpp
+++ b/keystore2/src/crypto/crypto.cpp
@@ -21,6 +21,7 @@
 #include <log/log.h>
 #include <openssl/aes.h>
 #include <openssl/evp.h>
+#include <openssl/rand.h>
 
 #include <vector>
 
@@ -60,6 +61,10 @@
     return cipher;
 }
 
+bool randomBytes(uint8_t* out, size_t len) {
+    return RAND_bytes(out, len);
+}
+
 /*
  * Encrypt 'len' data at 'in' with AES-GCM, using 128-bit or 256-bit key at 'key', 96-bit IV at
  * 'iv' and write output to 'out' (which may be the same location as 'in') and 128-bit tag to
@@ -172,13 +177,13 @@
 // Copied from system/security/keystore/user_state.cpp.
 
 void generateKeyFromPassword(uint8_t* key, size_t key_len, const char* pw, size_t pw_len,
-                             uint8_t* salt) {
+                             const uint8_t* salt) {
     size_t saltSize;
     if (salt != nullptr) {
         saltSize = SALT_SIZE;
     } else {
         // Pre-gingerbread used this hardwired salt, readMasterKey will rewrite these when found
-        salt = (uint8_t*)"keystore";
+        salt = reinterpret_cast<const uint8_t*>("keystore");
         // sizeof = 9, not strlen = 8
         saltSize = sizeof("keystore");
     }
diff --git a/keystore2/src/crypto/crypto.hpp b/keystore2/src/crypto/crypto.hpp
index 9a9bb2e..2e597f1 100644
--- a/keystore2/src/crypto/crypto.hpp
+++ b/keystore2/src/crypto/crypto.hpp
@@ -22,6 +22,7 @@
 #include <stddef.h>
 
 extern "C" {
+  bool randomBytes(uint8_t* out, size_t len);
   bool AES_gcm_encrypt(const uint8_t* in, uint8_t* out, size_t len,
                        const uint8_t* key, size_t key_size, const uint8_t* iv, uint8_t* tag);
   bool AES_gcm_decrypt(const uint8_t* in, uint8_t* out, size_t len,
@@ -34,7 +35,7 @@
   bool CreateKeyId(const uint8_t* key_blob, size_t len, km_id_t* out_id);
 
   void generateKeyFromPassword(uint8_t* key, size_t key_len, const char* pw,
-                               size_t pw_len, uint8_t* salt);
+                               size_t pw_len, const uint8_t* salt);
 }
 
 #endif  //  __CRYPTO_H__
diff --git a/keystore2/src/crypto/error.rs b/keystore2/src/crypto/error.rs
new file mode 100644
index 0000000..2eb97b9
--- /dev/null
+++ b/keystore2/src/crypto/error.rs
@@ -0,0 +1,59 @@
+// 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 Error for the keystore2_crypto library.
+
+/// Crypto specific error codes.
+#[derive(Debug, thiserror::Error, Eq, PartialEq)]
+pub enum Error {
+    /// This is returned if the C/C++ implementation of AES_gcm_decrypt returned false.
+    #[error("Failed to decrypt.")]
+    DecryptionFailed,
+
+    /// This is returned if the C/C++ implementation of AES_gcm_encrypt returned false.
+    #[error("Failed to encrypt.")]
+    EncryptionFailed,
+
+    /// The initialization vector has the wrong length.
+    #[error("Invalid IV length.")]
+    InvalidIvLength,
+
+    /// The aead tag has the wrong length.
+    #[error("Invalid AEAD tag length.")]
+    InvalidAeadTagLength,
+
+    /// The key has the wrong length.
+    #[error("Invalid key length.")]
+    InvalidKeyLength,
+
+    /// Invalid data length.
+    #[error("Invalid data length.")]
+    InvalidDataLength,
+
+    /// Invalid salt length.
+    #[error("Invalid salt length.")]
+    InvalidSaltLength,
+
+    /// Random number generation failed.
+    #[error("Random number generation failed.")]
+    RandomNumberGenerationFailed,
+
+    /// ZVec construction failed.
+    #[error(transparent)]
+    LayoutError(#[from] std::alloc::LayoutErr),
+
+    /// Nix error.
+    #[error(transparent)]
+    NixError(#[from] nix::Error),
+}
diff --git a/keystore2/src/crypto/lib.rs b/keystore2/src/crypto/lib.rs
index 6ec5edb..338bdb9 100644
--- a/keystore2/src/crypto/lib.rs
+++ b/keystore2/src/crypto/lib.rs
@@ -12,17 +12,188 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// TODO: Once this is complete, remove this and document everything public.
-#![allow(missing_docs)]
+//! This module implements safe wrappers for some crypto operations required by
+//! Keystore 2.0.
+
+mod error;
+mod zvec;
+pub use error::Error;
+use keystore2_crypto_bindgen::{
+    generateKeyFromPassword, randomBytes, size_t, AES_gcm_decrypt, AES_gcm_encrypt,
+};
+pub use zvec::ZVec;
+
+/// Length of the expected initialization vector.
+pub const IV_LENGTH: usize = 16;
+/// Length of the expected AEAD TAG.
+pub const TAG_LENGTH: usize = 16;
+/// Length of an AES 256 key in bytes.
+pub const AES_256_KEY_LENGTH: usize = 32;
+/// Length of an AES 128 key in bytes.
+pub const AES_128_KEY_LENGTH: usize = 16;
+/// Length of the expected salt for key from password generation.
+pub const SALT_LENGTH: usize = 16;
+
+// This is the number of bytes of the GCM IV that is expected to be initialized
+// with random bytes.
+const GCM_IV_LENGTH: usize = 12;
+
+/// Generate an AES256 key, essentially 32 random bytes from the underlying
+/// boringssl library discretely stuffed into a ZVec.
+pub fn generate_aes256_key() -> Result<ZVec, Error> {
+    // Safety: key has the same length as the requested number of random bytes.
+    let mut key = ZVec::new(AES_256_KEY_LENGTH)?;
+    if unsafe { randomBytes(key.as_mut_ptr(), AES_256_KEY_LENGTH as size_t) } {
+        Ok(key)
+    } else {
+        Err(Error::RandomNumberGenerationFailed)
+    }
+}
+
+/// Generate a salt.
+pub fn generate_salt() -> Result<Vec<u8>, Error> {
+    // Safety: salt has the same length as the requested number of random bytes.
+    let mut salt = vec![0; SALT_LENGTH];
+    if unsafe { randomBytes(salt.as_mut_ptr(), SALT_LENGTH as size_t) } {
+        Ok(salt)
+    } else {
+        Err(Error::RandomNumberGenerationFailed)
+    }
+}
+
+/// Uses AES GCM to decipher a message given an initialization vector, aead tag, and key.
+/// This function accepts 128 and 256-bit keys and uses AES128 and AES256 respectively based
+/// on the key length.
+/// This function returns the plaintext message in a ZVec because it is assumed that
+/// it contains sensitive information that should be zeroed from memory before its buffer is
+/// freed. Input key is taken as a slice for flexibility, but it is recommended that it is held
+/// in a ZVec as well.
+pub fn aes_gcm_decrypt(data: &[u8], iv: &[u8], tag: &[u8], key: &[u8]) -> Result<ZVec, Error> {
+    if iv.len() != IV_LENGTH {
+        return Err(Error::InvalidIvLength);
+    }
+
+    if tag.len() != TAG_LENGTH {
+        return Err(Error::InvalidAeadTagLength);
+    }
+
+    match key.len() {
+        AES_128_KEY_LENGTH | AES_256_KEY_LENGTH => {}
+        _ => return Err(Error::InvalidKeyLength),
+    }
+
+    let mut result = ZVec::new(data.len())?;
+
+    // Safety: The first two arguments must point to buffers with a size given by the third
+    // argument. The key must have a size of 16 or 32 bytes which we check above.
+    // The iv and tag arguments must be 16 bytes, which we also check above.
+    match unsafe {
+        AES_gcm_decrypt(
+            data.as_ptr(),
+            result.as_mut_ptr(),
+            data.len() as size_t,
+            key.as_ptr(),
+            key.len() as size_t,
+            iv.as_ptr(),
+            tag.as_ptr(),
+        )
+    } {
+        true => Ok(result),
+        false => Err(Error::DecryptionFailed),
+    }
+}
+
+/// Uses AES GCM to encrypt a message given a key.
+/// This function accepts 128 and 256-bit keys and uses AES128 and AES256 respectively based on
+/// the key length. The function generates an initialization vector. The return value is a tuple
+/// of `(ciphertext, iv, tag)`.
+pub fn aes_gcm_encrypt(data: &[u8], key: &[u8]) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>), Error> {
+    let mut iv = vec![0; IV_LENGTH];
+    // Safety: iv is longer than GCM_IV_LENGTH, which is 12 while IV_LENGTH is 16.
+    // The iv needs to be 16 bytes long, but the last 4 bytes remain zeroed.
+    if !unsafe { randomBytes(iv.as_mut_ptr(), GCM_IV_LENGTH as size_t) } {
+        return Err(Error::RandomNumberGenerationFailed);
+    }
+
+    match key.len() {
+        AES_128_KEY_LENGTH | AES_256_KEY_LENGTH => {}
+        _ => return Err(Error::InvalidKeyLength),
+    }
+
+    let mut result: Vec<u8> = vec![0; data.len()];
+    let mut tag: Vec<u8> = vec![0; TAG_LENGTH];
+    match unsafe {
+        AES_gcm_encrypt(
+            data.as_ptr(),
+            result.as_mut_ptr(),
+            data.len() as size_t,
+            key.as_ptr(),
+            key.len() as size_t,
+            iv.as_ptr(),
+            tag.as_mut_ptr(),
+        )
+    } {
+        true => Ok((result, iv, tag)),
+        false => Err(Error::EncryptionFailed),
+    }
+}
+
+/// Generates a key from the given password and salt.
+/// The salt must be exactly 16 bytes long.
+/// Two key sizes are accepted: 16 and 32 bytes.
+pub fn derive_key_from_password(
+    pw: &[u8],
+    salt: Option<&[u8]>,
+    key_length: usize,
+) -> Result<ZVec, Error> {
+    let salt: *const u8 = match salt {
+        Some(s) => {
+            if s.len() != SALT_LENGTH {
+                return Err(Error::InvalidSaltLength);
+            }
+            s.as_ptr()
+        }
+        None => std::ptr::null(),
+    };
+
+    match key_length {
+        AES_128_KEY_LENGTH | AES_256_KEY_LENGTH => {}
+        _ => return Err(Error::InvalidKeyLength),
+    }
+
+    let mut result = ZVec::new(key_length)?;
+
+    unsafe {
+        generateKeyFromPassword(
+            result.as_mut_ptr(),
+            result.len() as size_t,
+            pw.as_ptr() as *const std::os::raw::c_char,
+            pw.len() as size_t,
+            salt,
+        )
+    };
+
+    Ok(result)
+}
 
 #[cfg(test)]
 mod tests {
 
+    use super::*;
     use keystore2_crypto_bindgen::{
         generateKeyFromPassword, AES_gcm_decrypt, AES_gcm_encrypt, CreateKeyId,
     };
 
     #[test]
+    fn test_wrapper_roundtrip() {
+        let key = generate_aes256_key().unwrap();
+        let message = b"totally awesome message";
+        let (cipher_text, iv, tag) = aes_gcm_encrypt(message, &key).unwrap();
+        let message2 = aes_gcm_decrypt(&cipher_text, &iv, &tag, &key).unwrap();
+        assert_eq!(message[..], message2[..])
+    }
+
+    #[test]
     fn test_encrypt_decrypt() {
         let input = vec![0; 16];
         let mut out = vec![0; 16];
diff --git a/keystore2/src/crypto/zvec.rs b/keystore2/src/crypto/zvec.rs
new file mode 100644
index 0000000..52addfc
--- /dev/null
+++ b/keystore2/src/crypto/zvec.rs
@@ -0,0 +1,102 @@
+// 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.
+
+#![allow(dead_code)]
+
+use crate::error::Error;
+use nix::sys::mman::{mlock, munlock};
+use std::convert::TryFrom;
+use std::fmt;
+use std::ops::{Deref, DerefMut};
+use std::ptr::write_volatile;
+
+/// A fixed size u8 vector that is zeroed when dropped. Also the data is
+/// pinned in memory with mlock.
+#[derive(Default, Eq, PartialEq)]
+pub struct ZVec(Box<[u8]>);
+
+impl ZVec {
+    /// Create a ZVec with the given size.
+    pub fn new(size: usize) -> Result<Self, Error> {
+        let v: Vec<u8> = vec![0; size];
+        let b = v.into_boxed_slice();
+        if size > 0 {
+            unsafe { mlock(b.as_ptr() as *const std::ffi::c_void, b.len()) }?;
+        }
+        Ok(Self(b))
+    }
+}
+
+impl Drop for ZVec {
+    fn drop(&mut self) {
+        for i in 0..self.0.len() {
+            unsafe { write_volatile(self.0.as_mut_ptr().add(i), 0) };
+        }
+        if !self.0.is_empty() {
+            if let Err(e) =
+                unsafe { munlock(self.0.as_ptr() as *const std::ffi::c_void, self.0.len()) }
+            {
+                log::error!("In ZVec::drop: `munlock` failed: {:?}.", e);
+            }
+        }
+    }
+}
+
+impl Deref for ZVec {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for ZVec {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl fmt::Debug for ZVec {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        if self.0.is_empty() {
+            write!(f, "Zvec empty")
+        } else {
+            write!(f, "Zvec size: {} [ Sensitive information redacted ]", self.0.len())
+        }
+    }
+}
+
+impl TryFrom<&[u8]> for ZVec {
+    type Error = Error;
+
+    fn try_from(v: &[u8]) -> Result<Self, Self::Error> {
+        let mut z = ZVec::new(v.len())?;
+        if !v.is_empty() {
+            z.clone_from_slice(v);
+        }
+        Ok(z)
+    }
+}
+
+impl TryFrom<Vec<u8>> for ZVec {
+    type Error = Error;
+
+    fn try_from(v: Vec<u8>) -> Result<Self, Self::Error> {
+        let b = v.into_boxed_slice();
+        if !b.is_empty() {
+            unsafe { mlock(b.as_ptr() as *const std::ffi::c_void, b.len()) }?;
+        }
+        Ok(Self(b))
+    }
+}