Merge "AndroidKeyStore: Add encrypted flag"
diff --git a/keystore/java/android/security/AndroidKeyPairGenerator.java b/keystore/java/android/security/AndroidKeyPairGenerator.java
index c42001b..6975583 100644
--- a/keystore/java/android/security/AndroidKeyPairGenerator.java
+++ b/keystore/java/android/security/AndroidKeyPairGenerator.java
@@ -49,10 +49,7 @@
  *
  * {@hide}
  */
-@SuppressWarnings("deprecation")
 public class AndroidKeyPairGenerator extends KeyPairGeneratorSpi {
-    public static final String NAME = "AndroidKeyPairGenerator";
-
     private android.security.KeyStore mKeyStore;
 
     private AndroidKeyPairGeneratorSpec mSpec;
@@ -79,12 +76,21 @@
                     "Must call initialize with an AndroidKeyPairGeneratorSpec first");
         }
 
+        if (((mSpec.getFlags() & KeyStore.FLAG_ENCRYPTED) != 0)
+                && (mKeyStore.state() != KeyStore.State.UNLOCKED)) {
+            throw new IllegalStateException(
+                    "Android keystore must be in initialized and unlocked state "
+                            + "if encryption is required");
+        }
+
         final String alias = mSpec.getKeystoreAlias();
 
         Credentials.deleteAllTypesForAlias(mKeyStore, alias);
 
         final String privateKeyAlias = Credentials.USER_PRIVATE_KEY + alias;
-        mKeyStore.generate(privateKeyAlias);
+        if (!mKeyStore.generate(privateKeyAlias, KeyStore.UID_SELF, mSpec.getFlags())) {
+            throw new IllegalStateException("could not generate key in keystore");
+        }
 
         final PrivateKey privKey;
         final OpenSSLEngine engine = OpenSSLEngine.getInstance("keystore");
@@ -131,7 +137,8 @@
             throw new IllegalStateException("Can't get encoding of certificate", e);
         }
 
