From c07355fefb84fc997383cee10a87e57ca86a2bde Mon Sep 17 00:00:00 2001
From: Trevor Bentley <trevor@yubico.com>
Date: Fri, 20 Oct 2017 12:40:09 +0200
Subject: [PATCH] Fix unit tests for NEO: use ECCP256 and detect attestation
 errors

---
 lib/tests/api.c     | 21 ++++++++++++++++++---
 lib/ykpiv.c         |  3 +++
 lib/ykpiv.h         |  4 +++-
 tool/tests/basic.sh |  4 ++--
 4 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/lib/tests/api.c b/lib/tests/api.c
index 305a086..36e1e5e 100644
--- a/lib/tests/api.c
+++ b/lib/tests/api.c
@@ -409,8 +409,15 @@ START_TEST(test_import_key) {
   {
     unsigned char attest[2048];
     size_t attest_len = sizeof(attest);
+    ykpiv_devmodel model;
+    model = ykpiv_util_devicemodel(g_state);
     res = ykpiv_attest(g_state, 0x9e, attest, &attest_len);
-    ck_assert_int_eq(res, YKPIV_GENERIC_ERROR);
+    if (model == DEVTYPE_YK4) {
+      ck_assert_int_eq(res, YKPIV_GENERIC_ERROR);
+    }
+    else {
+      ck_assert_int_eq(res, YKPIV_NOT_SUPPORTED);
+    }
   }
 }
 END_TEST
@@ -440,11 +447,19 @@ START_TEST(test_generate_key) {
 
   // Verify that imported key can be attested
   {
+    ykpiv_devmodel model;
     unsigned char attest[2048];
     size_t attest_len = sizeof(attest);
+    model = ykpiv_util_devicemodel(g_state);
     res = ykpiv_attest(g_state, YKPIV_KEY_AUTHENTICATION, attest, &attest_len);
-    ck_assert_int_eq(res, YKPIV_OK);
-    ck_assert_int_gt(attest_len, 0);
+    // Only works with YK4.  NEO should error.
+    if (model == DEVTYPE_YK4) {
+      ck_assert_int_eq(res, YKPIV_OK);
+      ck_assert_int_gt(attest_len, 0);
+    }
+    else {
+      ck_assert_int_eq(res, YKPIV_NOT_SUPPORTED);
+    }
   }
 }
 END_TEST
diff --git a/lib/ykpiv.c b/lib/ykpiv.c
index ed31764..7948157 100644
--- a/lib/ykpiv.c
+++ b/lib/ykpiv.c
@@ -1346,6 +1346,9 @@ ykpiv_rc ykpiv_attest(ykpiv_state *state, const unsigned char key, unsigned char
     return res;
   }
   else if(SW_SUCCESS != sw) {
+    if (sw == SW_ERR_NOT_SUPPORTED) {
+      return YKPIV_NOT_SUPPORTED;
+    }
     return YKPIV_GENERIC_ERROR;
   }
   if (data[0] != 0x30) {
diff --git a/lib/ykpiv.h b/lib/ykpiv.h
index c5cae45..a2ddfa4 100644
--- a/lib/ykpiv.h
+++ b/lib/ykpiv.h
@@ -61,7 +61,8 @@ extern "C"
     YKPIV_PIN_LOCKED = -13,
 
     YKPIV_ARGUMENT_ERROR = -14, //i.e. invalid input argument
-    YKPIV_RANGE_ERROR = -15 //i.e. value range error
+    YKPIV_RANGE_ERROR = -15, //i.e. value range error
+    YKPIV_NOT_SUPPORTED = -16
   } ykpiv_rc;
 
   typedef void* (*ykpiv_pfn_alloc)(void* alloc_data, size_t size);
@@ -217,6 +218,7 @@ extern "C"
 #define SW_ERR_INCORRECT_PARAM 0x6a80
 /* this is a custom sw for yubikey */
 #define SW_ERR_INCORRECT_SLOT 0x6b00
+#define SW_ERR_NOT_SUPPORTED 0x6d00
 
   /* Yubico vendor specific instructions */
 #define YKPIV_INS_SET_MGMKEY 0xff
diff --git a/tool/tests/basic.sh b/tool/tests/basic.sh
index b605795..95a886a 100755
--- a/tool/tests/basic.sh
+++ b/tool/tests/basic.sh
@@ -111,7 +111,7 @@ $BIN -achange-puk -P000000 -N00000000 || true
 $BIN -areset
 
 # Generate key on-board, issue certificate, and verify it
-$BIN -agenerate -s9a -AECCP384 -o key_9a.pub
+$BIN -agenerate -s9a -AECCP256 -o key_9a.pub
 $BIN -averify -P123456 -s9a -S'/CN=YubicoTest/OU=YubicoGenerated/O=yubico.com/' -aselfsign -i key_9a.pub -o cert_9a.pem
 $BIN -averify -P123456 -s9a -atest-signature -i cert_9a.pem
 $BIN -aimport-certificate -P123456 -s9a -i cert_9a.pem
@@ -128,7 +128,7 @@ $BIN -aimport-certificate -P123456 -s9e -i cert.pem
 STATUS=$($BIN -astatus)
 echo "$STATUS"
 ALGO_9A=$(echo "$STATUS" |grep "Slot 9a" -A 6 |grep "Algorithm" |tr -d "[:blank:]")
-if [[ "x$ALGO_9A" != "xAlgorithm:ECCP384" ]]; then
+if [[ "x$ALGO_9A" != "xAlgorithm:ECCP256" ]]; then
     echo "$ALGO_9A"
     echo "Generated algorithm incorrect." >/dev/stderr
     exit 1