diff --git a/robotests/Android.mk b/robotests/Android.mk
index a6b854f..6165044 100644
--- a/robotests/Android.mk
+++ b/robotests/Android.mk
@@ -14,7 +14,7 @@
 
 LOCAL_JAVA_LIBRARIES := \
     junit \
-    platform-robolectric-3.5.1-prebuilt \
+    platform-robolectric-3.6.1-prebuilt \
     telephony-common
 
 LOCAL_INSTRUMENTATION_FOR := KeyChain
@@ -42,4 +42,4 @@
 
 LOCAL_ROBOTEST_TIMEOUT := 36000
 
-include prebuilts/misc/common/robolectric/3.5.1/run_robotests.mk
+include prebuilts/misc/common/robolectric/3.6.1/run_robotests.mk
diff --git a/src/com/android/keychain/KeyChainService.java b/src/com/android/keychain/KeyChainService.java
index 3f695fe..4ce8378 100644
--- a/src/com/android/keychain/KeyChainService.java
+++ b/src/com/android/keychain/KeyChainService.java
@@ -39,6 +39,8 @@
 import android.security.keymaster.KeymasterCertificateChain;
 import android.security.keymaster.KeymasterDefs;
 import android.security.KeyStore;
+import android.security.keystore.AttestationUtils;
+import android.security.keystore.DeviceIdAttestationException;
 import android.security.keystore.KeyGenParameterSpec;
 import android.security.keystore.ParcelableKeyGenParameterSpec;
 import android.text.TextUtils;
@@ -89,6 +91,7 @@
         private final KeyStore mKeyStore = KeyStore.getInstance();
         private final TrustedCertificateStore mTrustedCertificateStore
                 = new TrustedCertificateStore();
+        private final Context mContext = KeyChainService.this;
 
         @Override
         public String requestPrivateKey(String alias) {
@@ -163,6 +166,7 @@
 
         @Override public boolean attestKey(
                 String alias, byte[] attestationChallenge,
+                int[] idAttestationFlags,
                 KeymasterCertificateChain attestationChain) {
             checkSystemCaller();
             validateAlias(alias);
@@ -172,8 +176,14 @@
                 return false;
             }
 
-            KeymasterArguments attestArgs = new KeymasterArguments();
-            attestArgs.addBytes(KeymasterDefs.KM_TAG_ATTESTATION_CHALLENGE, attestationChallenge);
+            final KeymasterArguments attestArgs;
+            try {
+                attestArgs = AttestationUtils.prepareAttestationArguments(
+                        mContext, idAttestationFlags, attestationChallenge);
+            } catch (DeviceIdAttestationException e) {
+                Log.e(TAG, "Failed collecting attestation data", e);
+                return false;
+            }
             final String keystoreAlias = Credentials.USER_PRIVATE_KEY + alias;
             final int errorCode = mKeyStore.attestKey(keystoreAlias, attestArgs, attestationChain);
             return errorCode == KeyStore.NO_ERROR;