-        if (!mKeyStore.put(Credentials.USER_CERTIFICATE + alias, certBytes)) {
+        if (!mKeyStore.put(Credentials.USER_CERTIFICATE + alias, certBytes, KeyStore.UID_SELF,
+                mSpec.getFlags())) {
             Credentials.deleteAllTypesForAlias(mKeyStore, alias);
             throw new IllegalStateException("Can't store certificate in AndroidKeyStore");
         }
diff --git a/keystore/java/android/security/AndroidKeyPairGeneratorSpec.java b/keystore/java/android/security/AndroidKeyPairGeneratorSpec.java
index 83faf35..8d6bfa68 100644
--- a/keystore/java/android/security/AndroidKeyPairGeneratorSpec.java
+++ b/keystore/java/android/security/AndroidKeyPairGeneratorSpec.java
@@ -32,10 +32,9 @@
  * {@code KeyPairGenerator} that works with <a href="{@docRoot}
  * guide/topics/security/keystore.html">Android KeyStore facility</a>. The
  * Android KeyStore facility is accessed through a
- * {@link java.security.KeyPairGenerator} API using the
- * {@code AndroidKeyPairGenerator} provider. The {@code context} passed in may
- * be used to pop up some UI to ask the user to unlock or initialize the Android
- * keystore facility.
+ * {@link java.security.KeyPairGenerator} API using the {@code AndroidKeyStore}
+ * provider. The {@code context} passed in may be used to pop up some UI to ask
+ * the user to unlock or initialize the Android KeyStore facility.
  * <p>
  * After generation, the {@code keyStoreAlias} is used with the
  * {@link java.security.KeyStore#getEntry(String, java.security.KeyStore.ProtectionParameter)}
@@ -47,12 +46,12 @@
  * Distinguished Name along with the other parameters specified with the
  * {@link Builder}.
  * <p>
- * The self-signed certificate may be replaced at a later time by a certificate
- * signed by a real Certificate Authority.
+ * The self-signed X.509 certificate may be replaced at a later time by a
+ * certificate signed by a real Certificate Authority.
  *
  * @hide
  */
-public class AndroidKeyPairGeneratorSpec implements AlgorithmParameterSpec {
+public final class AndroidKeyPairGeneratorSpec implements AlgorithmParameterSpec {
     private final String mKeystoreAlias;
 
     private final Context mContext;
@@ -65,6 +64,8 @@
 
     private final Date mEndDate;
 
+    private final int mFlags;
+
     /**
      * Parameter specification for the "{@code AndroidKeyPairGenerator}"
      * instance of the {@link java.security.KeyPairGenerator} API. The
@@ -95,7 +96,8 @@
      * @hide should be built with AndroidKeyPairGeneratorSpecBuilder
      */
     public AndroidKeyPairGeneratorSpec(Context context, String keyStoreAlias,
-            X500Principal subjectDN, BigInteger serialNumber, Date startDate, Date endDate) {
+            X500Principal subjectDN, BigInteger serialNumber, Date startDate, Date endDate,
+            int flags) {
         if (context == null) {
             throw new IllegalArgumentException("context == null");
         } else if (TextUtils.isEmpty(keyStoreAlias)) {
@@ -118,48 +120,69 @@
         mSerialNumber = serialNumber;
         mStartDate = startDate;
         mEndDate = endDate;
+        mFlags = flags;
     }
 
     /**
-     * @hide
+     * Returns the alias that will be used in the {@code java.security.KeyStore}
+     * in conjunction with the {@code AndroidKeyStore}.
      */
-    String getKeystoreAlias() {
+    public String getKeystoreAlias() {
         return mKeystoreAlias;
     }
 
     /**
-     * @hide
+     * Gets the Android context used for operations with this instance.
      */
-    Context getContext() {
+    public Context getContext() {
         return mContext;
     }
 
     /**
-     * @hide
+     * Gets the subject distinguished name to be used on the X.509 certificate
+     * that will be put in the {@link java.security.KeyStore}.
      */
-    X500Principal getSubjectDN() {
+    public X500Principal getSubjectDN() {
         return mSubjectDN;
     }
 
     /**
-     * @hide
+     * Gets the serial number to be used on the X.509 certificate that will be
+     * put in the {@link java.security.KeyStore}.
      */
-    BigInteger getSerialNumber() {
+    public BigInteger getSerialNumber() {
         return mSerialNumber;
     }
 
     /**
-     * @hide
+     * Gets the start date to be used on the X.509 certificate that will be put
+     * in the {@link java.security.KeyStore}.
      */
-    Date getStartDate() {
+    public Date getStartDate() {
         return mStartDate;
     }
 
     /**
+     * Gets the end date to be used on the X.509 certificate that will be put in
+     * the {@link java.security.KeyStore}.
+     */
+    public Date getEndDate() {
+        return mEndDate;
+    }
+
+    /**
      * @hide
      */
-    Date getEndDate() {
-        return mEndDate;
+    int getFlags() {
+        return mFlags;
+    }
+
+    /**
+     * Returns {@code true} if this parameter will require generated keys to be
+     * encrypted in the {@link java.security.KeyStore}.
+     */
+    public boolean isEncryptionRequired() {
+        return (mFlags & KeyStore.FLAG_ENCRYPTED) != 0;
     }
 
     /**
@@ -177,16 +200,17 @@
      * Calendar end = new Calendar();
      * end.add(1, Calendar.YEAR);
      *
-     * AndroidKeyPairGeneratorSpec spec = new AndroidKeyPairGeneratorSpec.Builder(mContext)
-     *         .setAlias("myKey")
-     *         .setSubject(new X500Principal("CN=myKey"))
-     *         .setSerial(BigInteger.valueOf(1337))
-     *         .setStartDate(start.getTime())
-     *         .setEndDate(end.getTime())
-     *         .build();
+     * AndroidKeyPairGeneratorSpec spec =
+     *         new AndroidKeyPairGeneratorSpec.Builder(mContext)
+     *                 .setAlias(&quot;myKey&quot;)
+     *                 .setSubject(new X500Principal(&quot;CN=myKey&quot;))
+     *                 .setSerial(BigInteger.valueOf(1337))
+     *                 .setStartDate(start.getTime())
+     *                 .setEndDate(end.getTime())
+     *                 .build();
      * </pre>
      */
-    public static class Builder {
+    public final static class Builder {
         private final Context mContext;
 
         private String mKeystoreAlias;
@@ -199,6 +223,14 @@
 
         private Date mEndDate;
 
+        private int mFlags;
+
+        /**
+         * Creates a new instance of the {@code Builder} with the given
+         * {@code context}. The {@code context} passed in may be used to pop up
+         * some UI to ask the user to unlock or initialize the Android KeyStore
+         * facility.
+         */
         public Builder(Context context) {
             if (context == null) {
                 throw new NullPointerException("context == null");
@@ -268,6 +300,17 @@
         }
 
         /**
+         * Indicates that this key must be encrypted at rest on storage. Note
+         * that enabling this will require that the user enable a strong lock
+         * screen (e.g., PIN, password) before creating or using the generated
+         * key is successful.
+         */
+        public Builder setEncryptionRequired() {
+            mFlags |= KeyStore.FLAG_ENCRYPTED;
+            return this;
+        }
+
+        /**
          * Builds the instance of the {@code AndroidKeyPairGeneratorSpec}.
          *
          * @throws IllegalArgumentException if a required field is missing
@@ -275,7 +318,7 @@
          */
         public AndroidKeyPairGeneratorSpec build() {
             return new AndroidKeyPairGeneratorSpec(mContext, mKeystoreAlias, mSubjectDN,
-                    mSerialNumber, mStartDate, mEndDate);
+                    mSerialNumber, mStartDate, mEndDate, mFlags);
         }
     }
 }
diff --git a/keystore/java/android/security/AndroidKeyStore.java b/keystore/java/android/security/AndroidKeyStore.java
index 8a9826b..dcc9516 100644
--- a/keystore/java/android/security/AndroidKeyStore.java
+++ b/keystore/java/android/security/AndroidKeyStore.java
@@ -27,6 +27,10 @@
 import java.io.OutputStream;
 import java.security.InvalidKeyException;
 import java.security.Key;
+import java.security.KeyStore.Entry;
+import java.security.KeyStore.PrivateKeyEntry;
+import java.security.KeyStore.ProtectionParameter;
+import java.security.KeyStore;
 import java.security.KeyStoreException;
 import java.security.KeyStoreSpi;
 import java.security.NoSuchAlgorithmException;
@@ -198,14 +202,14 @@
         }
 
         if (key instanceof PrivateKey) {
-            setPrivateKeyEntry(alias, (PrivateKey) key, chain);
+            setPrivateKeyEntry(alias, (PrivateKey) key, chain, null);
         } else {
             throw new KeyStoreException("Only PrivateKeys are supported");
         }
     }
 
-    private void setPrivateKeyEntry(String alias, PrivateKey key, Certificate[] chain)
-            throws KeyStoreException {
+    private void setPrivateKeyEntry(String alias, PrivateKey key, Certificate[] chain,
+            AndroidKeyStoreParameter params) throws KeyStoreException {
         byte[] keyBytes = null;
 
         final String pkeyAlias;
@@ -317,15 +321,20 @@
             Credentials.deleteCertificateTypesForAlias(mKeyStore, alias);
         }
 
+        final int flags = (params == null) ? 0 : params.getFlags();
+
         if (shouldReplacePrivateKey
-                && !mKeyStore.importKey(Credentials.USER_PRIVATE_KEY + alias, keyBytes)) {
+                && !mKeyStore.importKey(Credentials.USER_PRIVATE_KEY + alias, keyBytes,
+                        android.security.KeyStore.UID_SELF, flags)) {
             Credentials.deleteAllTypesForAlias(mKeyStore, alias);
             throw new KeyStoreException("Couldn't put private key in keystore");
-        } else if (!mKeyStore.put(Credentials.USER_CERTIFICATE + alias, userCertBytes)) {
+        } else if (!mKeyStore.put(Credentials.USER_CERTIFICATE + alias, userCertBytes,
+                android.security.KeyStore.UID_SELF, flags)) {
             Credentials.deleteAllTypesForAlias(mKeyStore, alias);
             throw new KeyStoreException("Couldn't put certificate #1 in keystore");
         } else if (chainBytes != null
-                && !mKeyStore.put(Credentials.CA_CERTIFICATE + alias, chainBytes)) {
+                && !mKeyStore.put(Credentials.CA_CERTIFICATE + alias, chainBytes,
+                        android.security.KeyStore.UID_SELF, flags)) {
             Credentials.deleteAllTypesForAlias(mKeyStore, alias);
             throw new KeyStoreException("Couldn't put certificate chain in keystore");
         }
@@ -355,7 +364,8 @@
             throw new KeyStoreException(e);
         }
 
-        if (!mKeyStore.put(Credentials.CA_CERTIFICATE + alias, encoded)) {
+        if (!mKeyStore.put(Credentials.CA_CERTIFICATE + alias, encoded,
+                android.security.KeyStore.UID_SELF, android.security.KeyStore.FLAG_NONE)) {
             throw new KeyStoreException("Couldn't insert certificate; is KeyStore initialized?");
         }
     }
@@ -517,4 +527,37 @@
         mKeyStore = android.security.KeyStore.getInstance();
     }
 
+    @Override
+    public void engineSetEntry(String alias, Entry entry, ProtectionParameter param)
+            throws KeyStoreException {
+        if (entry == null) {
+            throw new KeyStoreException("entry == null");
+        }
+
+        if (engineContainsAlias(alias)) {
+            engineDeleteEntry(alias);
+        }
+
+        if (entry instanceof KeyStore.TrustedCertificateEntry) {
+            KeyStore.TrustedCertificateEntry trE = (KeyStore.TrustedCertificateEntry) entry;
+            engineSetCertificateEntry(alias, trE.getTrustedCertificate());
+            return;
+        }
+
+        if (param != null && !(param instanceof AndroidKeyStoreParameter)) {
+            throw new KeyStoreException("protParam should be AndroidKeyStoreParameter; was: "
+                    + param.getClass().getName());
+        }
+
+        if (entry instanceof PrivateKeyEntry) {
+            PrivateKeyEntry prE = (PrivateKeyEntry) entry;
+            setPrivateKeyEntry(alias, prE.getPrivateKey(), prE.getCertificateChain(),
+                    (AndroidKeyStoreParameter) param);
+            return;
+        }
+
+        throw new KeyStoreException(
+                "Entry must be a PrivateKeyEntry or TrustedCertificateEntry; was " + entry);
+    }
+
 }
diff --git a/keystore/java/android/security/AndroidKeyStoreParameter.java b/keystore/java/android/security/AndroidKeyStoreParameter.java
new file mode 100644
index 0000000..4ce9f09
--- /dev/null
+++ b/keystore/java/android/security/AndroidKeyStoreParameter.java
@@ -0,0 +1,125 @@
+/*
+ * Copyright (C) 2013 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.
+ */
+
+package android.security;
+
+import android.content.Context;
+import android.security.AndroidKeyPairGeneratorSpec.Builder;
+
+import java.security.KeyPairGenerator;
+import java.security.PrivateKey;
+import java.security.KeyStore.ProtectionParameter;
+import java.security.cert.Certificate;
+
+/**
+ * This provides the optional parameters that can be specified for
+ * {@code KeyStore} entries that work with <a href="{@docRoot}
+ * guide/topics/security/keystore.html">Android KeyStore facility</a>. The
+ * Android KeyStore facility is accessed through a
+ * {@link java.security.KeyStore} API using the {@code AndroidKeyStore}
+ * provider. The {@code context} passed in may be used to pop up some UI to ask
+ * the user to unlock or initialize the Android KeyStore facility.
+ * <p>
+ * Any entries placed in the {@code KeyStore} may be retrieved later. Note that
+ * there is only one logical instance of the {@code KeyStore} per application
+ * UID so apps using the {@code sharedUid} facility will also share a
+ * {@code KeyStore}.
+ * <p>
+ * Keys may be generated using the {@link KeyPairGenerator} facility with a
+ * {@link AndroidKeyPairGeneratorSpec} to specify the entry's {@code alias}. A
+ * self-signed X.509 certificate will be attached to generated entries, but that
+ * may be replaced at a later time by a certificate signed by a real Certificate
+ * Authority.
+ *
+ * @hide
+ */
+public final class AndroidKeyStoreParameter implements ProtectionParameter {
+    private int mFlags;
+
+    private AndroidKeyStoreParameter(int flags) {
+        mFlags = flags;
+    }
+
+    /**
+     * @hide
+     */
+    public int getFlags() {
+        return mFlags;
+    }
+
+    /**
+     * Returns {@code true} if this parameter requires entries to be encrypted
+     * on the disk.
+     */
+    public boolean isEncryptionRequired() {
+        return (mFlags & KeyStore.FLAG_ENCRYPTED) != 0;
+    }
+
+    /**
+     * Builder class for {@link AndroidKeyStoreParameter} objects.
+     * <p>
+     * This will build protection parameters for use with the <a
+     * href="{@docRoot} guide/topics/security/keystore.html">Android KeyStore
+     * facility</a>.
+     * <p>
+     * This can be used to require that KeyStore entries be stored encrypted.
+     * <p>
+     * Example:
+     *
+     * <pre class="prettyprint">
+     * AndroidKeyStoreParameter params =
+     *         new AndroidKeyStoreParameter.Builder(mContext).setEncryptionRequired().build();
+     * </pre>
+     */
+    public final static class Builder {
+        private int mFlags;
+
+        /**
+         * Creates a new instance of the {@code Builder} with the given
+         * {@code context}. The {@code context} passed in may be used to pop up
+         * some UI to ask the user to unlock or initialize the Android KeyStore
+         * facility.
+         */
+        public Builder(Context context) {
+            if (context == null) {
+                throw new NullPointerException("context == null");
+            }
+
+            // Context is currently not used, but will be in the future.
+        }
+
+        /**
+         * Indicates that this key must be encrypted at rest on storage. Note
+         * that enabling this will require that the user enable a strong lock
+         * screen (e.g., PIN, password) before creating or using the generated
+         * key is successful.
+         */
+        public Builder setEncryptionRequired() {
+            mFlags |= KeyStore.FLAG_ENCRYPTED;
+            return this;
+        }
+
+        /**
+         * Builds the instance of the {@code AndroidKeyPairGeneratorSpec}.
+         *
+         * @throws IllegalArgumentException if a required field is missing
+         * @return built instance of {@code AndroidKeyPairGeneratorSpec}
+         */
+        public AndroidKeyStoreParameter build() {
+            return new AndroidKeyStoreParameter(mFlags);
+        }
+    }
+}
diff --git a/keystore/java/android/security/AndroidKeyStoreProvider.java b/keystore/java/android/security/AndroidKeyStoreProvider.java
index 40d7e1a..8ca301e 100644
--- a/keystore/java/android/security/AndroidKeyStoreProvider.java
+++ b/keystore/java/android/security/AndroidKeyStoreProvider.java
@@ -33,7 +33,6 @@
         put("KeyStore." + AndroidKeyStore.NAME, AndroidKeyStore.class.getName());
 
         // java.security.KeyPairGenerator
-        put("KeyPairGenerator." + AndroidKeyPairGenerator.NAME,
-                AndroidKeyPairGenerator.class.getName());
+        put("KeyPairGenerator." + AndroidKeyStore.NAME, AndroidKeyPairGenerator.class.getName());
     }
 }
diff --git a/keystore/java/android/security/KeyStore.java b/keystore/java/android/security/KeyStore.java
index 309d3d3..45385ee 100644
--- a/keystore/java/android/security/KeyStore.java
+++ b/keystore/java/android/security/KeyStore.java
@@ -40,7 +40,11 @@
     public static final int UNDEFINED_ACTION = 9;
     public static final int WRONG_PASSWORD = 10;
 
-    // Flags for "put" and "import"
+    // Used for UID field to indicate the calling UID.
+    public static final int UID_SELF = -1;
+
+    // Flags for "put" "import" and "generate"
+    public static final int FLAG_NONE = 0;
     public static final int FLAG_ENCRYPTED = 1;
 
     // States
@@ -104,7 +108,7 @@
     }
 
     public boolean put(String key, byte[] value) {
-        return put(key, value, -1);
+        return put(key, value, UID_SELF);
     }
 
     public boolean delete(String key, int uid) {
@@ -117,7 +121,7 @@
     }
 
     public boolean delete(String key) {
-        return delete(key, -1);
+        return delete(key, UID_SELF);
     }
 
     public boolean contains(String key, int uid) {
@@ -130,7 +134,7 @@
     }
 
     public boolean contains(String key) {
-        return contains(key, -1);
+        return contains(key, UID_SELF);
     }
 
     public String[] saw(String prefix, int uid) {
@@ -143,7 +147,7 @@
     }
 
     public String[] saw(String prefix) {
-        return saw(prefix, -1);
+        return saw(prefix, UID_SELF);
     }
 
     public boolean reset() {
@@ -206,7 +210,7 @@
     }
 
     public boolean generate(String key) {
-        return generate(key, -1);
+        return generate(key, UID_SELF);
     }
 
     public boolean importKey(String keyName, byte[] key, int uid, int flags) {
@@ -223,7 +227,7 @@
     }
 
     public boolean importKey(String keyName, byte[] key) {
-        return importKey(keyName, key, -1);
+        return importKey(keyName, key, UID_SELF);
     }
 
     public byte[] getPubkey(String key) {
@@ -245,7 +249,7 @@
     }
 
     public boolean delKey(String key) {
-        return delKey(key, -1);
+        return delKey(key, UID_SELF);
     }
 
     public byte[] sign(String key, byte[] data) {
diff --git a/keystore/tests/src/android/security/AndroidKeyPairGeneratorSpecTest.java b/keystore/tests/src/android/security/AndroidKeyPairGeneratorSpecTest.java
index 3d275cd..5d4ab9c 100644
--- a/keystore/tests/src/android/security/AndroidKeyPairGeneratorSpecTest.java
+++ b/keystore/tests/src/android/security/AndroidKeyPairGeneratorSpecTest.java
@@ -39,8 +39,9 @@
     private static final Date NOW_PLUS_10_YEARS = new Date(NOW.getYear() + 10, 0, 1);
 
     public void testConstructor_Success() throws Exception {
-        AndroidKeyPairGeneratorSpec spec = new AndroidKeyPairGeneratorSpec(getContext(),
-                TEST_ALIAS_1, TEST_DN_1, SERIAL_1, NOW, NOW_PLUS_10_YEARS);
+        AndroidKeyPairGeneratorSpec spec =
+                new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1, TEST_DN_1, SERIAL_1,
+                        NOW, NOW_PLUS_10_YEARS, 0);
 
         assertEquals("Context should be the one specified", getContext(), spec.getContext());
 
@@ -60,6 +61,7 @@
                 .setSerialNumber(SERIAL_1)
                 .setStartDate(NOW)
                 .setEndDate(NOW_PLUS_10_YEARS)
+                .setEncryptionRequired()
                 .build();
 
         assertEquals("Context should be the one specified", getContext(), spec.getContext());
@@ -71,12 +73,14 @@
         assertEquals("startDate should be the one specified", NOW, spec.getStartDate());
 
         assertEquals("endDate should be the one specified", NOW_PLUS_10_YEARS, spec.getEndDate());
+
+        assertEquals("encryption flag should be on", KeyStore.FLAG_ENCRYPTED, spec.getFlags());
     }
 
     public void testConstructor_NullContext_Failure() throws Exception {
         try {
             new AndroidKeyPairGeneratorSpec(null, TEST_ALIAS_1, TEST_DN_1, SERIAL_1, NOW,
-                    NOW_PLUS_10_YEARS);
+                    NOW_PLUS_10_YEARS, 0);
             fail("Should throw IllegalArgumentException when context is null");
         } catch (IllegalArgumentException success) {
         }
@@ -85,7 +89,7 @@
     public void testConstructor_NullKeystoreAlias_Failure() throws Exception {
         try {
             new AndroidKeyPairGeneratorSpec(getContext(), null, TEST_DN_1, SERIAL_1, NOW,
-                    NOW_PLUS_10_YEARS);
+                    NOW_PLUS_10_YEARS, 0);
             fail("Should throw IllegalArgumentException when keystoreAlias is null");
         } catch (IllegalArgumentException success) {
         }
@@ -94,7 +98,7 @@
     public void testConstructor_NullSubjectDN_Failure() throws Exception {
         try {
             new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1, null, SERIAL_1, NOW,
-                    NOW_PLUS_10_YEARS);
+                    NOW_PLUS_10_YEARS, 0);
             fail("Should throw IllegalArgumentException when subjectDN is null");
         } catch (IllegalArgumentException success) {
         }
@@ -103,7 +107,7 @@
     public void testConstructor_NullSerial_Failure() throws Exception {
         try {
             new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1, TEST_DN_1, null, NOW,
-                    NOW_PLUS_10_YEARS);
+                    NOW_PLUS_10_YEARS, 0);
             fail("Should throw IllegalArgumentException when startDate is null");
         } catch (IllegalArgumentException success) {
         }
@@ -112,7 +116,7 @@
     public void testConstructor_NullStartDate_Failure() throws Exception {
         try {
             new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1, TEST_DN_1, SERIAL_1, null,
-                    NOW_PLUS_10_YEARS);
+                    NOW_PLUS_10_YEARS, 0);
             fail("Should throw IllegalArgumentException when startDate is null");
         } catch (IllegalArgumentException success) {
         }
@@ -121,7 +125,7 @@
     public void testConstructor_NullEndDate_Failure() throws Exception {
         try {
             new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1, TEST_DN_1, SERIAL_1, NOW,
-                    null);
+                    null, 0);
             fail("Should throw IllegalArgumentException when keystoreAlias is null");
         } catch (IllegalArgumentException success) {
         }
