From 640b060792c130be5e48bc167a7552b69a76883c Mon Sep 17 00:00:00 2001
From: Sean Parkinson <sean@wolfssl.com>
Date: Tue, 20 May 2025 09:25:21 +1000
Subject: [PATCH] LMS: Key ID fixup

Fix implementation for extracting from private key data.
Add implementation that gets Key ID from wc_LmsKey.
---
 wolfcrypt/src/ext_lms.c |  9 +++++++++
 wolfcrypt/src/wc_lms.c  | 32 ++++++++++++++++++++++++++++++--
 wolfcrypt/test/test.c   | 34 ++++++++++++++++++++++++++++++++++
 wolfssl/wolfcrypt/lms.h |  5 +++++
 4 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/wolfcrypt/src/ext_lms.c b/wolfcrypt/src/ext_lms.c
index 00a3e55d3..2435db454 100644
--- a/wolfcrypt/src/ext_lms.c
+++ b/wolfcrypt/src/ext_lms.c
@@ -1043,6 +1043,15 @@ int wc_LmsKey_Verify(LmsKey * key, const byte * sig, word32 sigSz,
     return 0;
 }
 
+int wc_LmsKey_GetKid(LmsKey * key, const byte ** kid, word32* kidSz)
+{
+    if ((key == NULL) || (kid == NULL) || (kidSz == NULL)) {
+        return BAD_FUNC_ARG;
+    }
+
+    return NOT_COMPILED_IN;
+}
+
 const byte * wc_LmsKey_GetKidFromPrivRaw(const byte * priv, word32 privSz)
 {
     if ((priv == NULL) || (privSz < 16)) {
diff --git a/wolfcrypt/src/wc_lms.c b/wolfcrypt/src/wc_lms.c
index 9de58da27..5b87b7e1f 100644
--- a/wolfcrypt/src/wc_lms.c
+++ b/wolfcrypt/src/wc_lms.c
@@ -1258,6 +1258,34 @@ int wc_LmsKey_Verify(LmsKey* key, const byte* sig, word32 sigSz,
     return ret;
 }
 
+/* Get the Key ID from the LMS key.
+ *
+ * PRIV = Q | PARAMS | SEED | I
+ * where I is the Key ID.
+ *
+ * @param [in]  key    LMS key.
+ * @param [out] kid    Key ID data.
+ * @param [out] kidSz  Size of key ID.
+ * @return  0 on success.
+ * @return  BAD_FUNC_ARG when a key, kid or kidSz is NULL.
+ */
+int wc_LmsKey_GetKid(LmsKey * key, const byte ** kid, word32* kidSz)
+{
+    word32 offset;
+
+    if ((key == NULL) || (kid == NULL) || (kidSz == NULL)) {
+        return BAD_FUNC_ARG;
+    }
+
+    /* SEED length is hash length. */
+    offset = HSS_Q_LEN + HSS_PRIV_KEY_PARAM_SET_LEN + key->params->hash_len;
+    *kid = key->priv_raw + offset;
+    *kidSz = HSS_PRIVATE_KEY_LEN(key->params->hash_len) - offset;
+
+    return 0;
+}
+
+
 /* Get the Key ID from the raw private key data.
  *
  * PRIV = Q | PARAMS | SEED | I
@@ -1270,7 +1298,7 @@ int wc_LmsKey_Verify(LmsKey* key, const byte* sig, word32 sigSz,
  */
 const byte * wc_LmsKey_GetKidFromPrivRaw(const byte * priv, word32 privSz)
 {
-    word32 seedSz = privSz - LMS_Q_LEN + HSS_PRIV_KEY_PARAM_SET_LEN - LMS_I_LEN;
+    word32 seedSz = privSz - HSS_Q_LEN - HSS_PRIV_KEY_PARAM_SET_LEN - LMS_I_LEN;
 
     if (priv == NULL) {
         return NULL;
@@ -1279,7 +1307,7 @@ const byte * wc_LmsKey_GetKidFromPrivRaw(const byte * priv, word32 privSz)
             (seedSz != WC_SHA256_DIGEST_SIZE)) {
         return NULL;
     }
-    return priv - LMS_I_LEN;
+    return priv + privSz - LMS_I_LEN;
 }
 
 #endif /* WOLFSSL_HAVE_LMS && WOLFSSL_WC_LMS */
diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index d7de46b44..a2382718f 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -47891,6 +47891,11 @@ WOLFSSL_TEST_SUBROUTINE wc_test_ret_t lms_test(void)
 #else
     byte          sig[WC_TEST_LMS_SIG_LEN];
 #endif
+#if !defined(HAVE_LIBLMS)
+    const byte *  kid;
+    word32        kidSz;
+#endif
+
     WOLFSSL_ENTER("lms_test");
 
     XMEMSET(priv, 0, sizeof(priv));
@@ -47939,6 +47944,35 @@ WOLFSSL_TEST_SUBROUTINE wc_test_ret_t lms_test(void)
 
     XMEMCPY(old_priv, priv, sizeof(priv));
 
+#if !defined(HAVE_LIBLMS)
+    ret = wc_LmsKey_GetKid(NULL, NULL, NULL);
+    if (ret != WC_NO_ERR_TRACE(BAD_FUNC_ARG))
+        ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+    ret = wc_LmsKey_GetKid(&signingKey, NULL, NULL);
+    if (ret != WC_NO_ERR_TRACE(BAD_FUNC_ARG))
+        ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+    ret = wc_LmsKey_GetKid(NULL, &kid, NULL);
+    if (ret != WC_NO_ERR_TRACE(BAD_FUNC_ARG))
+        ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+    ret = wc_LmsKey_GetKid(NULL, NULL, &kidSz);
+    if (ret != WC_NO_ERR_TRACE(BAD_FUNC_ARG))
+        ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+    ret = wc_LmsKey_GetKid(&signingKey, &kid, NULL);
+    if (ret != WC_NO_ERR_TRACE(BAD_FUNC_ARG))
+        ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+    ret = wc_LmsKey_GetKid(&signingKey, NULL, &kidSz);
+    if (ret != WC_NO_ERR_TRACE(BAD_FUNC_ARG))
+        ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+    ret = wc_LmsKey_GetKid(NULL, &kid, &kidSz);
+    if (ret != WC_NO_ERR_TRACE(BAD_FUNC_ARG))
+        ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+    ret = wc_LmsKey_GetKid(&signingKey, &kid, &kidSz);
+    if (ret != 0) { ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); }
+    if (kidSz != WC_LMS_I_LEN) {
+        ERROR_OUT(WC_TEST_RET_ENC_I(kidSz), out);
+    }
+#endif
+
     ret = wc_LmsKey_ExportPub(&verifyKey, &signingKey);
     if (ret != 0) { ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); }
 
diff --git a/wolfssl/wolfcrypt/lms.h b/wolfssl/wolfcrypt/lms.h
index ee4ccc5e2..40d8641d4 100644
--- a/wolfssl/wolfcrypt/lms.h
+++ b/wolfssl/wolfcrypt/lms.h
@@ -31,6 +31,9 @@
 
 #ifdef WOLFSSL_HAVE_LMS
 
+/* Length of the Key ID. */
+#define WC_LMS_I_LEN    16
+
 typedef struct LmsKey LmsKey;
 
 /* Private key write and read callbacks. */
@@ -187,6 +190,8 @@ WOLFSSL_API int  wc_LmsKey_Verify(LmsKey * key, const byte * sig, word32 sigSz,
 WOLFSSL_API const char * wc_LmsKey_ParmToStr(enum wc_LmsParm lmsParm);
 WOLFSSL_API const char * wc_LmsKey_RcToStr(enum wc_LmsRc lmsRc);
 
+WOLFSSL_API int wc_LmsKey_GetKid(LmsKey * key, const byte ** kid,
+    word32* kidSz);
 WOLFSSL_API const byte * wc_LmsKey_GetKidFromPrivRaw(const byte * priv,
     word32 privSz);
 #ifdef __cplusplus