From 277a2dd3bae4685e10b1a621e31aabd93db1d303 Mon Sep 17 00:00:00 2001
From: jordan <jordan@wolfssl.com>
Date: Fri, 26 Jan 2024 17:35:27 -0600
Subject: [PATCH] Update xmss example.

---
 pq/stateful_hash_sig/Makefile       |  7 +++++++
 pq/stateful_hash_sig/README.md      | 11 +++++++----
 pq/stateful_hash_sig/xmss_example.c | 25 ++++++++++++++++++-------
 3 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/pq/stateful_hash_sig/Makefile b/pq/stateful_hash_sig/Makefile
index 52c34244..13e3b360 100644
--- a/pq/stateful_hash_sig/Makefile
+++ b/pq/stateful_hash_sig/Makefile
@@ -1,4 +1,8 @@
 # PQ Stateful Hash-Based Signature Examples Makefile
+#
+# XMSS_INC, XMSS_LIB, and WOLF_STATIC_LIB are only required if
+# building with --with-libxmss.
+#
 CC       = gcc
 LIB_PATH = /usr/local
 HSS_INC  =
@@ -35,7 +39,10 @@ lms_example: lms_example.c
 	$(CC) -o $@ $< $(CFLAGS) -I$(HSS_INC) $(LIBS) $(WOLF_STATIC_LIB) $(HSS_LIB)
 
 xmss_example: xmss_example.c
+# If building with  --enable-xmss=wolfssl:
 	$(CC) -o $@ $< $(CFLAGS) -I$(XMSS_INC) $(LIBS) $(WOLF_STATIC_LIB) $(XMSS_LIB)
+# If building with  --enable-xmss --with-libxmss=<path>:
+#	$(CC) -o $@ $< $(CFLAGS) $(LIBS) $(WOLF_DYN_LIB)
 
 xmss_example_verifyonly: xmss_example.c
 	$(CC) -o $@ $< $(CFLAGS) -I$(XMSS_INC) -DWOLFSSL_XMSS_VERIFY_ONLY $(LIBS) $(WOLF_STATIC_LIB) $(XMSS_LIB)
diff --git a/pq/stateful_hash_sig/README.md b/pq/stateful_hash_sig/README.md
index 354da9fb..17ce5750 100644
--- a/pq/stateful_hash_sig/README.md
+++ b/pq/stateful_hash_sig/README.md
@@ -8,7 +8,8 @@ This directory contains:
 
 - An example that uses wolfCrypt XMSS/XMSS^MT hooks to sign and verify a message
   with a configurable XMSS/XMSS^MT parameter string. Requires wolfssl with `--enable-xmss=yes`
-  and `--with-libxmss=<path to patched xmss-reference install>`.
+  and `--with-libxmss=<path to patched xmss-reference install>`, or wolfssl
+  with `--enable-xmss=wolfssl`.
 
 # Prerequisites
 
@@ -18,8 +19,9 @@ in the wolfSSL repo's INSTALL file.
 
 https://github.com/wolfSSL/wolfssl/blob/master/INSTALL
 
-The XMSS/XMSS^MT example requires that the xmss-reference repository has been
-cloned, patched, and built. Please see item 20 in the wolfSSL repo's INSTALL file.
+If building with `--with-libxmss=<path>`, the XMSS/XMSS^MT example requires
+that the xmss-reference repository has been cloned, patched, and built. Please
+see item 20 in the wolfSSL repo's INSTALL file.
 
 The patch to use is `0001-Patch-to-support-wolfSSL-xmss-reference-integration.patch` from this XMSS/XMSS^MT example.
 This patch includes an addendum readme, `patch_readme.md`, that lists all changes made and explains their rationale.
@@ -79,7 +81,8 @@ description:
 
 # Building the XMSS/XMSS^MT example
 
-Configure the Makefile to point to your xmss install:
+If building with `--with-libxmss=<path>`, configure the Makefile to point to
+your xmss install:
 
 ```
 XMSS_INC = <path to patched xmss install>
diff --git a/pq/stateful_hash_sig/xmss_example.c b/pq/stateful_hash_sig/xmss_example.c
index 709f6f42..223d0f92 100644
--- a/pq/stateful_hash_sig/xmss_example.c
+++ b/pq/stateful_hash_sig/xmss_example.c
@@ -24,12 +24,16 @@
 #include <wolfssl/options.h>
 #include <wolfssl/wolfcrypt/error-crypt.h>
 
-#ifdef HAVE_LIBXMSS
+#ifdef WOLFSSL_HAVE_XMSS
 
 #include <wolfssl/wolfcrypt/xmss.h>
-#include <wolfssl/wolfcrypt/ext_xmss.h>
+#ifdef HAVE_LIBXMSS
+    #include <wolfssl/wolfcrypt/ext_xmss.h>
+#else
+    #include <wolfssl/wolfcrypt/wc_xmss.h>
+#endif
 
-static void dump_hex(const char * what, const uint8_t * buf, size_t len);
+static void dump_hex(const char * what, const byte * buf, size_t len);
 static void print_usage(void);
 #if !defined WOLFSSL_XMSS_VERIFY_ONLY
 static int  do_xmss_example(const char * params, size_t sigs_to_do);
@@ -304,6 +308,13 @@ do_xmss_example(const char * params,
 
     printf("signing and verifying %zu signatures...\n", sigs_to_do);
     for (size_t i = 0; i < sigs_to_do; ++i) {
+        ret = wc_XmssKey_SigsLeft(&signingKey);
+
+        if (ret <= 0) {
+            fprintf(stderr, "info: %zu: wc_XmssKey_SigsLeft returned %d\n", i, ret);
+            break;
+        }
+
         ret = wc_XmssKey_Sign(&signingKey, sig, &sigSz,(byte *) msg,
                              strlen(msg));
         if (ret) {
@@ -524,9 +535,9 @@ read_file(byte *       data,
 #endif /* if !defined WOLFSSL_XMSS_VERIFY_ONLY */
 
 static void
-dump_hex(const char *    what,
-         const uint8_t * buf,
-         size_t          len)
+dump_hex(const char * what,
+         const byte * buf,
+         size_t       len)
 {
     printf("%s\n", what);
     for (size_t i = 0; i < len; ++i) {
@@ -547,7 +558,7 @@ dump_hex(const char *    what,
 #else
 
 int main(int argc, char** argv) {
-    printf("This requires the --with-libxmss flag.\n");
+    printf("This requires --enable-xmss.\n");
     return 0;
 }
 #endif /* WITH_LIBXMSS */