@@ -130,7 +134,7 @@
     public void testConstructor_EndBeforeStart_Failure() throws Exception {
         try {
             new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1, TEST_DN_1, SERIAL_1,
-                    NOW_PLUS_10_YEARS, NOW);
+                    NOW_PLUS_10_YEARS, NOW, 0);
             fail("Should throw IllegalArgumentException when end is before start");
         } catch (IllegalArgumentException success) {
         }
diff --git a/keystore/tests/src/android/security/AndroidKeyPairGeneratorTest.java b/keystore/tests/src/android/security/AndroidKeyPairGeneratorTest.java
index 69007c4..c5cf514 100644
--- a/keystore/tests/src/android/security/AndroidKeyPairGeneratorTest.java
+++ b/keystore/tests/src/android/security/AndroidKeyPairGeneratorTest.java
@@ -27,6 +27,7 @@
 import java.security.cert.Certificate;
 import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
+import java.text.SimpleDateFormat;
 import java.util.Date;
 
 import javax.security.auth.x500.X500Principal;
@@ -64,22 +65,34 @@
 
         assertFalse(mAndroidKeyStore.isUnlocked());
 
+        mGenerator = java.security.KeyPairGenerator.getInstance("AndroidKeyStore");
+    }
+
+    private void setupPassword() {
         assertTrue(mAndroidKeyStore.password("1111"));
         assertTrue(mAndroidKeyStore.isUnlocked());
 
         String[] aliases = mAndroidKeyStore.saw("");
         assertNotNull(aliases);
         assertEquals(0, aliases.length);
-
-        mGenerator = java.security.KeyPairGenerator.getInstance(AndroidKeyPairGenerator.NAME);
     }
 
