From 048a03e8bfde14f353ca4661e64b72bfeb990ed1 Mon Sep 17 00:00:00 2001
From: Josh Holtrop <josh@wolfssl.com>
Date: Thu, 26 Mar 2026 15:20:02 -0400
Subject: [PATCH] Add buffer size and callback checks to wc_LmsKey_Sign

Fixes ZD#21439
---
 wolfcrypt/src/wc_lms.c | 21 +++++++++++++++++++++
 wolfcrypt/test/test.c  | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/wolfcrypt/src/wc_lms.c b/wolfcrypt/src/wc_lms.c
index 7035060621c..f6aa7098c0c 100644
--- a/wolfcrypt/src/wc_lms.c
+++ b/wolfcrypt/src/wc_lms.c
@@ -963,6 +963,11 @@ int wc_LmsKey_GetPrivLen(const LmsKey* key, word32* len)
  * @return  0 on success.
  * @return  BAD_FUNC_ARG when key, sig, sigSz or msg is NULL.
  * @return  BAD_FUNC_ARG when msgSz is not greater than 0.
+ * @return  BAD_FUNC_ARG when a write private key is not set.
+ * @return  BAD_FUNC_ARG when a read/write private key context is not set.
+ * @return  BUFFER_E when sigSz is too small.
+ * @return  BAD_STATE_E when wrong state for operation.
+ * @return  IO_FAILED_E when reading or writing private key failed.
  */
 int wc_LmsKey_Sign(LmsKey* key, byte* sig, word32* sigSz, const byte* msg,
     int msgSz)
@@ -987,6 +992,22 @@ int wc_LmsKey_Sign(LmsKey* key, byte* sig, word32* sigSz, const byte* msg,
         WOLFSSL_MSG("error: can't sign, LMS key not in good state");
         ret = BAD_STATE_E;
     }
+    /* Check signature buffer size. */
+    if ((ret == 0) && (*sigSz < key->params->sig_len)) {
+        /* Signature buffer too small. */
+        WOLFSSL_MSG("error: LMS sig buffer too small");
+        ret = BUFFER_E;
+    }
+    /* Check read and write callbacks available. */
+    if ((ret == 0) && (key->write_private_key == NULL)) {
+        WOLFSSL_MSG("error: LmsKey write/read callbacks are not set");
+        ret = BAD_FUNC_ARG;
+    }
+    /* Check read/write callback context available. */
+    if ((ret == 0) && (key->context == NULL)) {
+        WOLFSSL_MSG("error: LmsKey context is not set");
+        ret = BAD_FUNC_ARG;
+    }
 
     if (ret == 0) {
         WC_DECLARE_VAR(state, LmsState, 1, 0);
diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index a3d6dcfcda3..960012ae3b4 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -51593,6 +51593,39 @@ WOLFSSL_TEST_SUBROUTINE wc_test_ret_t lms_test(void)
         ERROR_OUT(WC_TEST_RET_ENC_I(sigSz), out);
     }
 
+    /* Test wc_LmsKey_Sign input validation. */
+    {
+        word32 smallSz = 1;
+        wc_lms_write_private_key_cb saved_write_cb;
+        void*                       saved_ctx;
+
+        /* Undersized sig buffer should return BUFFER_E. */
+        ret = wc_LmsKey_Sign(&signingKey, sig, &smallSz, (byte *) msg, msgSz);
+        if (ret != WC_NO_ERR_TRACE(BUFFER_E)) {
+            ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+        }
+
+        /* NULL write callback should return BAD_FUNC_ARG. */
+        saved_write_cb = signingKey.write_private_key;
+        signingKey.write_private_key = NULL;
+        ret = wc_LmsKey_Sign(&signingKey, sig, &sigSz, (byte *) msg, msgSz);
+        signingKey.write_private_key = saved_write_cb;
+        if (ret != WC_NO_ERR_TRACE(BAD_FUNC_ARG)) {
+            ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+        }
+
+        /* NULL context should return BAD_FUNC_ARG. */
+        saved_ctx = signingKey.context;
+        signingKey.context = NULL;
+        ret = wc_LmsKey_Sign(&signingKey, sig, &sigSz, (byte *) msg, msgSz);
+        signingKey.context = saved_ctx;
+        if (ret != WC_NO_ERR_TRACE(BAD_FUNC_ARG)) {
+            ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out);
+        }
+
+        ret = 0;
+    }
+
     /* 2 ** 5 should be the max number of signatures */
     for (i = 0; i < 32; ++i) {
         /* We should have remaining signstures. */