Repository: randombit/botan
Pull request #3263 (2023-02-08 17:00:10)

Refactor: Make xmss_wots.h and xmss_hash.h internal
This moves some implementation details of `XMSS_PrivateKey` out of the public interface. With that `xmss_wots.h` and `xmss_hash.h` can be internal headers (and not marked as public API classes anymore).

I'm hoping to reduce the public cross section of the XMSS implementation further in later PRs. Particularly, I don't really see a reason to expose implementation details with public getters **and setters** 😲.

To me, parts of the implementation (particularly the WOTS public/private key handling) seem more complicated than they need to be. Here is a first step towards cleaning this up.

Later, we're planning to adapt the XMSS key generation to the recommendations in [NIST SP.800-208](https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-208.pdf) -- which should be backward compatible even for existing private keys generated by Botan 2.x and add the parameter sets described in the same document.

Eventually, an implementation for XMSS-MT is planned, but I wouldn't want to commit that for Botan 3.0.


Patch:
diff --git a/src/lib/pubkey/xmss/info.txt b/src/lib/pubkey/xmss/info.txt
index 7786999c5..8dbde490e 100644
--- a/src/lib/pubkey/xmss/info.txt
+++ b/src/lib/pubkey/xmss/info.txt
@@ -8,8 +8,6 @@ name -> "XMSS"
 
 <header:public>
 xmss.h
-xmss_hash.h
-xmss_wots.h
 xmss_parameters.h
 </header:public>
 
@@ -17,11 +15,13 @@ xmss_parameters.h
 atomic.h
 xmss_address.h
 xmss_common_ops.h
+xmss_hash.h
 xmss_index_registry.h
 xmss_signature.h
 xmss_signature_operation.h
 xmss_tools.h
 xmss_verification_operation.h
+xmss_wots.h
 </header:internal>
 
 <requires>
diff --git a/src/lib/pubkey/xmss/xmss.h b/src/lib/pubkey/xmss/xmss.h
index c83a54e5f..148459509 100644
--- a/src/lib/pubkey/xmss/xmss.h
+++ b/src/lib/pubkey/xmss/xmss.h
@@ -9,15 +9,21 @@
 #ifndef BOTAN_XMSS_H_
 #define BOTAN_XMSS_H_
 
+#include <span>
+#include <memory>
+
 #include <botan/pk_keys.h>
 #include <botan/exceptn.h>
 #include <botan/xmss_parameters.h>