-    public void testKeyPairGenerator_Initialize_Params_Success() throws Exception {
-        mGenerator.initialize(new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1,
-                TEST_DN_1, TEST_SERIAL_1, NOW, NOW_PLUS_10_YEARS));
+    public void testKeyPairGenerator_Initialize_Params_Encrypted_Success() throws Exception {
+        setupPassword();
+
+        mGenerator.initialize(new AndroidKeyPairGeneratorSpec.Builder(getContext())
+                .setAlias(TEST_ALIAS_1)
+                .setSubject(TEST_DN_1)
+                .setSerialNumber(TEST_SERIAL_1)
+                .setStartDate(NOW)
+                .setEndDate(NOW_PLUS_10_YEARS)
+                .setEncryptionRequired()
+                .build());
     }
 
-    public void testKeyPairGenerator_Initialize_KeySize_Failure() throws Exception {
+    public void testKeyPairGenerator_Initialize_KeySize_Encrypted_Failure() throws Exception {
+        setupPassword();
+
         try {
             mGenerator.initialize(1024);
             fail("KeyPairGenerator should not support setting the key size");
@@ -87,7 +100,10 @@
         }
     }
 
-    public void testKeyPairGenerator_Initialize_KeySizeAndSecureRandom_Failure() throws Exception {
+    public void testKeyPairGenerator_Initialize_KeySizeAndSecureRandom_Encrypted_Failure()
+            throws Exception {
+        setupPassword();
+
         try {
             mGenerator.initialize(1024, new SecureRandom());
             fail("KeyPairGenerator should not support setting the key size");
@@ -95,14 +111,48 @@
         }
     }
 
-    public void testKeyPairGenerator_Initialize_ParamsAndSecureRandom_Failure() throws Exception {
-        mGenerator.initialize(new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1,
-                TEST_DN_1, TEST_SERIAL_1, NOW, NOW_PLUS_10_YEARS), new SecureRandom());
+    public void testKeyPairGenerator_Initialize_ParamsAndSecureRandom_Encrypted_Failure()
+            throws Exception {
+        setupPassword();
+
+        mGenerator.initialize(
+                new AndroidKeyPairGeneratorSpec.Builder(getContext())
+                        .setAlias(TEST_ALIAS_1)
+                        .setSubject(TEST_DN_1)
+                        .setSerialNumber(TEST_SERIAL_1)
+                        .setStartDate(NOW)
+                        .setEndDate(NOW_PLUS_10_YEARS)
+                        .setEncryptionRequired()
+                        .build(),
+                new SecureRandom());
     }
 
-    public void testKeyPairGenerator_GenerateKeyPair_Success() throws Exception {
-        mGenerator.initialize(new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1,
-                TEST_DN_1, TEST_SERIAL_1, NOW, NOW_PLUS_10_YEARS));
+    public void testKeyPairGenerator_GenerateKeyPair_Encrypted_Success() throws Exception {
+        setupPassword();
+
+        mGenerator.initialize(new AndroidKeyPairGeneratorSpec.Builder(getContext())
+                .setAlias(TEST_ALIAS_1)
+                .setSubject(TEST_DN_1)
+                .setSerialNumber(TEST_SERIAL_1)
+                .setStartDate(NOW)
+                .setEndDate(NOW_PLUS_10_YEARS)
+                .setEncryptionRequired()
+                .build());
+
+        final KeyPair pair = mGenerator.generateKeyPair();
+        assertNotNull("The KeyPair returned should not be null", pair);
+
+        assertKeyPairCorrect(pair, TEST_ALIAS_1, TEST_DN_1, TEST_SERIAL_1, NOW, NOW_PLUS_10_YEARS);
+    }
+
+    public void testKeyPairGenerator_GenerateKeyPair_Unencrypted_Success() throws Exception {
+        mGenerator.initialize(new AndroidKeyPairGeneratorSpec.Builder(getContext())
+                .setAlias(TEST_ALIAS_1)
+                .setSubject(TEST_DN_1)
+                .setSerialNumber(TEST_SERIAL_1)
+                .setStartDate(NOW)
+                .setEndDate(NOW_PLUS_10_YEARS)
+                .build());
 
         final KeyPair pair = mGenerator.generateKeyPair();
         assertNotNull("The KeyPair returned should not be null", pair);
@@ -113,8 +163,13 @@
     public void testKeyPairGenerator_GenerateKeyPair_Replaced_Success() throws Exception {
         // Generate the first key
         {
-            mGenerator.initialize(new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_1,
-                    TEST_DN_1, TEST_SERIAL_1, NOW, NOW_PLUS_10_YEARS));
+            mGenerator.initialize(new AndroidKeyPairGeneratorSpec.Builder(getContext())
+                    .setAlias(TEST_ALIAS_1)
+                    .setSubject(TEST_DN_1)
+                    .setSerialNumber(TEST_SERIAL_1)
+                    .setStartDate(NOW)
+                    .setEndDate(NOW_PLUS_10_YEARS)
+                    .build());
             final KeyPair pair1 = mGenerator.generateKeyPair();
             assertNotNull("The KeyPair returned should not be null", pair1);
             assertKeyPairCorrect(pair1, TEST_ALIAS_1, TEST_DN_1, TEST_SERIAL_1, NOW,
@@ -123,8 +178,13 @@
 
         // Replace the original key
         {
-            mGenerator.initialize(new AndroidKeyPairGeneratorSpec(getContext(), TEST_ALIAS_2,
-                    TEST_DN_2, TEST_SERIAL_2, NOW, NOW_PLUS_10_YEARS));
+            mGenerator.initialize(new AndroidKeyPairGeneratorSpec.Builder(getContext())
+                    .setAlias(TEST_ALIAS_2)
+                    .setSubject(TEST_DN_2)
+                    .setSerialNumber(TEST_SERIAL_2)
+                    .setStartDate(NOW)
+                    .setEndDate(NOW_PLUS_10_YEARS)
+                    .build());
             final KeyPair pair2 = mGenerator.generateKeyPair();
             assertNotNull("The KeyPair returned should not be null", pair2);
             assertKeyPairCorrect(pair2, TEST_ALIAS_2, TEST_DN_2, TEST_SERIAL_2, NOW,
@@ -132,6 +192,49 @@
         }
     }
 
+    public void testKeyPairGenerator_GenerateKeyPair_Replaced_UnencryptedToEncrypted_Success()
+            throws Exception {
+        // Generate the first key
+        {
+            mGenerator.initialize(new AndroidKeyPairGeneratorSpec.Builder(getContext())
+                    .setAlias(TEST_ALIAS_1)
+                    .setSubject(TEST_DN_1)
+                    .setSerialNumber(TEST_SERIAL_1)
+                    .setStartDate(NOW)
+                    .setEndDate(NOW_PLUS_10_YEARS)
+                    .build());
+            final KeyPair pair1 = mGenerator.generateKeyPair();
+            assertNotNull("The KeyPair returned should not be null", pair1);
+            assertKeyPairCorrect(pair1, TEST_ALIAS_1, TEST_DN_1, TEST_SERIAL_1, NOW,
+                    NOW_PLUS_10_YEARS);
+        }
+
+        // Attempt to replace previous key
+        {
+            mGenerator.initialize(new AndroidKeyPairGeneratorSpec.Builder(getContext())
+                    .setAlias(TEST_ALIAS_1)
+                    .setSubject(TEST_DN_2)
+                    .setSerialNumber(TEST_SERIAL_2)
+                    .setStartDate(NOW)
+                    .setEndDate(NOW_PLUS_10_YEARS)
+                    .setEncryptionRequired()
+                    .build());
+            try {
+                mGenerator.generateKeyPair();
+                fail("Should not be able to generate encrypted key while not initialized");
+            } catch (IllegalStateException expected) {
+            }
+
+            assertTrue(mAndroidKeyStore.password("1111"));
+            assertTrue(mAndroidKeyStore.isUnlocked());
+
+            final KeyPair pair2 = mGenerator.generateKeyPair();
+            assertNotNull("The KeyPair returned should not be null", pair2);
+            assertKeyPairCorrect(pair2, TEST_ALIAS_1, TEST_DN_2, TEST_SERIAL_2, NOW,
+                    NOW_PLUS_10_YEARS);
+        }
+    }
+
     private void assertKeyPairCorrect(KeyPair pair, String alias, X500Principal dn,
             BigInteger serial, Date start, Date end) throws Exception {
         final PublicKey pubKey = pair.getPublic();
@@ -163,10 +266,10 @@
         assertEquals("The Serial should be the one passed into the params", serial,
                 x509userCert.getSerialNumber());
 
-        assertEquals("The notBefore date should be the one passed into the params", start,
+        assertDateEquals("The notBefore date should be the one passed into the params", start,
                 x509userCert.getNotBefore());
 
-        assertEquals("The notAfter date should be the one passed into the params", end,
+        assertDateEquals("The notAfter date should be the one passed into the params", end,
                 x509userCert.getNotAfter());
 
         x509userCert.verify(pubKey);
@@ -178,4 +281,13 @@
         assertNotNull("The keystore should return the public key for the generated key",
                 pubKeyBytes);
     }
+
+    private static void assertDateEquals(String message, Date date1, Date date2) throws Exception {
+        SimpleDateFormat formatter = new SimpleDateFormat("dd MMM yyyy HH:mm:ss");
+
+        String result1 = formatter.format(date1);
+        String result2 = formatter.format(date2);
+
+        assertEquals(message, result1, result2);
+    }
 }
diff --git a/keystore/tests/src/android/security/AndroidKeyStoreTest.java b/keystore/tests/src/android/security/AndroidKeyStoreTest.java
index 8928e06..05ffe109 100644
--- a/keystore/tests/src/android/security/AndroidKeyStoreTest.java
+++ b/keystore/tests/src/android/security/AndroidKeyStoreTest.java
@@ -469,12 +469,14 @@
         assertTrue(mAndroidKeyStore.reset());
         assertFalse(mAndroidKeyStore.isUnlocked());
 
+        mKeyStore = java.security.KeyStore.getInstance("AndroidKeyStore");
+    }
+
+    private void setupPassword() {
         assertTrue(mAndroidKeyStore.password("1111"));
         assertTrue(mAndroidKeyStore.isUnlocked());
 
         assertEquals(0, mAndroidKeyStore.saw("").length);
-
-        mKeyStore = java.security.KeyStore.getInstance(AndroidKeyStore.NAME);
     }
 
     private void assertAliases(final String[] expectedAliases) throws KeyStoreException {
@@ -495,7 +497,9 @@
                 expectedAliases.length, count);
     }
 
-    public void testKeyStore_Aliases_Success() throws Exception {
+    public void testKeyStore_Aliases_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertAliases(new String[] {});
@@ -509,7 +513,9 @@
         assertAliases(new String[] { TEST_ALIAS_1, TEST_ALIAS_2 });
     }
 
-    public void testKeyStore_Aliases_NotInitialized_Failure() throws Exception {
+    public void testKeyStore_Aliases_NotInitialized_Encrypted_Failure() throws Exception {
+        setupPassword();
+
         try {
             mKeyStore.aliases();
             fail("KeyStore should throw exception when not initialized");
@@ -517,7 +523,9 @@
         }
     }
 
-    public void testKeyStore_ContainsAliases_PrivateAndCA_Success() throws Exception {
+    public void testKeyStore_ContainsAliases_PrivateAndCA_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertAliases(new String[] {});
@@ -534,7 +542,9 @@
                 mKeyStore.containsAlias(TEST_ALIAS_3));
     }
 
-    public void testKeyStore_ContainsAliases_CAOnly_Success() throws Exception {
+    public void testKeyStore_ContainsAliases_CAOnly_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_2, FAKE_CA_1));
@@ -542,13 +552,17 @@
         assertTrue("Should contain added CA certificate", mKeyStore.containsAlias(TEST_ALIAS_2));
     }
 
-    public void testKeyStore_ContainsAliases_NonExistent_Failure() throws Exception {
+    public void testKeyStore_ContainsAliases_NonExistent_Encrypted_Failure() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertFalse("Should contain added CA certificate", mKeyStore.containsAlias(TEST_ALIAS_1));
     }
 
-    public void testKeyStore_DeleteEntry_Success() throws Exception {
+    public void testKeyStore_DeleteEntry_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         // TEST_ALIAS_1
@@ -578,14 +592,18 @@
         assertAliases(new String[] { });
     }
 
-    public void testKeyStore_DeleteEntry_EmptyStore_Success() throws Exception {
+    public void testKeyStore_DeleteEntry_EmptyStore_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         // Should not throw when a non-existent entry is requested for delete.
         mKeyStore.deleteEntry(TEST_ALIAS_1);
     }
 
-    public void testKeyStore_DeleteEntry_NonExistent_Success() throws Exception {
+    public void testKeyStore_DeleteEntry_NonExistent_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         // TEST_ALIAS_1
@@ -598,7 +616,9 @@
         mKeyStore.deleteEntry(TEST_ALIAS_2);
     }
 
-    public void testKeyStore_GetCertificate_Single_Success() throws Exception {
+    public void testKeyStore_GetCertificate_Single_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1));
@@ -618,14 +638,18 @@
         assertEquals("Actual and retrieved certificates should be the same", actual, retrieved);
     }
 
