Refactoring in LMS

bcgit · Nov 29, 2024 · b788586 · b788586
1 parent ae9ac87
commit b788586
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 29 deletions.
diff --git a/crypto/src/pqc/crypto/lms/HSSPrivateKeyParameters.cs b/crypto/src/pqc/crypto/lms/HSSPrivateKeyParameters.cs
@@ -155,7 +155,7 @@ protected void UpdateHierarchy(IList<LmsPrivateKeyParameters> newKeys, IList<Lms
 
         public long IndexLimit => m_indexLimit;
 
-        public long GetUsagesRemaining() => m_indexLimit - m_index;
+        public long GetUsagesRemaining() => IndexLimit - GetIndex();
 
         internal LmsPrivateKeyParameters GetRootKey() => m_keys[0];
 
@@ -172,22 +172,22 @@ public HssPrivateKeyParameters ExtractKeyShard(int usageCount)
         {
             lock (this)
             {
-                if (GetUsagesRemaining() < usageCount)
-                    throw new ArgumentException("usageCount exceeds usages remaining in current leaf");
+                if (usageCount < 0)
+                    throw new ArgumentOutOfRangeException("cannot be negative", nameof(usageCount));
+                if (usageCount > m_indexLimit - m_index)
+                    throw new ArgumentException("exceeds usages remaining in current leaf", nameof(usageCount));
 
-                long maxIndexForShard = m_index + usageCount;
-                long shardStartIndex = m_index;
+                long shardIndex = m_index;
+                long shardIndexLimit = m_index + usageCount;
 
-                //
                 // Move this key's index along
-                //
-                m_index += usageCount;
+                m_index = shardIndexLimit;
 
                 var keys = new List<LmsPrivateKeyParameters>(this.GetKeys());
                 var sig = new List<LmsSignature>(this.GetSig());
 
                 HssPrivateKeyParameters shard = MakeCopy(
-                    new HssPrivateKeyParameters(m_level, keys, sig, shardStartIndex, maxIndexForShard, true));
+                    new HssPrivateKeyParameters(m_level, keys, sig, shardIndex, shardIndexLimit, isShard: true));
 
                 ResetKeyToIndex();
 

diff --git a/crypto/src/pqc/crypto/lms/LMSPrivateKeyParameters.cs b/crypto/src/pqc/crypto/lms/LMSPrivateKeyParameters.cs
@@ -211,13 +211,18 @@ public LmsPrivateKeyParameters ExtractKeyShard(int usageCount)
         {
             lock (this)
             {
-                if (q + usageCount >= maxQ)
-                    throw new ArgumentException("usageCount exceeds usages remaining");
+                if (usageCount < 0)
+                    throw new ArgumentOutOfRangeException("cannot be negative", nameof(usageCount));
+                if (usageCount > maxQ - q)
+                    throw new ArgumentException("exceeds usages remaining", nameof(usageCount));
 
-                LmsPrivateKeyParameters keyParameters = new LmsPrivateKeyParameters(this, q, q + usageCount);
-                q += usageCount;
+                int shardIndex = q;
+                int shardIndexLimit = q + usageCount;
 
-                return keyParameters;
+                // Move this key's index along
+                q = shardIndexLimit;
+
+                return new LmsPrivateKeyParameters(this, shardIndex, shardIndexLimit);
             }
         }
 
@@ -235,7 +240,10 @@ public LmsPrivateKeyParameters ExtractKeyShard(int usageCount)
 
         public byte[] GetMasterSecret() => Arrays.Clone(masterSecret);
 
-        public long GetUsagesRemaining() => maxQ - GetIndex();
+        public int IndexLimit => maxQ;
+
+        // TODO[api] Only needs 'int'
+        public long GetUsagesRemaining() => IndexLimit - GetIndex();
 
         public LmsPublicKeyParameters GetPublicKey()
         {

diff --git a/crypto/src/pqc/crypto/utils/PqcPublicKeyFactory.cs b/crypto/src/pqc/crypto/utils/PqcPublicKeyFactory.cs
@@ -4,7 +4,6 @@
 
 using Org.BouncyCastle.Asn1;
 using Org.BouncyCastle.Asn1.BC;
-using Org.BouncyCastle.Asn1.Nist;
 using Org.BouncyCastle.Asn1.Pkcs;
 using Org.BouncyCastle.Asn1.X509;
 using Org.BouncyCastle.Crypto;
@@ -232,21 +231,53 @@ internal static DilithiumPublicKeyParameters GetDilithiumPublicKey(DilithiumPara
 
         private static AsymmetricKeyParameter LmsConverter(SubjectPublicKeyInfo keyInfo, object defaultParams)
         {
-            byte[] keyEnc = Asn1OctetString.GetInstance(keyInfo.ParsePublicKey()).GetOctets();
+            DerBitString publicKey = keyInfo.PublicKey;
 
-            if (Pack.BE_To_UInt32(keyEnc, 0) == 1U)
-            {
-                return LmsPublicKeyParameters.GetInstance(Arrays.CopyOfRange(keyEnc, 4, keyEnc.Length));
-            }
-            else
+            if (publicKey.IsOctetAligned())
             {
-                // public key with extra tree height
-                if (keyEnc.Length == 64)
+                //int expectedLength = ???;
+
+                //int bytesLength = publicKey.GetBytesLength();
+                //if (bytesLength == expectedLength)
+                //    return GetLmsKeyParameters(publicKey.GetOctets());
+
+                // TODO[pqc] Remove support for legacy/prototype formats?
+                //if (bytesLength > expectedLength)
                 {
-                    keyEnc = Arrays.CopyOfRange(keyEnc, 4, keyEnc.Length);
+                    try
+                    {
+                        Asn1Object obj = Asn1Object.FromMemoryStream(publicKey.GetOctetMemoryStream());
+                        if (obj is Asn1OctetString oct)
+                        {
+                            //if (oct.GetOctetsLength() == expectedLength)
+                            {
+                                return GetLmsKeyParameters(oct.GetOctets());
+                            }
+                        }
+                    }
+                    catch (Exception)
+                    {
+                    }
                 }
-                return HssPublicKeyParameters.GetInstance(keyEnc);
+
+                return GetLmsKeyParameters(publicKey.GetOctets());
             }
+
+            throw new ArgumentException($"invalid LMS public key");
+        }
+
+        private static LmsKeyParameters GetLmsKeyParameters(byte[] keyEnc)
+        {
+            if (Pack.BE_To_UInt32(keyEnc, 0) == 1U)
+                return LmsPublicKeyParameters.GetInstance(Arrays.CopyOfRange(keyEnc, 4, keyEnc.Length));
+
+            // public key with extra tree height
+            if (keyEnc.Length == 64)
+            {
+                keyEnc = Arrays.CopyOfRange(keyEnc, 4, keyEnc.Length);
+            }
+
+            return HssPublicKeyParameters.GetInstance(keyEnc);
         }
 
 #pragma warning disable CS0618 // Type or member is obsolete

diff --git a/crypto/src/pqc/crypto/utils/PqcSubjectPublicKeyInfoFactory.cs b/crypto/src/pqc/crypto/utils/PqcSubjectPublicKeyInfoFactory.cs
@@ -43,15 +43,15 @@ public static SubjectPublicKeyInfo CreateSubjectPublicKeyInfo(AsymmetricKeyParam
                 byte[] encoding = Composer.Compose().U32Str(1).Bytes(lmsPublicKeyParameters).Build();
 
                 AlgorithmIdentifier algorithmIdentifier = new AlgorithmIdentifier(PkcsObjectIdentifiers.IdAlgHssLmsHashsig);
-                return new SubjectPublicKeyInfo(algorithmIdentifier, new DerOctetString(encoding));
+                return new SubjectPublicKeyInfo(algorithmIdentifier, encoding);
             }
             if (publicKey is HssPublicKeyParameters hssPublicKeyParameters)
             {
                 int L = hssPublicKeyParameters.Level;
                 byte[] encoding = Composer.Compose().U32Str(L).Bytes(hssPublicKeyParameters.LmsPublicKey).Build();
 
                 AlgorithmIdentifier algorithmIdentifier = new AlgorithmIdentifier(PkcsObjectIdentifiers.IdAlgHssLmsHashsig);
-                return new SubjectPublicKeyInfo(algorithmIdentifier, new DerOctetString(encoding));
+                return new SubjectPublicKeyInfo(algorithmIdentifier, encoding);
             }
 #pragma warning disable CS0618 // Type or member is obsolete
             if (publicKey is SphincsPlusPublicKeyParameters sphincsPlusPublicKeyParameters)

diff --git a/crypto/test/src/pqc/crypto/lms/test/HssTests.cs b/crypto/test/src/pqc/crypto/lms/test/HssTests.cs
@@ -507,6 +507,7 @@ public void TestVectorsFromReference_Expanded()
 
                 Assert.AreEqual(1024, keyPair.GetUsagesRemaining());
                 Assert.AreEqual(1024, keyPair.IndexLimit);
+                Assert.AreEqual(0, keyPair.GetIndex());
 
                 //
                 // Split the space up with a shard.
@@ -603,6 +604,7 @@ public void TestRemaining()
 
             HssPrivateKeyParameters shard = keyPair.ExtractKeyShard(10);
 
+            Assert.True(10 == shard.GetUsagesRemaining());
             Assert.True(15 == shard.IndexLimit);
             Assert.True(5 == shard.GetIndex());
 

diff --git a/crypto/test/src/pqc/crypto/test/LMSTest.cs b/crypto/test/src/pqc/crypto/test/LMSTest.cs
@@ -3,6 +3,7 @@
 using NUnit.Framework;
 
 using Org.BouncyCastle.Asn1.Pkcs;
+using Org.BouncyCastle.Asn1.X509;
 using Org.BouncyCastle.Crypto;
 using Org.BouncyCastle.Pqc.Crypto.Lms;
 using Org.BouncyCastle.Pqc.Crypto.Utilities;
@@ -97,6 +98,7 @@ public void TestKeyGenAndSignTwoSigsWithShard()
             LmsSigner signer = new LmsSigner();
 
             Assert.AreEqual(2, privKey.GetUsagesRemaining());
+            Assert.AreEqual(2, privKey.IndexLimit);
             Assert.AreEqual(0, privKey.GetIndex());
 
             signer.Init(true, privKey);
@@ -141,12 +143,19 @@ public void TestKeyGenAndSignTwoSigsWithShard()
 
             Assert.True(signer.VerifySignature(msg1, sig1));
 
-            PrivateKeyInfo pInfo = PqcPrivateKeyInfoFactory.CreatePrivateKeyInfo(kp.Private);//TODO
+            PrivateKeyInfo pInfo = PqcPrivateKeyInfoFactory.CreatePrivateKeyInfo(kp.Private);
             AsymmetricKeyParameter pKey = PqcPrivateKeyFactory.CreateKey(pInfo.GetEncoded());
 
             signer.Init(false, ((LmsPrivateKeyParameters)pKey).GetPublicKey());
 
             Assert.True(signer.VerifySignature(msg1, sig1));
+
+            SubjectPublicKeyInfo spki = PqcSubjectPublicKeyInfoFactory.CreateSubjectPublicKeyInfo(kp.Public);
+            AsymmetricKeyParameter publicKeyRT = PqcPublicKeyFactory.CreateKey(spki.GetEncoded());
+
+            signer.Init(false, publicKeyRT);
+
+            Assert.True(signer.VerifySignature(msg1, sig1));
         }
 
         [Test]