From b4dfb0239f338283cfa9df40191bae1497ac3210 Mon Sep 17 00:00:00 2001
From: Aveen Ismail <aveen.ismail@yubico.com>
Date: Thu, 5 Sep 2024 14:17:53 +0200
Subject: [PATCH] YubiHSM Auth: Add YubiKey firmware version check before
 getting device challenge

---
 src/commands.c        | 21 ++++++++++++++++++---
 ykhsmauth/ykhsmauth.c | 37 +++++++++++++++++++++++++++----------
 ykhsmauth/ykhsmauth.h |  3 +++
 3 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/src/commands.c b/src/commands.c
index dacc03cd..1ae6d996 100644
--- a/src/commands.c
+++ b/src/commands.c
@@ -1857,9 +1857,24 @@ int yh_com_open_yksession(yubihsm_context *ctx, Argument *argv,
   uint8_t host_challenge[YH_EC_P256_PUBKEY_LEN] = {0};
   size_t host_challenge_len = sizeof(host_challenge);
 
-  ykhsmauthrc = ykhsmauth_get_challenge_ex(ctx->state, argv[1].s,
-                                        argv[2].x, argv[2].len, 
-                                        host_challenge, &host_challenge_len);
+  uint8_t major = 0, minor = 0, patch = 0;
+  ykhsmauthrc = ykhsmauth_get_version_ex(ctx->state, &major, &minor, &patch);
+  if (ykhsmauthrc != YKHSMAUTHR_SUCCESS) {
+    fprintf(stderr, "Failed to get YubiKey firmware version: %s\n",
+            ykhsmauth_strerror(ykhsmauthrc));
+    ykhsmauth_disconnect(ctx->state);
+    return -1;
+  }
+
+  if (major > 5 || (major == 5 && minor > 7) ||
+      (major == 5 && minor == 7 && patch >= 1)) {
+    ykhsmauthrc =
+      ykhsmauth_get_challenge_ex(ctx->state, argv[1].s, argv[2].x, argv[2].len,
+                                 host_challenge, &host_challenge_len);
+  } else {
+    ykhsmauthrc = ykhsmauth_get_challenge(ctx->state, argv[1].s, host_challenge,
+                                          &host_challenge_len);
+  }
   if (ykhsmauthrc != YKHSMAUTHR_SUCCESS) {
     fprintf(stderr, "Failed to get host challenge from the YubiKey: %s\n",
             ykhsmauth_strerror(ykhsmauthrc));
diff --git a/ykhsmauth/ykhsmauth.c b/ykhsmauth/ykhsmauth.c
index 34a82b4a..76b4a306 100644
--- a/ykhsmauth/ykhsmauth.c
+++ b/ykhsmauth/ykhsmauth.c
@@ -296,9 +296,9 @@ ykhsmauth_rc ykhsmauth_list_readers(ykhsmauth_state *state, char *readers,
   return YKHSMAUTHR_SUCCESS;
 }
 
-ykhsmauth_rc ykhsmauth_get_version(ykhsmauth_state *state, char *version,
-                                   size_t len) {
-  if (state == NULL || version == NULL) {
+ykhsmauth_rc ykhsmauth_get_version_ex(ykhsmauth_state *state, uint8_t *major,
+                                      uint8_t *minor, uint8_t *patch) {
+  if (state == NULL || major == NULL || minor == NULL || patch == NULL) {
     return YKHSMAUTHR_INVALID_PARAMS;
   }
 
@@ -307,22 +307,37 @@ ykhsmauth_rc ykhsmauth_get_version(ykhsmauth_state *state, char *version,
   DWORD recv_len = sizeof(data);
   uint16_t sw = 0;
   ykhsmauth_rc res;
-
   if ((res = send_data(state, &apdu, data, &recv_len, &sw)) !=
       YKHSMAUTHR_SUCCESS) {
     return res;
   } else if (sw == SW_SUCCESS && recv_len == 3) {
-    int result = snprintf(version, len, "%d.%d.%d", data[0], data[1], data[2]);
+    *major = data[0];
+    *minor = data[1];
+    *patch = data[2];
+    return YKHSMAUTHR_SUCCESS;
+  } else {
+    return translate_error(sw, NULL);
+  }
+}
+
+ykhsmauth_rc ykhsmauth_get_version(ykhsmauth_state *state, char *version,
+                                   size_t len) {
+  if (version == NULL) {
+    return YKHSMAUTHR_INVALID_PARAMS;
+  }
+
+  uint8_t v[3] = {0};
+  ykhsmauth_rc res = ykhsmauth_get_version_ex(state, &v[0], &v[1], &v[2]);
+  if(res == YKHSMAUTHR_SUCCESS) {
+    int result = snprintf(version, len, "%d.%d.%d", v[0], v[1], v[2]);
     if (result < 0) {
       if (state->verbose) {
         fprintf(stderr, "Version buffer too small\n");
       }
       return YKHSMAUTHR_GENERIC_ERROR;
     }
-    return YKHSMAUTHR_SUCCESS;
-  } else {
-    return translate_error(sw, NULL);
   }
+  return res;
 }
 
 ykhsmauth_rc ykhsmauth_put(ykhsmauth_state *state, const uint8_t *mgmkey,
@@ -603,10 +618,12 @@ ykhsmauth_rc ykhsmauth_get_challenge_ex(ykhsmauth_state *state, const char *labe
     return YKHSMAUTHR_INVALID_PARAMS;
   }
 
+  ykhsmauth_rc rc;
+
   APDU apdu = {{0, YKHSMAUTH_INS_GET_CHALLENGE, 0, 0, 0, {0}}};
 
   add_tag(&apdu, YKHSMAUTH_TAG_LABEL, label, strlen(label), 0);
-  if(cpw && cpw_len) {
+  if (cpw && cpw_len) {
     add_tag(&apdu, YKHSMAUTH_TAG_PW, cpw, cpw_len, YKHSMAUTH_PW_LEN - cpw_len);
   }
 
@@ -614,7 +631,7 @@ ykhsmauth_rc ykhsmauth_get_challenge_ex(ykhsmauth_state *state, const char *labe
   DWORD recv_len = sizeof(data);
   uint16_t sw = 0;
 
-  ykhsmauth_rc rc = send_data(state, &apdu, data, &recv_len, &sw);
+  rc = send_data(state, &apdu, data, &recv_len, &sw);
   if (rc != YKHSMAUTHR_SUCCESS) {
     return rc;
   } else if (sw != SW_SUCCESS) {
diff --git a/ykhsmauth/ykhsmauth.h b/ykhsmauth/ykhsmauth.h
index 98246eee..39d2322f 100644
--- a/ykhsmauth/ykhsmauth.h
+++ b/ykhsmauth/ykhsmauth.h
@@ -143,6 +143,9 @@ ykhsmauth_rc ykhsmauth_list_readers(ykhsmauth_state *state, char *readers,
                                     size_t *len);
 ykhsmauth_rc ykhsmauth_disconnect(ykhsmauth_state *state);
 
+ykhsmauth_rc ykhsmauth_get_version_ex(ykhsmauth_state *state, uint8_t *major,
+                                      uint8_t *minor, uint8_t *patch);
+
 ykhsmauth_rc ykhsmauth_get_version(ykhsmauth_state *state, char *version,
                                    size_t len);