-    public void testKeyStore_GetCertificate_NonExist_Failure() throws Exception {
+    public void testKeyStore_GetCertificate_NonExist_Encrypted_Failure() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertNull("Certificate should not exist in keystore",
                 mKeyStore.getCertificate(TEST_ALIAS_1));
     }
 
-    public void testKeyStore_GetCertificateAlias_CAEntry_Success() throws Exception {
+    public void testKeyStore_GetCertificateAlias_CAEntry_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1));
@@ -637,7 +661,10 @@
                 mKeyStore.getCertificateAlias(actual));
     }
 
-    public void testKeyStore_GetCertificateAlias_PrivateKeyEntry_Success() throws Exception {
+    public void testKeyStore_GetCertificateAlias_PrivateKeyEntry_Encrypted_Success()
+            throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
@@ -652,8 +679,10 @@
                 mKeyStore.getCertificateAlias(actual));
     }
 
-    public void testKeyStore_GetCertificateAlias_CAEntry_WithPrivateKeyUsingCA_Success()
+    public void testKeyStore_GetCertificateAlias_CAEntry_WithPrivateKeyUsingCA_Encrypted_Success()
             throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         // Insert TrustedCertificateEntry with CA name
@@ -672,7 +701,10 @@
                 mKeyStore.getCertificateAlias(actual));
     }
 