-#include <botan/xmss_wots.h>
 
 namespace Botan {
 
 class RandomNumberGenerator;
+class XMSS_Address;
+class XMSS_Hash;
+class XMSS_PrivateKey_Internal;
 class XMSS_Verification_Operation;
+class XMSS_WOTS_PrivateKey;
 
 /**
  * An XMSS: Extended Hash-Based Signature public key.
@@ -50,7 +56,7 @@ class BOTAN_PUBLIC_API(2,0) XMSS_PublicKey : public virtual Public_Key
        *
        * @param key_bits DER encoded public key bits
        */
-      XMSS_PublicKey(const std::vector<uint8_t>& key_bits);
+      XMSS_PublicKey(std::span<const uint8_t> key_bits);
 
       /**
        * Creates a new XMSS public key for a chosen XMSS signature method as
@@ -204,7 +210,7 @@ class BOTAN_PUBLIC_API(2,0) XMSS_PublicKey : public virtual Public_Key
        **/
       virtual size_t size() const
          {
-         return sizeof(uint32_t) + 2 * m_xmss_params.element_size();
+         return m_xmss_params.raw_public_key_size();
          }
 
       /**
@@ -264,7 +270,7 @@ class BOTAN_PUBLIC_API(2,0) XMSS_PrivateKey final : public virtual XMSS_PublicKe
        *
        * @param raw_key An XMSS private key serialized using raw_private_key().
        **/
-      XMSS_PrivateKey(const secure_vector<uint8_t>& raw_key);
+      XMSS_PrivateKey(std::span<const uint8_t> raw_key);
 
       /**
        * Creates a new XMSS private key for the chosen XMSS signature method
@@ -318,10 +324,7 @@ class BOTAN_PUBLIC_API(2,0) XMSS_PrivateKey final : public virtual XMSS_PublicKe
        *
        * @return WOTS+ private key.
        **/
-      const XMSS_WOTS_PrivateKey& wots_private_key() const
-         {
-         return m_wots_priv_key;
-         }
+      const XMSS_WOTS_PrivateKey& wots_private_key() const;
 
       /**
        * Winternitz One Time Signature Scheme key utilized for signing
@@ -329,31 +332,13 @@ class BOTAN_PUBLIC_API(2,0) XMSS_PrivateKey final : public virtual XMSS_PublicKe
        *
        * @return WOTS+ private key.
        **/
-      XMSS_WOTS_PrivateKey& wots_private_key()
-         {
-         return m_wots_priv_key;
-         }
+      XMSS_WOTS_PrivateKey& wots_private_key();
 
-      const secure_vector<uint8_t>& prf() const
-         {
-         return m_prf;
-         }
-
-      secure_vector<uint8_t>& prf()
-         {
-         return m_prf;
-         }
-
-      void set_public_seed(secure_vector<uint8_t> public_seed) override
-         {
-         m_public_seed = std::move(public_seed);
-         m_wots_priv_key.set_public_seed(m_public_seed);
-         }
+      const secure_vector<uint8_t>& prf() const;
+      secure_vector<uint8_t>& prf();
 
-      const secure_vector<uint8_t>& public_seed() const override
-         {
-         return m_public_seed;
-         }
+      void set_public_seed(secure_vector<uint8_t> public_seed) override;
+      const secure_vector<uint8_t>& public_seed() const override;
 
       std::unique_ptr<PK_Ops::Signature>
       create_signature_op(RandomNumberGenerator&,
@@ -364,9 +349,7 @@ class BOTAN_PUBLIC_API(2,0) XMSS_PrivateKey final : public virtual XMSS_PublicKe
 
       size_t size() const override
          {
-         return XMSS_PublicKey::size() +
-                sizeof(uint32_t) +
-                2 * XMSS_PublicKey::m_xmss_params.element_size();
+         return XMSS_PublicKey::m_xmss_params.raw_private_key_size();
          }
 
       /**
@@ -378,6 +361,7 @@ class BOTAN_PUBLIC_API(2,0) XMSS_PrivateKey final : public virtual XMSS_PublicKe
        *         8-byte unused leaf index, n-byte prf seed, n-byte private seed.
        **/
       secure_vector<uint8_t> raw_private_key() const;
+
       /**
        * Algorithm 9: "treeHash"
        * Computes the internal n-byte nodes of a Merkle tree.
@@ -396,18 +380,10 @@ class BOTAN_PUBLIC_API(2,0) XMSS_PrivateKey final : public virtual XMSS_PublicKe
          XMSS_Address& adrs);
 
    private:
-      /**
-       * Fetches shared unused leaf index from the index registry
-       **/
-      std::shared_ptr<Atomic<size_t>> recover_global_leaf_index() const;
-
-      inline void tree_hash_subtree(secure_vector<uint8_t>& result,
-                                    size_t start_idx,
-                                    size_t target_node_height,
-                                    XMSS_Address& adrs)
-         {
-         return tree_hash_subtree(result, start_idx, target_node_height, adrs, m_hash);
-         }
+      void tree_hash_subtree(secure_vector<uint8_t>& result,
+                             size_t start_idx,
+                             size_t target_node_height,
+                             XMSS_Address& adrs);
 
 
       /**
@@ -419,10 +395,7 @@ class BOTAN_PUBLIC_API(2,0) XMSS_PrivateKey final : public virtual XMSS_PublicKe
                              XMSS_Address& adrs,
                              XMSS_Hash& hash);
 
-      XMSS_WOTS_PrivateKey m_wots_priv_key;
-      XMSS_Hash m_hash;
-      secure_vector<uint8_t> m_prf;
-      XMSS_Index_Registry& m_index_reg;
+      std::shared_ptr<XMSS_PrivateKey_Internal> m_private;
    };
 
 }
diff --git a/src/lib/pubkey/xmss/xmss_common_ops.cpp b/src/lib/pubkey/xmss/xmss_common_ops.cpp
index bdc03e40f..0a7399f56 100644
--- a/src/lib/pubkey/xmss/xmss_common_ops.cpp
+++ b/src/lib/pubkey/xmss/xmss_common_ops.cpp
@@ -7,7 +7,7 @@
  **/
 
 #include <botan/internal/xmss_common_ops.h>
-#include <botan/xmss_hash.h>
+#include <botan/internal/xmss_hash.h>
 
 namespace Botan {
 
diff --git a/src/lib/pubkey/xmss/xmss_hash.cpp b/src/lib/pubkey/xmss/xmss_hash.cpp
index 8e4de06fd..3b8c779a7 100644
--- a/src/lib/pubkey/xmss/xmss_hash.cpp
+++ b/src/lib/pubkey/xmss/xmss_hash.cpp
@@ -7,7 +7,7 @@
  * Botan is released under the Simplified BSD License (see license.txt)
  **/
 
-#include <botan/xmss_hash.h>
+#include <botan/internal/xmss_hash.h>
 #include <botan/exceptn.h>
 
 namespace Botan {
diff --git a/src/lib/pubkey/xmss/xmss_hash.h b/src/lib/pubkey/xmss/xmss_hash.h
index db8b8d2cb..f45432d59 100644
--- a/src/lib/pubkey/xmss/xmss_hash.h
+++ b/src/lib/pubkey/xmss/xmss_hash.h
@@ -10,15 +10,13 @@
 
 #include <botan/hash.h>
 
-//BOTAN_FUTURE_INTERNAL_HEADER(xmss_hash.h)
-
 namespace Botan {
 
 /**
  * A collection of pseudorandom hash functions required for XMSS and WOTS
  * computations.
  **/
-class BOTAN_PUBLIC_API(2,0) XMSS_Hash final
+class XMSS_Hash final
    {
    public:
       XMSS_Hash(const std::string& h_func_name);
diff --git a/src/lib/pubkey/xmss/xmss_parameters.h b/src/lib/pubkey/xmss/xmss_parameters.h
index 2f186ac53..69531fb15 100644
--- a/src/lib/pubkey/xmss/xmss_parameters.h
+++ b/src/lib/pubkey/xmss/xmss_parameters.h
@@ -8,11 +8,116 @@
 #ifndef BOTAN_XMSS_PARAMETERS_H_
 #define BOTAN_XMSS_PARAMETERS_H_
 
-#include <botan/xmss_wots.h>
 #include <string>
+#include <map>
+
+#include <botan/build.h>
+#include <botan/secmem.h>
 
 namespace Botan {
 
+/**
+ * Descibes a signature method for XMSS Winternitz One Time Signatures,
+ * as defined in:
+ * [1] XMSS: Extended Hash-Based Signatures,
+ *     Request for Comments: 8391
+ *     Release: May 2018.
+ *     https://datatracker.ietf.org/doc/rfc8391/
+ **/
+class BOTAN_PUBLIC_API(2,0) XMSS_WOTS_Parameters final
+   {
+   public:
+      enum ots_algorithm_t
+         {
+         WOTSP_SHA2_256 = 0x00000001,
+         WOTSP_SHA2_512 = 0x00000002,
+         WOTSP_SHAKE_256 = 0x00000003,
+         WOTSP_SHAKE_512 = 0x00000004
+         };
+
+      XMSS_WOTS_Parameters(const std::string& algo_name);
+      XMSS_WOTS_Parameters(ots_algorithm_t ots_spec);
+
+      static ots_algorithm_t xmss_wots_id_from_string(const std::string& param_set);
+
+      /**
+       * Algorithm 1: convert input string to base.
+       *
+       * @param msg Input string (referred to as X in [1]).
+       * @param out_size size of message in base w.
+       *
+       * @return Input string converted to the given base.
+       **/
+      secure_vector<uint8_t> base_w(const secure_vector<uint8_t>& msg, size_t out_size) const;
+
+      secure_vector<uint8_t> base_w(size_t value) const;
+
+      void append_checksum(secure_vector<uint8_t>& data) const;
+
+      /**
+       * @return XMSS WOTS registry name for the chosen parameter set.
+       **/
+      const std::string& name() const
+         {
+         return m_name;
+         }
+
+      /**
+       * @return Botan name for the hash function used.
+       **/
+      const std::string& hash_function_name() const
+         {
+         return m_hash_name;
+         }
+
+      /**
+       * Retrieves the uniform length of a message, and the size of
+       * each node. This correlates to XMSS parameter "n" defined
+       * in [1].
+       *
+       * @return element length in bytes.
+       **/
+      size_t element_size() const { return m_element_size; }
+
+      /**
+       * The Winternitz parameter.
+       *
+       * @return numeric base used for internal representation of
+       *         data.
+       **/
+      size_t wots_parameter() const { return m_w; }
+
+      size_t len() const { return m_len; }
+
+      size_t len_1() const { return m_len_1; }
+
+      size_t len_2() const { return m_len_2; }
+
+      size_t lg_w() const { return m_lg_w; }
+
+      ots_algorithm_t oid() const { return m_oid; }
+
+      size_t estimated_strength() const { return m_strength; }
+
+      bool operator==(const XMSS_WOTS_Parameters& p) const
+         {
+         return m_oid == p.m_oid;
+         }
+
+   private:
+      static const std::map<std::string, ots_algorithm_t> m_oid_name_lut;
+      ots_algorithm_t m_oid;
+      std::string m_name;
+      std::string m_hash_name;
+      size_t m_element_size;
+      size_t m_w;
+      size_t m_len_1;
+      size_t m_len_2;
+      size_t m_len;
+      size_t m_strength;
+      uint8_t m_lg_w;
+   };
+
 /**
  * Descibes a signature method for XMSS, as defined in:
  * [1] XMSS: Extended Hash-Based Signatures,
@@ -97,6 +202,18 @@ class BOTAN_PUBLIC_API(2,0) XMSS_Parameters
          return m_strength;
          }
 
+      size_t raw_public_key_size() const
+         {
+         return sizeof(uint32_t) + 2 * element_size();
+         }
+
+      size_t raw_private_key_size() const
+         {
+         return raw_public_key_size() +
+                sizeof(uint32_t) +
+                2 * element_size();
+         }
+
       bool operator==(const XMSS_Parameters& p) const
          {
          return m_oid == p.m_oid;
diff --git a/src/lib/pubkey/xmss/xmss_privatekey.cpp b/src/lib/pubkey/xmss/xmss_privatekey.cpp
index 1355718d4..f3648d026 100644
--- a/src/lib/pubkey/xmss/xmss_privatekey.cpp
+++ b/src/lib/pubkey/xmss/xmss_privatekey.cpp
@@ -19,6 +19,8 @@
 #include <botan/internal/xmss_signature_operation.h>
 #include <botan/internal/xmss_index_registry.h>
 #include <botan/internal/xmss_common_ops.h>
+#include <botan/internal/loadstor.h>
+#include <botan/internal/stl_util.h>
 #include <botan/ber_dec.h>
 #include <botan/der_enc.h>
 #include <iterator>
@@ -32,82 +34,182 @@ namespace Botan {
 namespace {
 
 // fall back to raw decoding for previous versions, which did not encode an OCTET STRING
-secure_vector<uint8_t> extract_raw_key(const secure_vector<uint8_t>& key_bits)
+secure_vector<uint8_t> extract_raw_private_key(std::span<const uint8_t> key_bits)
    {
    secure_vector<uint8_t> raw_key;
    try
       {
-      BER_Decoder(key_bits).decode(raw_key, ASN1_Type::OctetString);
+      DataSource_Memory src(key_bits);
+      BER_Decoder(src).decode(raw_key, ASN1_Type::OctetString);
       }
    catch(Decoding_Error&)
       {
-      raw_key = key_bits;
+      raw_key.assign(key_bits.begin(), key_bits.end());
       }
    return raw_key;
    }
 
 }
 
-XMSS_PrivateKey::XMSS_PrivateKey(const secure_vector<uint8_t>& key_bits)
-   : XMSS_PublicKey(unlock(key_bits)),
-     m_wots_priv_key(m_wots_params.oid(), m_public_seed),
-     m_hash(xmss_hash_function()),
-     m_index_reg(XMSS_Index_Registry::get_instance())
+class XMSS_PrivateKey_Internal
    {
-   /*
-   The code requires sizeof(size_t) >= ceil(tree_height / 8)
-
-   Maximum supported tree height is 20, ceil(20/8) == 3, so 4 byte
-   size_t is sufficient for all defined parameters, or even a
-   (hypothetical) tree height 32, which would be extremely slow to
-   compute.
-   */
-   static_assert(sizeof(size_t) >= 4, "size_t is big enough to support leaf index");
+   public:
+      XMSS_PrivateKey_Internal(const XMSS_Parameters& xmss_params,
+                               const XMSS_WOTS_Parameters& wots_params,
+                               secure_vector<uint8_t> public_seed,
+                               secure_vector<uint8_t> private_seed,
+                               secure_vector<uint8_t> prf)
+         : m_xmss_params(xmss_params)
+         , m_wots_params(wots_params)
+         , m_wots_priv_key(wots_params.oid(), std::move(public_seed), std::move(private_seed))
+         , m_hash(m_xmss_params.hash_function_name())
+         , m_prf(std::move(prf))
+         , m_index_reg(XMSS_Index_Registry::get_instance()) {}
+
+      XMSS_PrivateKey_Internal(const XMSS_Parameters& xmss_params,
+                               const XMSS_WOTS_Parameters& wots_params,
+                               secure_vector<uint8_t> public_seed,
+                               std::span<const uint8_t> key_bits)
+         : m_xmss_params(xmss_params)
+         , m_wots_params(wots_params)
+         , m_wots_priv_key(wots_params.oid(), std::move(public_seed))
+         , m_hash(m_xmss_params.hash_function_name())
+         , m_index_reg(XMSS_Index_Registry::get_instance())
+         {
+         /*
+         The code requires sizeof(size_t) >= ceil(tree_height / 8)
+
+         Maximum supported tree height is 20, ceil(20/8) == 3, so 4 byte
+         size_t is sufficient for all defined parameters, or even a
+         (hypothetical) tree height 32, which would be extremely slow to
+         compute.
+         */
+         static_assert(sizeof(size_t) >= 4, "size_t is big enough to support leaf index");
+
+         const secure_vector<uint8_t> raw_key = extract_raw_private_key(key_bits);
+
+         if(raw_key.size() != m_xmss_params.raw_private_key_size())
+            {
+            throw Decoding_Error("Invalid XMSS private key size");
+            }
+
+         // extract & copy unused leaf index from raw_key
+         uint64_t unused_leaf = 0;
+         auto begin = (raw_key.begin() + m_xmss_params.raw_public_key_size());
+         auto end = raw_key.begin() + m_xmss_params.raw_public_key_size() + sizeof(uint32_t);
+
+         for(auto& i = begin; i != end; i++)
+            {
+            unused_leaf = ((unused_leaf << 8) | *i);
+            }
+
+         if(unused_leaf >= (1ull << m_xmss_params.tree_height()))
+            {
+            throw Decoding_Error("XMSS private key leaf index out of bounds");
+            }
+
+         begin = end;
+         end = begin + m_xmss_params.element_size();
+         m_prf.clear();
+         m_prf.reserve(m_xmss_params.element_size());
+         std::copy(begin, end, std::back_inserter(m_prf));
+
+         begin = end;
+         end = begin + m_wots_params.element_size();
+         m_wots_priv_key.set_private_seed(secure_vector<uint8_t>(begin, end));
+         set_unused_leaf_index(static_cast<size_t>(unused_leaf));
+         }
 
-   secure_vector<uint8_t> raw_key = extract_raw_key(key_bits);
+      secure_vector<uint8_t> serialize(std::vector<uint8_t> raw_public_key) const {
+         std::vector<uint8_t> unused_index(4);
+         store_be(static_cast<uint32_t>(unused_leaf_index()), unused_index.data());
 
-   if(raw_key.size() != XMSS_PrivateKey::size())
-      {
-      throw Decoding_Error("Invalid XMSS private key size");
+         return concat_as<secure_vector<uint8_t>>(
+            raw_public_key,
+            unused_index,
+            m_prf,
+            m_wots_priv_key.private_seed());
       }
 
-   // extract & copy unused leaf index from raw_key
-   uint64_t unused_leaf = 0;
-   auto begin = (raw_key.begin() + XMSS_PublicKey::size());
-   auto end = raw_key.begin() + XMSS_PublicKey::size() + sizeof(uint32_t);
+      XMSS_Hash& hash() { return m_hash; }
+      const secure_vector<uint8_t>& prf_value() const { return m_prf; }
+      secure_vector<uint8_t>& prf_value() { return m_prf; }
+      const XMSS_WOTS_PrivateKey& wots_private_key() const { return m_wots_priv_key; }
+      XMSS_WOTS_PrivateKey& wots_private_key() { return m_wots_priv_key; }
+      XMSS_Index_Registry& index_registry() { return m_index_reg; }
 
-   for(auto& i = begin; i != end; i++)
-      {
-      unused_leaf = ((unused_leaf << 8) | *i);
-      }
+      std::shared_ptr<Atomic<size_t>>
+      recover_global_leaf_index() const
+         {
+         BOTAN_ASSERT(m_wots_priv_key.private_seed().size() ==
+                     m_xmss_params.element_size() &&
+                     m_prf.size() == m_xmss_params.element_size(),
+                     "Trying to retrieve index for partially initialized key");
+         return m_index_reg.get(m_wots_priv_key.private_seed(), m_prf);
+         }
 
-   if(unused_leaf >= (1ull << XMSS_PublicKey::m_xmss_params.tree_height()))
-      {
-      throw Decoding_Error("XMSS private key leaf index out of bounds");
-      }
+      void set_unused_leaf_index(size_t idx)
+         {
+         if(idx >= (1ull << m_xmss_params.tree_height()))
+            {
+            throw Decoding_Error("XMSS private key leaf index out of bounds");
+            }
+         else
+            {
+            std::atomic<size_t>& index =
+               static_cast<std::atomic<size_t>&>(*recover_global_leaf_index());
+            size_t current = 0;
+
+            do
+               {
+               current = index.load();
+               if(current > idx)
+                  { return; }
+               }
+            while(!index.compare_exchange_strong(current, idx));
+            }
+         }
+
+      size_t reserve_unused_leaf_index()
+         {
+         size_t idx = (static_cast<std::atomic<size_t>&>(
+                        *recover_global_leaf_index())).fetch_add(1);
+         if(idx >= (1ull << m_xmss_params.tree_height()))
+            {
+            throw Decoding_Error("XMSS private key, one time signatures exhaused");
+            }
+         return idx;
+         }
 
-   begin = end;
-   end = begin + XMSS_PublicKey::m_xmss_params.element_size();
-   m_prf.clear();
-   m_prf.reserve(XMSS_PublicKey::m_xmss_params.element_size());
-   std::copy(begin, end, std::back_inserter(m_prf));
+      size_t unused_leaf_index() const
+         {
+         return *recover_global_leaf_index();
+         }
 
-   begin = end;
-   end = begin + m_wots_params.element_size();
-   m_wots_priv_key.set_private_seed(secure_vector<uint8_t>(begin, end));
-   set_unused_leaf_index(static_cast<size_t>(unused_leaf));
-   }
+   private:
+      const XMSS_Parameters& m_xmss_params;
+      const XMSS_WOTS_Parameters& m_wots_params;
+
+      XMSS_WOTS_PrivateKey m_wots_priv_key;
+      XMSS_Hash m_hash;
+      secure_vector<uint8_t> m_prf;
+      XMSS_Index_Registry& m_index_reg;
+   };
+
+
+XMSS_PrivateKey::XMSS_PrivateKey(std::span<const uint8_t> key_bits)
+   : XMSS_PublicKey(key_bits)
+   , m_private(std::make_shared<XMSS_PrivateKey_Internal>(
+         m_xmss_params, m_wots_params, m_public_seed, key_bits)) {}
 
 XMSS_PrivateKey::XMSS_PrivateKey(
    XMSS_Parameters::xmss_algorithm_t xmss_algo_id,
    RandomNumberGenerator& rng)
-   : XMSS_PublicKey(xmss_algo_id, rng),
-     m_wots_priv_key(XMSS_PublicKey::m_xmss_params.ots_oid(),
-                     public_seed(),
-                     rng),
-     m_hash(xmss_hash_function()),
-     m_prf(rng.random_vec(XMSS_PublicKey::m_xmss_params.element_size())),
-     m_index_reg(XMSS_Index_Registry::get_instance())
+   : XMSS_PublicKey(xmss_algo_id, rng)
+   , m_private(std::make_shared<XMSS_PrivateKey_Internal>(
+      m_xmss_params, m_wots_params, m_public_seed,
+      rng.random_vec(m_xmss_params.element_size()),
+      rng.random_vec(m_xmss_params.element_size())))
    {
    XMSS_Address adrs;
    set_root(tree_hash(0,
@@ -115,22 +217,18 @@ XMSS_PrivateKey::XMSS_PrivateKey(
                       adrs));
    }
 
-
 XMSS_PrivateKey::XMSS_PrivateKey(XMSS_Parameters::xmss_algorithm_t xmss_algo_id,
                                  size_t idx_leaf,
                                  secure_vector<uint8_t> wots_priv_seed,
                                  secure_vector<uint8_t> prf,
                                  secure_vector<uint8_t> root,
                                  secure_vector<uint8_t> public_seed)
-   : XMSS_PublicKey(xmss_algo_id, std::move(root), public_seed /* copied on purpose */),
-     m_wots_priv_key(XMSS_PublicKey::m_xmss_params.ots_oid(),
-                     std::move(public_seed),
-                     std::move(wots_priv_seed)),
-     m_hash(XMSS_PublicKey::m_xmss_params.hash_function_name()),
-     m_prf(std::move(prf)),
-     m_index_reg(XMSS_Index_Registry::get_instance())
+   : XMSS_PublicKey(xmss_algo_id, std::move(root), std::move(public_seed))
+   , m_private(std::make_shared<XMSS_PrivateKey_Internal>(
+      m_xmss_params, m_wots_params, m_public_seed,
+      std::move(wots_priv_seed), std::move(prf)))
    {
-   set_unused_leaf_index(idx_leaf);
+   m_private->set_unused_leaf_index(idx_leaf);
    }
 
 secure_vector<uint8_t>
@@ -171,7 +269,7 @@ XMSS_PrivateKey::tree_hash(size_t start_idx,
        subtrees,
        secure_vector<uint8_t>(XMSS_PublicKey::m_xmss_params.element_size()));
    std::vector<XMSS_Address> node_addresses(subtrees, adrs);
-   std::vector<XMSS_Hash> xmss_hash(subtrees, m_hash);
+   std::vector<XMSS_Hash> xmss_hash(subtrees, m_private->hash());
    std::vector<std::future<void>> work;
 
    // Calculate multiple subtrees in parallel.
@@ -243,16 +341,25 @@ XMSS_PrivateKey::tree_hash(size_t start_idx,
                                         nodes[1],
                                         node_addresses[0],
                                         this->public_seed(),
-                                        m_hash,
+                                        m_private->hash(),
                                         m_xmss_params);
    return nodes[0];
 #else
    secure_vector<uint8_t> result;
-   tree_hash_subtree(result, start_idx, target_node_height, adrs, m_hash);
+   tree_hash_subtree(result, start_idx, target_node_height, adrs, m_private->hash());
    return result;
 #endif
    }
 
+
+void XMSS_PrivateKey::tree_hash_subtree(secure_vector<uint8_t>& result,
+                                        size_t start_idx,
+                                        size_t target_node_height,
+                                        XMSS_Address& adrs)
+   {
+   return tree_hash_subtree(result, start_idx, target_node_height, adrs, m_private->hash());
+   }
+
 void
 XMSS_PrivateKey::tree_hash_subtree(secure_vector<uint8_t>& result,
                                    size_t start_idx,
@@ -273,7 +380,7 @@ XMSS_PrivateKey::tree_hash_subtree(secure_vector<uint8_t>& result,
    std::vector<uint8_t> node_levels(target_node_height + 1);
 
    uint8_t level = 0; // current level on the node stack.
-   XMSS_WOTS_PublicKey pk(m_wots_priv_key.wots_parameters().oid(), seed);
+   XMSS_WOTS_PublicKey pk(m_private->wots_private_key().wots_parameters().oid(), seed);
    const size_t last_idx = (static_cast<size_t>(1) << target_node_height) + start_idx;
 
    for(size_t i = start_idx; i < last_idx; i++)
@@ -321,73 +428,56 @@ secure_vector<uint8_t> XMSS_PrivateKey::private_key_bits() const
    return DER_Encoder().encode(raw_private_key(), ASN1_Type::OctetString).get_contents();
    }
 
-std::shared_ptr<Atomic<size_t>>
-XMSS_PrivateKey::recover_global_leaf_index() const
+void XMSS_PrivateKey::set_unused_leaf_index(size_t idx)
    {
-   BOTAN_ASSERT(m_wots_priv_key.private_seed().size() ==
-                XMSS_PublicKey::m_xmss_params.element_size() &&
-                m_prf.size() == XMSS_PublicKey::m_xmss_params.element_size(),
-                "Trying to retrieve index for partially initialized key");
-   return m_index_reg.get(m_wots_priv_key.private_seed(), m_prf);
+   return m_private->set_unused_leaf_index(idx);
    }
 
-void XMSS_PrivateKey::set_unused_leaf_index(size_t idx)
+size_t XMSS_PrivateKey::reserve_unused_leaf_index()
    {
-   if(idx >= (1ull << XMSS_PublicKey::m_xmss_params.tree_height()))
-      {
-      throw Decoding_Error("XMSS private key leaf index out of bounds");
-      }
-   else
-      {
-      std::atomic<size_t>& index =
-         static_cast<std::atomic<size_t>&>(*recover_global_leaf_index());
-      size_t current = 0;
+   return m_private->reserve_unused_leaf_index();
+   }
 
-      do
-         {
-         current = index.load();
-         if(current > idx)
-            { return; }
-         }
-      while(!index.compare_exchange_strong(current, idx));
-      }
+size_t XMSS_PrivateKey::unused_leaf_index() const
+   {
+   return m_private->unused_leaf_index();
    }
 
-size_t XMSS_PrivateKey::reserve_unused_leaf_index()
+const XMSS_WOTS_PrivateKey& XMSS_PrivateKey::wots_private_key() const
    {
-   size_t idx = (static_cast<std::atomic<size_t>&>(
-                    *recover_global_leaf_index())).fetch_add(1);
-   if(idx >= (1ull << XMSS_PublicKey::m_xmss_params.tree_height()))
-      {
-      throw Decoding_Error("XMSS private key, one time signatures exhaused");
-      }
-   return idx;
+   return m_private->wots_private_key();
    }
 
-size_t XMSS_PrivateKey::unused_leaf_index() const
+XMSS_WOTS_PrivateKey& XMSS_PrivateKey::wots_private_key()
    {
-   return *recover_global_leaf_index();
+   return m_private->wots_private_key();
    }
 
-secure_vector<uint8_t> XMSS_PrivateKey::raw_private_key() const
+const secure_vector<uint8_t>& XMSS_PrivateKey::prf() const
    {
-   std::vector<uint8_t> pk { raw_public_key() };
-   secure_vector<uint8_t> result(pk.begin(), pk.end());
-   result.reserve(size());
+   return m_private->prf_value();
+   }
 
-   for(int i = 3; i >= 0; i--)
-      {
-      result.push_back(
-         static_cast<uint8_t>(
-            static_cast<uint64_t>(unused_leaf_index()) >> 8 * i));
-      }
+secure_vector<uint8_t>& XMSS_PrivateKey::prf()
+   {
+   // TODO: do we even need this?
+   return m_private->prf_value();
+   }
+
+void XMSS_PrivateKey::set_public_seed(secure_vector<uint8_t> public_seed)
+   {
+   m_public_seed = std::move(public_seed);
+   m_private->wots_private_key().set_public_seed(m_public_seed);
+   }
 
-   std::copy(m_prf.begin(), m_prf.end(), std::back_inserter(result));
-   std::copy(m_wots_priv_key.private_seed().begin(),
-             m_wots_priv_key.private_seed().end(),
-             std::back_inserter(result));
+const secure_vector<uint8_t>& XMSS_PrivateKey::public_seed() const
+   {
+   return m_public_seed;
+   }
 
-   return result;
+secure_vector<uint8_t> XMSS_PrivateKey::raw_private_key() const
+   {
+   return m_private->serialize(raw_public_key());
    }
 
 std::unique_ptr<Public_Key> XMSS_PrivateKey::public_key() const
diff --git a/src/lib/pubkey/xmss/xmss_publickey.cpp b/src/lib/pubkey/xmss/xmss_publickey.cpp
index 7b9ee69a2..85cc09930 100644
--- a/src/lib/pubkey/xmss/xmss_publickey.cpp
+++ b/src/lib/pubkey/xmss/xmss_publickey.cpp
@@ -25,16 +25,17 @@ namespace Botan {
 namespace {
 
 // fall back to raw decoding for previous versions, which did not encode an OCTET STRING
-std::vector<uint8_t> extract_raw_key(const std::vector<uint8_t>& key_bits)
+std::vector<uint8_t> extract_raw_public_key(std::span<const uint8_t> key_bits)
    {
    std::vector<uint8_t> raw_key;
    try
       {
-      BER_Decoder(key_bits).decode(raw_key, ASN1_Type::OctetString);
+      DataSource_Memory src(key_bits);
+      BER_Decoder(src).decode(raw_key, ASN1_Type::OctetString);
       }
    catch(Decoding_Error&)
       {
-      raw_key = key_bits;
+      raw_key.assign(key_bits.begin(), key_bits.end());
       }
    return raw_key;
    }
@@ -48,8 +49,8 @@ XMSS_PublicKey::XMSS_PublicKey(XMSS_Parameters::xmss_algorithm_t xmss_oid,
      m_public_seed(rng.random_vec(m_xmss_params.element_size()))
    {}
 
-XMSS_PublicKey::XMSS_PublicKey(const std::vector<uint8_t>& key_bits)
-   : m_raw_key(extract_raw_key(key_bits)),
+XMSS_PublicKey::XMSS_PublicKey(std::span<const uint8_t> key_bits)
+   : m_raw_key(extract_raw_public_key(key_bits)),
      m_xmss_params(XMSS_PublicKey::deserialize_xmss_oid(m_raw_key)),
      m_wots_params(m_xmss_params.ots_oid())
    {
diff --git a/src/lib/pubkey/xmss/xmss_signature.h b/src/lib/pubkey/xmss/xmss_signature.h
index 58c408464..355254c17 100644
--- a/src/lib/pubkey/xmss/xmss_signature.h
+++ b/src/lib/pubkey/xmss/xmss_signature.h
@@ -13,7 +13,7 @@
 #include <botan/types.h>
 #include <botan/secmem.h>
 #include <botan/xmss_parameters.h>
-#include <botan/xmss_wots.h>
+#include <botan/internal/xmss_wots.h>
 
 namespace Botan {
 
diff --git a/src/lib/pubkey/xmss/xmss_signature_operation.h b/src/lib/pubkey/xmss/xmss_signature_operation.h
index 5236b5bbe..a86f8b352 100644
--- a/src/lib/pubkey/xmss/xmss_signature_operation.h
+++ b/src/lib/pubkey/xmss/xmss_signature_operation.h
@@ -12,7 +12,7 @@
 #include <botan/xmss.h>
 #include <botan/internal/xmss_address.h>
 #include <botan/internal/xmss_signature.h>
-#include <botan/xmss_wots.h>
+#include <botan/internal/xmss_wots.h>
 
 namespace Botan {
 
diff --git a/src/lib/pubkey/xmss/xmss_wots.h b/src/lib/pubkey/xmss/xmss_wots.h
index 5675b84aa..726d715a4 100644
--- a/src/lib/pubkey/xmss/xmss_wots.h
+++ b/src/lib/pubkey/xmss/xmss_wots.h
@@ -13,7 +13,8 @@
 #include <botan/pk_keys.h>
 #include <botan/rng.h>
 #include <botan/secmem.h>
-#include <botan/xmss_hash.h>
+#include <botan/internal/xmss_hash.h>
+#include <botan/xmss_parameters.h>
 #include <map>
 #include <memory>
 #include <string>
@@ -21,108 +22,6 @@
 
 namespace Botan {
 
-/**
- * Descibes a signature method for XMSS Winternitz One Time Signatures,
- * as defined in:
- * [1] XMSS: Extended Hash-Based Signatures,
- *     Request for Comments: 8391
- *     Release: May 2018.
- *     https://datatracker.ietf.org/doc/rfc8391/
- **/
-class BOTAN_PUBLIC_API(2,0) XMSS_WOTS_Parameters final
-   {
-   public:
-      enum ots_algorithm_t
-         {
-         WOTSP_SHA2_256 = 0x00000001,
-         WOTSP_SHA2_512 = 0x00000002,
-         WOTSP_SHAKE_256 = 0x00000003,
-         WOTSP_SHAKE_512 = 0x00000004
-         };
-
-      XMSS_WOTS_Parameters(const std::string& algo_name);
-      XMSS_WOTS_Parameters(ots_algorithm_t ots_spec);
-
-      static ots_algorithm_t xmss_wots_id_from_string(const std::string& param_set);
-
-      /**
-       * Algorithm 1: convert input string to base.
-       *
-       * @param msg Input string (referred to as X in [1]).
-       * @param out_size size of message in base w.
-       *
-       * @return Input string converted to the given base.
-       **/
-      secure_vector<uint8_t> base_w(const secure_vector<uint8_t>& msg, size_t out_size) const;
-
-      secure_vector<uint8_t> base_w(size_t value) const;
-
-      void append_checksum(secure_vector<uint8_t>& data) const;
-
-      /**
-       * @return XMSS WOTS registry name for the chosen parameter set.
-       **/
-      const std::string& name() const
-         {
-         return m_name;
-         }
-
-      /**
-       * @return Botan name for the hash function used.
-       **/
-      const std::string& hash_function_name() const
-         {
-         return m_hash_name;
-         }
-
-      /**
-       * Retrieves the uniform length of a message, and the size of
-       * each node. This correlates to XMSS parameter "n" defined
-       * in [1].
-       *
-       * @return element length in bytes.
-       **/
-      size_t element_size() const { return m_element_size; }
-
-      /**
-       * The Winternitz parameter.
-       *
-       * @return numeric base used for internal representation of
-       *         data.
-       **/
-      size_t wots_parameter() const { return m_w; }
-
-      size_t len() const { return m_len; }
-
-      size_t len_1() const { return m_len_1; }
-
-      size_t len_2() const { return m_len_2; }
-
-      size_t lg_w() const { return m_lg_w; }
-
-      ots_algorithm_t oid() const { return m_oid; }
-
-      size_t estimated_strength() const { return m_strength; }
-
-      bool operator==(const XMSS_WOTS_Parameters& p) const
-         {
-         return m_oid == p.m_oid;
-         }
-
-   private:
-      static const std::map<std::string, ots_algorithm_t> m_oid_name_lut;
-      ots_algorithm_t m_oid;
-      std::string m_name;
-      std::string m_hash_name;
-      size_t m_element_size;
-      size_t m_w;
-      size_t m_len_1;
-      size_t m_len_2;
-      size_t m_len;
-      size_t m_strength;
-      uint8_t m_lg_w;
-   };
-
 class XMSS_Address;
 
 typedef std::vector<secure_vector<uint8_t>> wots_keysig_t;
@@ -131,7 +30,7 @@ typedef std::vector<secure_vector<uint8_t>> wots_keysig_t;
  * A Winternitz One Time Signature public key for use with Extended Hash-Based
  * Signatures.
  **/
-class BOTAN_PUBLIC_API(2,0) XMSS_WOTS_PublicKey : virtual public Public_Key
+class XMSS_WOTS_PublicKey : virtual public Public_Key
    {
    public:
       class TreeSignature final
@@ -418,7 +317,7 @@ class BOTAN_PUBLIC_API(2,0) XMSS_WOTS_PublicKey : virtual public Public_Key
 /** A Winternitz One Time Signature private key for use with Extended Hash-Based
  * Signatures.
  **/
-class BOTAN_PUBLIC_API(2,0) XMSS_WOTS_PrivateKey final : public virtual XMSS_WOTS_PublicKey,
+class XMSS_WOTS_PrivateKey final : public virtual XMSS_WOTS_PublicKey,
    public virtual Private_Key
    {
    public:
diff --git a/src/lib/pubkey/xmss/xmss_wots_parameters.cpp b/src/lib/pubkey/xmss/xmss_wots_parameters.cpp
index 4601b5c4e..e8564e8ea 100644
--- a/src/lib/pubkey/xmss/xmss_wots_parameters.cpp
+++ b/src/lib/pubkey/xmss/xmss_wots_parameters.cpp
@@ -12,7 +12,7 @@
  * Botan is released under the Simplified BSD License (see license.txt)
  **/
 
-#include <botan/xmss_wots.h>
+#include <botan/internal/xmss_wots.h>
 #include <botan/internal/xmss_tools.h>
 #include <botan/exceptn.h>
 #include <cmath>
diff --git a/src/lib/pubkey/xmss/xmss_wots_privatekey.cpp b/src/lib/pubkey/xmss/xmss_wots_privatekey.cpp
index ace70d333..89484c0b4 100644
--- a/src/lib/pubkey/xmss/xmss_wots_privatekey.cpp
+++ b/src/lib/pubkey/xmss/xmss_wots_privatekey.cpp
@@ -8,7 +8,7 @@
  * Botan is released under the Simplified BSD License (see license.txt)
  **/
 
-#include <botan/xmss_wots.h>
+#include <botan/internal/xmss_wots.h>
 #include <botan/internal/xmss_tools.h>
 #include <botan/internal/xmss_address.h>
 
diff --git a/src/lib/pubkey/xmss/xmss_wots_publickey.cpp b/src/lib/pubkey/xmss/xmss_wots_publickey.cpp
index 2c50e1387..d5b30b4da 100644
--- a/src/lib/pubkey/xmss/xmss_wots_publickey.cpp
+++ b/src/lib/pubkey/xmss/xmss_wots_publickey.cpp
@@ -8,7 +8,7 @@
  * Botan is released under the Simplified BSD License (see license.txt)
  **/
 
-#include <botan/xmss_wots.h>
+#include <botan/internal/xmss_wots.h>
 #include <botan/internal/xmss_address.h>
 
 namespace Botan {


Test patch:
diff --git a/src/scripts/test_cli.py b/src/scripts/test_cli.py
index fd54dc017..a3cadab93 100755
--- a/src/scripts/test_cli.py
+++ b/src/scripts/test_cli.py
@@ -405,27 +405,27 @@ def cli_xmss_sign_tests(tmp_dir):
     test_cli("hash", ["--no-fsname", msg], "E3B0C44298FC1C149AFBF4C8996FB92427AE41E4649B934CA495991B7852B855")
 
     test_cli("keygen", ["--algo=XMSS", "--output=%s" % (priv_key)], "")
-    test_cli("hash", ["--no-fsname", priv_key], "5B38F737BA41BE7F40433DB30EAEF7C41ABB0F7D9E7A09DEB5FDCE7B6811693F")
+    test_cli("hash", ["--no-fsname", priv_key], "8E239EC9CB28F58174F0C443D9884F921575B46A870FB277197758D120DA1FA1")
 
     test_cli("pkcs8", "--pub-out --output=%s %s" % (pub_key, priv_key), "")
     test_cli("fingerprint", ['--no-fsname', pub_key],
-             "B0:F4:98:6E:D8:4E:05:63:A1:D8:4B:37:61:5A:A0:41:78:7E:DE:0E:72:46:E0:A8:D6:CF:09:54:08:DA:A4:22")
+             "1E:10:69:84:37:63:4F:66:84:34:E2:4B:2B:A1:7A:3C:92:0F:B3:BA:D1:5F:CA:A8:66:47:89:90:7E:91:D1:6A")
 
     # verify the key is updated after each signature:
     test_cli("sign", [priv_key, msg, "--output=%s" % (sig1)], "")
     test_cli("verify", [pub_key, msg, sig1], "Signature is valid")
-    test_cli("hash", ["--no-fsname", sig1], "04AF45451C7A9AF2D828E1AD6EC262E012436F4087C5DA6F32C689D781E597D0")
-    test_cli("hash", ["--no-fsname", priv_key], "67929FAEC636E43DE828C1CD7E2D11CE7C3388CE90DD0A0F687C6627FFA850CD")
+    test_cli("hash", ["--no-fsname", sig1], "1506DD8AB43C714CD8EA20AA599B77D916A2FB554AD5B19A68DFCF00847F6001")
+    test_cli("hash", ["--no-fsname", priv_key], "B76BCED7BD52D0B314373589E7AA638B331740CD24F686ACEBAD0BA7D5384048")
 
     test_cli("sign", [priv_key, msg, "--output=%s" % (sig2)], "")
     test_cli("verify", [pub_key, msg, sig2], "Signature is valid")
-    test_cli("hash", ["--no-fsname", sig2], "0785A6AD54CC7D01F2BE2BC6463A3EAA1159792E52210ED754992C5068E8F24F")
-    test_cli("hash", ["--no-fsname", priv_key], "1940945D68B1CF54D79E05DD7913A4D0B4959183F1E12B81A4E43EF4E63FBD20")
+    test_cli("hash", ["--no-fsname", sig2], "E3FCB5D2AB12159A99E7619C933A97808077C67CCEC8A2504EFC8720465A68E7")
+    test_cli("hash", ["--no-fsname", priv_key], "E2EFFE9164ED6F0CBF39C30CC3996BC0B596B9C987C0FE01113FB7DC15F1E5ED")
 
     # private key updates, public key is unchanged:
     test_cli("pkcs8", "--pub-out --output=%s %s" % (pub_key2, priv_key), "")
     test_cli("fingerprint", ['--no-fsname', pub_key2],
-             "B0:F4:98:6E:D8:4E:05:63:A1:D8:4B:37:61:5A:A0:41:78:7E:DE:0E:72:46:E0:A8:D6:CF:09:54:08:DA:A4:22")
+             "1E:10:69:84:37:63:4F:66:84:34:E2:4B:2B:A1:7A:3C:92:0F:B3:BA:D1:5F:CA:A8:66:47:89:90:7E:91:D1:6A")
 
 def cli_pbkdf_tune_tests(_tmp_dir):
     if not check_for_command("pbkdf_tune"):