Improve LMS parsing

9 files changed, 42 insertions, 23 deletions

diff --git a/crypto/src/pqc/crypto/lms/HSSSigner.cs b/crypto/src/pqc/crypto/lms/HSSSigner.cs
index 295caa00e..539226aba 100644
--- a/crypto/src/pqc/crypto/lms/HSSSigner.cs
+++ b/crypto/src/pqc/crypto/lms/HSSSigner.cs
@@ -41,6 +41,10 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
             {
                 return Hss.VerifySignature(pubKey, HssSignature.GetInstance(signature, pubKey.Level), message);
             }
+            catch (InvalidDataException e)
+            {
+                throw new Exception($"unable to decode signature: {e.Message}");
+            }
             catch (IOException e)
             {
                 throw new Exception($"unable to decode signature: {e.Message}");
diff --git a/crypto/src/pqc/crypto/lms/LMOtsParameters.cs b/crypto/src/pqc/crypto/lms/LMOtsParameters.cs
index ce472c18d..0342c0863 100644
--- a/crypto/src/pqc/crypto/lms/LMOtsParameters.cs
+++ b/crypto/src/pqc/crypto/lms/LMOtsParameters.cs
@@ -1,9 +1,11 @@
 using System;
 using System.Collections.Generic;
+using System.IO;
 
 using Org.BouncyCastle.Asn1;
 using Org.BouncyCastle.Asn1.Nist;
 using Org.BouncyCastle.Utilities.Collections;
+using Org.BouncyCastle.Utilities.IO;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Lms
 {
@@ -29,7 +31,7 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
         public static LMOtsParameters shake256_n24_w4 = new LMOtsParameters(15, 24, 4, 51, 4, 1500, NistObjectIdentifiers.IdShake256Len);
         public static LMOtsParameters shake256_n24_w8 = new LMOtsParameters(16, 24, 8, 26, 0, 1020, NistObjectIdentifiers.IdShake256Len);
 
-        private static Dictionary<object, LMOtsParameters> Suppliers = new Dictionary<object, LMOtsParameters>
+        private static Dictionary<int, LMOtsParameters> ParametersByID = new Dictionary<int, LMOtsParameters>
         {
             { sha256_n32_w1.ID, sha256_n32_w1 },
             { sha256_n32_w2.ID, sha256_n32_w2 },
@@ -53,7 +55,15 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
         };
 
         public static LMOtsParameters GetParametersByID(int id) =>
-            CollectionUtilities.GetValueOrNull(Suppliers, id);
+            CollectionUtilities.GetValueOrNull(ParametersByID, id);
+
+        internal static LMOtsParameters ParseByID(BinaryReader binaryReader)
+        {
+            int id = BinaryReaders.ReadInt32BigEndian(binaryReader);
+            if (!ParametersByID.TryGetValue(id, out var parameters))
+                throw new InvalidDataException($"unknown LMOtsParameters {id}");
+            return parameters;
+        }
 
         private readonly int m_id;
         private readonly int m_n;
diff --git a/crypto/src/pqc/crypto/lms/LMOtsPublicKey.cs b/crypto/src/pqc/crypto/lms/LMOtsPublicKey.cs
index 0f5d956e8..93519bb05 100644
--- a/crypto/src/pqc/crypto/lms/LMOtsPublicKey.cs
+++ b/crypto/src/pqc/crypto/lms/LMOtsPublicKey.cs
@@ -42,8 +42,7 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
 
         internal static LMOtsPublicKey Parse(BinaryReader binaryReader)
         {
-            int index = BinaryReaders.ReadInt32BigEndian(binaryReader);
-            LMOtsParameters parameter = LMOtsParameters.GetParametersByID(index);
+            LMOtsParameters parameter = LMOtsParameters.ParseByID(binaryReader);
 
             byte[] I = BinaryReaders.ReadBytesFully(binaryReader, 16);
 
diff --git a/crypto/src/pqc/crypto/lms/LMOtsSignature.cs b/crypto/src/pqc/crypto/lms/LMOtsSignature.cs
index f337461ac..58c5b5a46 100644
--- a/crypto/src/pqc/crypto/lms/LMOtsSignature.cs
+++ b/crypto/src/pqc/crypto/lms/LMOtsSignature.cs
@@ -40,8 +40,7 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
 
         internal static LMOtsSignature Parse(BinaryReader binaryReader)
         {
-            int index = BinaryReaders.ReadInt32BigEndian(binaryReader);
-            LMOtsParameters parameter = LMOtsParameters.GetParametersByID(index);
+            LMOtsParameters parameter = LMOtsParameters.ParseByID(binaryReader);
 
             byte[] C = BinaryReaders.ReadBytesFully(binaryReader, parameter.N);
 
diff --git a/crypto/src/pqc/crypto/lms/LMSPrivateKeyParameters.cs b/crypto/src/pqc/crypto/lms/LMSPrivateKeyParameters.cs
index 5b34f9713..b640d31bd 100644
--- a/crypto/src/pqc/crypto/lms/LMSPrivateKeyParameters.cs
+++ b/crypto/src/pqc/crypto/lms/LMSPrivateKeyParameters.cs
@@ -103,11 +103,8 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
             if (version != 0)
                 throw new Exception("unknown version for LMS private key");
 
-            int sigParamType = BinaryReaders.ReadInt32BigEndian(binaryReader);
-            LMSigParameters sigParameter = LMSigParameters.GetParametersByID(sigParamType);
-
-            int otsParamType = BinaryReaders.ReadInt32BigEndian(binaryReader);
-            LMOtsParameters otsParameter = LMOtsParameters.GetParametersByID(otsParamType);
+            LMSigParameters sigParameter = LMSigParameters.ParseByID(binaryReader);
+            LMOtsParameters otsParameter = LMOtsParameters.ParseByID(binaryReader);
 
             byte[] I = BinaryReaders.ReadBytesFully(binaryReader, 16);
 
diff --git a/crypto/src/pqc/crypto/lms/LMSPublicKeyParameters.cs b/crypto/src/pqc/crypto/lms/LMSPublicKeyParameters.cs
index 506b78096..9eda4b4fa 100644
--- a/crypto/src/pqc/crypto/lms/LMSPublicKeyParameters.cs
+++ b/crypto/src/pqc/crypto/lms/LMSPublicKeyParameters.cs
@@ -42,17 +42,14 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
 
         internal static LmsPublicKeyParameters Parse(BinaryReader binaryReader)
         {
-            int pubType = BinaryReaders.ReadInt32BigEndian(binaryReader);
-            LMSigParameters lmsParameter = LMSigParameters.GetParametersByID(pubType);
-
-            int index = BinaryReaders.ReadInt32BigEndian(binaryReader);
-            LMOtsParameters ostTypeCode = LMOtsParameters.GetParametersByID(index);
+            LMSigParameters sigParameter = LMSigParameters.ParseByID(binaryReader);
+            LMOtsParameters otsParameter = LMOtsParameters.ParseByID(binaryReader);
 
             byte[] I = BinaryReaders.ReadBytesFully(binaryReader, 16);
 
-            byte[] T1 = BinaryReaders.ReadBytesFully(binaryReader, lmsParameter.M);
+            byte[] T1 = BinaryReaders.ReadBytesFully(binaryReader, sigParameter.M);
 
-            return new LmsPublicKeyParameters(lmsParameter, ostTypeCode, T1, I);
+            return new LmsPublicKeyParameters(sigParameter, otsParameter, T1, I);
         }
 
         public override byte[] GetEncoded() => ToByteArray();
diff --git a/crypto/src/pqc/crypto/lms/LMSSignature.cs b/crypto/src/pqc/crypto/lms/LMSSignature.cs
index 75b0bdfbb..7031b156c 100644
--- a/crypto/src/pqc/crypto/lms/LMSSignature.cs
+++ b/crypto/src/pqc/crypto/lms/LMSSignature.cs
@@ -46,17 +46,16 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
 
             LMOtsSignature otsSignature = LMOtsSignature.Parse(binaryReader);
 
-            int index = BinaryReaders.ReadInt32BigEndian(binaryReader);
-            LMSigParameters type = LMSigParameters.GetParametersByID(index);
+            LMSigParameters sigParameter = LMSigParameters.ParseByID(binaryReader);
 
-            byte[][] path = new byte[type.H][];
+            byte[][] path = new byte[sigParameter.H][];
             for (int h = 0; h < path.Length; h++)
             {
-                path[h] = new byte[type.M];
+                path[h] = new byte[sigParameter.M];
                 binaryReader.Read(path[h], 0, path[h].Length);
             }
 
-            return new LmsSignature(q, otsSignature, type, path);
+            return new LmsSignature(q, otsSignature, sigParameter, path);
         }
 
         // TODO[api] Fix parameter name
diff --git a/crypto/src/pqc/crypto/lms/LMSSigner.cs b/crypto/src/pqc/crypto/lms/LMSSigner.cs
index edbfd7330..0c3b06674 100644
--- a/crypto/src/pqc/crypto/lms/LMSSigner.cs
+++ b/crypto/src/pqc/crypto/lms/LMSSigner.cs
@@ -41,6 +41,10 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
             {
                 return Lms.VerifySignature(m_publicKey, LmsSignature.GetInstance(signature), message);
             }
+            catch (InvalidDataException e)
+            {
+                throw new Exception($"unable to decode signature: {e.Message}");
+            }
             catch (IOException e)
             {
                 throw new Exception($"unable to decode signature: {e.Message}");
diff --git a/crypto/src/pqc/crypto/lms/LMSigParameters.cs b/crypto/src/pqc/crypto/lms/LMSigParameters.cs
index 721560cb8..de3918a09 100644
--- a/crypto/src/pqc/crypto/lms/LMSigParameters.cs
+++ b/crypto/src/pqc/crypto/lms/LMSigParameters.cs
@@ -1,8 +1,10 @@
 using System.Collections.Generic;
+using System.IO;
 
 using Org.BouncyCastle.Asn1;
 using Org.BouncyCastle.Asn1.Nist;
 using Org.BouncyCastle.Utilities.Collections;
+using Org.BouncyCastle.Utilities.IO;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Lms
 {
@@ -62,6 +64,14 @@ namespace Org.BouncyCastle.Pqc.Crypto.Lms
         public static LMSigParameters GetParametersByID(int id) =>
             CollectionUtilities.GetValueOrNull(ParametersByID, id);
 
+        internal static LMSigParameters ParseByID(BinaryReader binaryReader)
+        {
+            int id = BinaryReaders.ReadInt32BigEndian(binaryReader);
+            if (!ParametersByID.TryGetValue(id, out var parameters))
+                throw new InvalidDataException($"unknown LMSigParameters {id}");
+            return parameters;
+        }
+
         private readonly int m_id;
         private readonly int m_m;
         private readonly int m_h;