-    public void testKeyStore_GetCertificateAlias_NonExist_Empty_Failure() throws Exception {
+    public void testKeyStore_GetCertificateAlias_NonExist_Empty_Encrypted_Failure()
+            throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         CertificateFactory f = CertificateFactory.getInstance("X.509");
@@ -682,7 +714,9 @@
                 mKeyStore.getCertificateAlias(actual));
     }
 
-    public void testKeyStore_GetCertificateAlias_NonExist_Failure() throws Exception {
+    public void testKeyStore_GetCertificateAlias_NonExist_Encrypted_Failure() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1));
@@ -694,7 +728,9 @@
                 mKeyStore.getCertificateAlias(userCert));
     }
 
-    public void testKeyStore_GetCertificateChain_SingleLength_Success() throws Exception {
+    public void testKeyStore_GetCertificateChain_SingleLength_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
@@ -720,14 +756,18 @@
                 mKeyStore.getCertificateChain(TEST_ALIAS_2));
     }
 
-    public void testKeyStore_GetCertificateChain_NonExist_Failure() throws Exception {
+    public void testKeyStore_GetCertificateChain_NonExist_Encrypted_Failure() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertNull("Stored certificate alias should not be found",
                 mKeyStore.getCertificateChain(TEST_ALIAS_1));
     }
 
-    public void testKeyStore_GetCreationDate_PrivateKeyEntry_Success() throws Exception {
+    public void testKeyStore_GetCreationDate_PrivateKeyEntry_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
@@ -745,7 +785,29 @@
         assertTrue("Time should be close to current time", actual.after(expectedAfter));
     }
 
-    public void testKeyStore_GetCreationDate_CAEntry_Success() throws Exception {
+    public void testKeyStore_GetCreationDate_PrivateKeyEntry_Unencrypted_Success() throws Exception {
+        mKeyStore.load(null, null);
+
+        assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
+                FAKE_KEY_1, KeyStore.UID_SELF, KeyStore.FLAG_NONE));
+        assertTrue(mAndroidKeyStore.put(Credentials.USER_CERTIFICATE + TEST_ALIAS_1, FAKE_USER_1,
+                KeyStore.UID_SELF, KeyStore.FLAG_NONE));
+        assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1,
+                KeyStore.UID_SELF, KeyStore.FLAG_NONE));
+
+        Date now = new Date();
+        Date actual = mKeyStore.getCreationDate(TEST_ALIAS_1);
+
+        Date expectedAfter = new Date(now.getTime() - SLOP_TIME_MILLIS);
+        Date expectedBefore = new Date(now.getTime() + SLOP_TIME_MILLIS);
+
+        assertTrue("Time should be close to current time", actual.before(expectedBefore));
+        assertTrue("Time should be close to current time", actual.after(expectedAfter));
+    }
+
+    public void testKeyStore_GetCreationDate_CAEntry_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1));
@@ -761,7 +823,9 @@
         assertTrue("Time should be close to current time", actual.after(expectedAfter));
     }
 
-    public void testKeyStore_GetEntry_NullParams_Success() throws Exception {
+    public void testKeyStore_GetEntry_NullParams_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
@@ -779,6 +843,26 @@
         assertPrivateKeyEntryEquals(keyEntry, FAKE_KEY_1, FAKE_USER_1, FAKE_CA_1);
     }
 
+    public void testKeyStore_GetEntry_NullParams_Unencrypted_Success() throws Exception {
+        mKeyStore.load(null, null);
+
+        assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
+                FAKE_KEY_1, KeyStore.UID_SELF, KeyStore.FLAG_NONE));
+        assertTrue(mAndroidKeyStore.put(Credentials.USER_CERTIFICATE + TEST_ALIAS_1, FAKE_USER_1,
+                KeyStore.UID_SELF, KeyStore.FLAG_NONE));
+        assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1,
+                KeyStore.UID_SELF, KeyStore.FLAG_NONE));
+
+        Entry entry = mKeyStore.getEntry(TEST_ALIAS_1, null);
+        assertNotNull("Entry should exist", entry);
+
+        assertTrue("Should be a PrivateKeyEntry", entry instanceof PrivateKeyEntry);
+
+        PrivateKeyEntry keyEntry = (PrivateKeyEntry) entry;
+
+        assertPrivateKeyEntryEquals(keyEntry, FAKE_KEY_1, FAKE_USER_1, FAKE_CA_1);
+    }
+
     @SuppressWarnings("unchecked")
     private void assertPrivateKeyEntryEquals(PrivateKeyEntry keyEntry, byte[] key, byte[] cert,
             byte[] ca) throws Exception {
@@ -801,8 +885,9 @@
 
     private void assertPrivateKeyEntryEquals(PrivateKeyEntry keyEntry, PrivateKey expectedKey,
             Certificate expectedCert, Collection<Certificate> expectedChain) throws Exception {
-        assertEquals("Returned PrivateKey should be what we inserted", expectedKey,
-                keyEntry.getPrivateKey());
+        assertEquals("Returned PrivateKey should be what we inserted",
+                ((RSAPrivateKey) expectedKey).getModulus(),
+                ((RSAPrivateKey) keyEntry.getPrivateKey()).getModulus());
 
         assertEquals("Returned Certificate should be what we inserted", expectedCert,
                 keyEntry.getCertificate());
@@ -823,14 +908,25 @@
         }
     }
 
-    public void testKeyStore_GetEntry_Nonexistent_NullParams_Failure() throws Exception {
+    public void testKeyStore_GetEntry_Nonexistent_NullParams_Encrypted_Failure() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertNull("A non-existent entry should return null",
                 mKeyStore.getEntry(TEST_ALIAS_1, null));
     }
 
-    public void testKeyStore_GetKey_NoPassword_Success() throws Exception {
+    public void testKeyStore_GetEntry_Nonexistent_NullParams_Unencrypted_Failure() throws Exception {
+        mKeyStore.load(null, null);
+
+        assertNull("A non-existent entry should return null",
+                mKeyStore.getEntry(TEST_ALIAS_1, null));
+    }
+
+    public void testKeyStore_GetKey_NoPassword_Encrypted_Success() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
@@ -848,10 +944,37 @@
         KeyFactory keyFact = KeyFactory.getInstance("RSA");
         PrivateKey expectedKey = keyFact.generatePrivate(new PKCS8EncodedKeySpec(FAKE_KEY_1));
 
-        assertEquals("Inserted key should be same as retrieved key", actualKey, expectedKey);
+        assertEquals("Inserted key should be same as retrieved key",
+                ((RSAPrivateKey) expectedKey).getModulus(), actualKey.getModulus());
     }
 
-    public void testKeyStore_GetKey_Certificate_Failure() throws Exception {
+    public void testKeyStore_GetKey_NoPassword_Unencrypted_Success() throws Exception {
+        mKeyStore.load(null, null);
+
+        assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
+                FAKE_KEY_1, KeyStore.UID_SELF, KeyStore.FLAG_NONE));
+        assertTrue(mAndroidKeyStore.put(Credentials.USER_CERTIFICATE + TEST_ALIAS_1, FAKE_USER_1,
+                KeyStore.UID_SELF, KeyStore.FLAG_NONE));
+        assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1,
+                KeyStore.UID_SELF, KeyStore.FLAG_NONE));
+
+        Key key = mKeyStore.getKey(TEST_ALIAS_1, null);
+        assertNotNull("Key should exist", key);
+
+        assertTrue("Should be a RSAPrivateKey", key instanceof RSAPrivateKey);
+
+        RSAPrivateKey actualKey = (RSAPrivateKey) key;
+
+        KeyFactory keyFact = KeyFactory.getInstance("RSA");
+        PrivateKey expectedKey = keyFact.generatePrivate(new PKCS8EncodedKeySpec(FAKE_KEY_1));
+
+        assertEquals("Inserted key should be same as retrieved key",
+                ((RSAPrivateKey) expectedKey).getModulus(), actualKey.getModulus());
+    }
+
+    public void testKeyStore_GetKey_Certificate_Encrypted_Failure() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1));
@@ -859,21 +982,28 @@
         assertNull("Certificate entries should return null", mKeyStore.getKey(TEST_ALIAS_1, null));
     }
 
-    public void testKeyStore_GetKey_NonExistent_Failure() throws Exception {
+    public void testKeyStore_GetKey_NonExistent_Encrypted_Failure() throws Exception {
+        setupPassword();
+
         mKeyStore.load(null, null);
 
         assertNull("A non-existent entry should return null", mKeyStore.getKey(TEST_ALIAS_1, null));
     }
 
-    public void testKeyStore_GetProvider_Success() throws Exception {
+    public void testKeyStore_GetProvider_Encrypted_Success() throws Exception {
+        assertEquals(AndroidKeyStoreProvider.PROVIDER_NAME, mKeyStore.getProvider().getName());
+        setupPassword();
         assertEquals(AndroidKeyStoreProvider.PROVIDER_NAME, mKeyStore.getProvider().getName());
     }
 
-    public void testKeyStore_GetType_Success() throws Exception {
+    public void testKeyStore_GetType_Encrypted_Success() throws Exception {
+        assertEquals(AndroidKeyStore.NAME, mKeyStore.getType());
+        setupPassword();
         assertEquals(AndroidKeyStore.NAME, mKeyStore.getType());
     }
 
-    public void testKeyStore_IsCertificateEntry_CA_Success() throws Exception {
+    public void testKeyStore_IsCertificateEntry_CA_Encrypted_Success() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1));
@@ -882,7 +1012,8 @@
                 mKeyStore.isCertificateEntry(TEST_ALIAS_1));
     }
 
-    public void testKeyStore_IsCertificateEntry_PrivateKey_Failure() throws Exception {
+    public void testKeyStore_IsCertificateEntry_PrivateKey_Encrypted_Failure() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
@@ -894,14 +1025,23 @@
                 mKeyStore.isCertificateEntry(TEST_ALIAS_1));
     }
 
-    public void testKeyStore_IsCertificateEntry_NonExist_Failure() throws Exception {
+    public void testKeyStore_IsCertificateEntry_NonExist_Encrypted_Failure() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         assertFalse("Should return false for non-existent entry",
                 mKeyStore.isCertificateEntry(TEST_ALIAS_1));
     }
 
-    public void testKeyStore_IsKeyEntry_PrivateKey_Success() throws Exception {
+    public void testKeyStore_IsCertificateEntry_NonExist_Unencrypted_Failure() throws Exception {
+        mKeyStore.load(null, null);
+
+        assertFalse("Should return false for non-existent entry",
+                mKeyStore.isCertificateEntry(TEST_ALIAS_1));
+    }
+
+    public void testKeyStore_IsKeyEntry_PrivateKey_Encrypted_Success() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
@@ -912,7 +1052,8 @@
         assertTrue("Should return true for PrivateKeyEntry", mKeyStore.isKeyEntry(TEST_ALIAS_1));
     }
 
-    public void testKeyStore_IsKeyEntry_CA_Failure() throws Exception {
+    public void testKeyStore_IsKeyEntry_CA_Encrypted_Failure() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1));
@@ -920,17 +1061,19 @@
         assertFalse("Should return false for CA certificate", mKeyStore.isKeyEntry(TEST_ALIAS_1));
     }
 
-    public void testKeyStore_IsKeyEntry_NonExist_Failure() throws Exception {
+    public void testKeyStore_IsKeyEntry_NonExist_Encrypted_Failure() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         assertFalse("Should return false for non-existent entry",
                 mKeyStore.isKeyEntry(TEST_ALIAS_1));
     }
 
-    public void testKeyStore_SetCertificate_CA_Success() throws Exception {
+    public void testKeyStore_SetCertificate_CA_Encrypted_Success() throws Exception {
         final CertificateFactory f = CertificateFactory.getInstance("X.509");
         final Certificate actual = f.generateCertificate(new ByteArrayInputStream(FAKE_CA_1));
 
+        setupPassword();
         mKeyStore.load(null, null);
 
         mKeyStore.setCertificateEntry(TEST_ALIAS_1, actual);
@@ -942,7 +1085,8 @@
                 retrieved);
     }
 
-    public void testKeyStore_SetCertificate_CAExists_Overwrite_Success() throws Exception {
+    public void testKeyStore_SetCertificate_CAExists_Overwrite_Encrypted_Success() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1));
@@ -958,7 +1102,8 @@
         assertAliases(new String[] { TEST_ALIAS_1 });
     }
 
-    public void testKeyStore_SetCertificate_PrivateKeyExists_Failure() throws Exception {
+    public void testKeyStore_SetCertificate_PrivateKeyExists_Encrypted_Failure() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.importKey(Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1,
@@ -978,7 +1123,8 @@
         }
     }
 
-    public void testKeyStore_SetEntry_PrivateKeyEntry_Success() throws Exception {
+    public void testKeyStore_SetEntry_PrivateKeyEntry_Encrypted_Success() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         KeyFactory keyFact = KeyFactory.getInstance("RSA");
@@ -1005,8 +1151,63 @@
         assertPrivateKeyEntryEquals(actual, FAKE_KEY_1, FAKE_USER_1, FAKE_CA_1);
     }
 
-    public void testKeyStore_SetEntry_PrivateKeyEntry_Overwrites_PrivateKeyEntry_Success()
+    public void testKeyStore_SetEntry_PrivateKeyEntry_Unencrypted_Success() throws Exception {
+        mKeyStore.load(null, null);
+
+        KeyFactory keyFact = KeyFactory.getInstance("RSA");
+        PrivateKey expectedKey = keyFact.generatePrivate(new PKCS8EncodedKeySpec(FAKE_KEY_1));
+
+        final CertificateFactory f = CertificateFactory.getInstance("X.509");
+
+        final Certificate[] expectedChain = new Certificate[2];
+        expectedChain[0] = f.generateCertificate(new ByteArrayInputStream(FAKE_USER_1));
+        expectedChain[1] = f.generateCertificate(new ByteArrayInputStream(FAKE_CA_1));
+
+        PrivateKeyEntry expected = new PrivateKeyEntry(expectedKey, expectedChain);
+
+        mKeyStore.setEntry(TEST_ALIAS_1, expected, null);
+
+        Entry actualEntry = mKeyStore.getEntry(TEST_ALIAS_1, null);
+        assertNotNull("Retrieved entry should exist", actualEntry);
+
+        assertTrue("Retrieved entry should be of type PrivateKeyEntry",
+                actualEntry instanceof PrivateKeyEntry);
+
+        PrivateKeyEntry actual = (PrivateKeyEntry) actualEntry;
+
+        assertPrivateKeyEntryEquals(actual, FAKE_KEY_1, FAKE_USER_1, FAKE_CA_1);
+    }
+
+    public void testKeyStore_SetEntry_PrivateKeyEntry_Params_Unencrypted_Failure() throws Exception {
+        mKeyStore.load(null, null);
+
+        KeyFactory keyFact = KeyFactory.getInstance("RSA");
+        PrivateKey expectedKey = keyFact.generatePrivate(new PKCS8EncodedKeySpec(FAKE_KEY_1));
+
+        final CertificateFactory f = CertificateFactory.getInstance("X.509");
+
+        final Certificate[] expectedChain = new Certificate[2];
+        expectedChain[0] = f.generateCertificate(new ByteArrayInputStream(FAKE_USER_1));
+        expectedChain[1] = f.generateCertificate(new ByteArrayInputStream(FAKE_CA_1));
+
+        PrivateKeyEntry entry = new PrivateKeyEntry(expectedKey, expectedChain);
+
+        try {
+            mKeyStore.setEntry(TEST_ALIAS_1, entry,
+                    new AndroidKeyStoreParameter.Builder(getContext())
+                    .setEncryptionRequired()
+                    .build());
+            fail("Shouldn't be able to insert encrypted entry when KeyStore uninitialized");
+        } catch (KeyStoreException expected) {
+        }
+
+        assertNull(mKeyStore.getEntry(TEST_ALIAS_1, null));
+    }
+
+    public void
+            testKeyStore_SetEntry_PrivateKeyEntry_Overwrites_PrivateKeyEntry_Encrypted_Success()
             throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         final KeyFactory keyFact = KeyFactory.getInstance("RSA");
@@ -1060,7 +1261,9 @@
         }
     }
 
-    public void testKeyStore_SetEntry_CAEntry_Overwrites_PrivateKeyEntry_Success() throws Exception {
+    public void testKeyStore_SetEntry_CAEntry_Overwrites_PrivateKeyEntry_Encrypted_Success()
+            throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         final CertificateFactory f = CertificateFactory.getInstance("X.509");
@@ -1104,7 +1307,9 @@
         }
     }
 
-    public void testKeyStore_SetEntry_PrivateKeyEntry_Overwrites_CAEntry_Success() throws Exception {
+    public void testKeyStore_SetEntry_PrivateKeyEntry_Overwrites_CAEntry_Encrypted_Success()
+            throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         final CertificateFactory f = CertificateFactory.getInstance("X.509");
@@ -1148,8 +1353,11 @@
         }
     }
 
-    public void testKeyStore_SetEntry_PrivateKeyEntry_Overwrites_ShortPrivateKeyEntry_Success()
+    public
+            void
+            testKeyStore_SetEntry_PrivateKeyEntry_Overwrites_ShortPrivateKeyEntry_Encrypted_Success()
             throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         final CertificateFactory f = CertificateFactory.getInstance("X.509");
@@ -1198,7 +1406,9 @@
         }
     }
 
-    public void testKeyStore_SetEntry_CAEntry_Overwrites_CAEntry_Success() throws Exception {
+    public void testKeyStore_SetEntry_CAEntry_Overwrites_CAEntry_Encrypted_Success()
+            throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         final CertificateFactory f = CertificateFactory.getInstance("X.509");
@@ -1239,7 +1449,8 @@
         }
     }
 
-    public void testKeyStore_SetKeyEntry_ProtectedKey_Failure() throws Exception {
+    public void testKeyStore_SetKeyEntry_ProtectedKey_Encrypted_Failure() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         final CertificateFactory f = CertificateFactory.getInstance("X.509");
@@ -1259,7 +1470,8 @@
         }
     }
 
-    public void testKeyStore_SetKeyEntry_Success() throws Exception {
+    public void testKeyStore_SetKeyEntry_Encrypted_Success() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         final CertificateFactory f = CertificateFactory.getInstance("X.509");
@@ -1285,7 +1497,8 @@
         assertPrivateKeyEntryEquals(actual, FAKE_KEY_1, FAKE_USER_1, FAKE_CA_1);
     }
 
-    public void testKeyStore_SetKeyEntry_Replaced_Success() throws Exception {
+    public void testKeyStore_SetKeyEntry_Replaced_Encrypted_Success() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         final CertificateFactory f = CertificateFactory.getInstance("X.509");
@@ -1376,7 +1589,8 @@
         return cert;
     }
 
-    public void testKeyStore_SetKeyEntry_ReplacedChain_Success() throws Exception {
+    public void testKeyStore_SetKeyEntry_ReplacedChain_Encrypted_Success() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         // Create key #1
@@ -1429,8 +1643,9 @@
         }
     }
 
-    public void testKeyStore_SetKeyEntry_ReplacedChain_DifferentPrivateKey_Failure()
+    public void testKeyStore_SetKeyEntry_ReplacedChain_DifferentPrivateKey_Encrypted_Failure()
             throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         // Create key #1
@@ -1472,7 +1687,48 @@
         }
     }
 
-    public void testKeyStore_Size_Success() throws Exception {
+    public void testKeyStore_SetKeyEntry_ReplacedChain_UnencryptedToEncrypted_Failure()
+            throws Exception {
+        mKeyStore.load(null, null);
+
+        // Create key #1
+        {
+            final String privateKeyAlias = Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1;
+            assertTrue(mAndroidKeyStore.generate(privateKeyAlias,
+                    android.security.KeyStore.UID_SELF, android.security.KeyStore.FLAG_NONE));
+
+            X509Certificate cert =
+                    generateCertificate(mAndroidKeyStore, TEST_ALIAS_1, TEST_SERIAL_1, TEST_DN_1,
+                            NOW, NOW_PLUS_10_YEARS);
+
+            assertTrue(mAndroidKeyStore.put(Credentials.USER_CERTIFICATE + TEST_ALIAS_1,
+                    cert.getEncoded(), android.security.KeyStore.UID_SELF,
+                    android.security.KeyStore.FLAG_NONE));
+        }
+
+        // Replace with one that requires encryption
+        {
+            Entry entry = mKeyStore.getEntry(TEST_ALIAS_1, null);
+
+            try {
+                mKeyStore.setEntry(TEST_ALIAS_1, entry, new AndroidKeyStoreParameter.Builder(
+                        getContext()).setEncryptionRequired().build());
+                fail("Should not allow setting of Entry without unlocked keystore");
+            } catch (KeyStoreException success) {
+            }
+
+            assertTrue(mAndroidKeyStore.password("1111"));
+            assertTrue(mAndroidKeyStore.isUnlocked());
+
+            mKeyStore.setEntry(TEST_ALIAS_1, entry,
+                    new AndroidKeyStoreParameter.Builder(getContext())
+                            .setEncryptionRequired()
+                            .build());
+        }
+    }
+
+    public void testKeyStore_Size_Encrypted_Success() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         assertTrue(mAndroidKeyStore.put(Credentials.CA_CERTIFICATE + TEST_ALIAS_1, FAKE_CA_1));
@@ -1501,7 +1757,8 @@
         assertAliases(new String[] { TEST_ALIAS_2 });
     }
 
-    public void testKeyStore_Store_LoadStoreParam_Failure() throws Exception {
+    public void testKeyStore_Store_LoadStoreParam_Encrypted_Failure() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         try {
@@ -1511,7 +1768,7 @@
         }
     }
 
-    public void testKeyStore_Load_InputStreamSupplied_Failure() throws Exception {
+    public void testKeyStore_Load_InputStreamSupplied_Encrypted_Failure() throws Exception {
         byte[] buf = "FAKE KEYSTORE".getBytes();
         ByteArrayInputStream is = new ByteArrayInputStream(buf);
 
@@ -1522,7 +1779,7 @@
         }
     }
 
-    public void testKeyStore_Load_PasswordSupplied_Failure() throws Exception {
+    public void testKeyStore_Load_PasswordSupplied_Encrypted_Failure() throws Exception {
         try {
             mKeyStore.load(null, "password".toCharArray());
             fail("Should throw IllegalArgumentException when password is supplied");
@@ -1530,7 +1787,8 @@
         }
     }
 
-    public void testKeyStore_Store_OutputStream_Failure() throws Exception {
+    public void testKeyStore_Store_OutputStream_Encrypted_Failure() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         OutputStream sink = new ByteArrayOutputStream();
@@ -1558,7 +1816,8 @@
                 cert.getEncoded()));
     }
 
-    public void testKeyStore_KeyOperations_Wrap_Success() throws Exception {
+    public void testKeyStore_KeyOperations_Wrap_Encrypted_Success() throws Exception {
+        setupPassword();
         mKeyStore.load(null, null);
 
         setupKey();