Give IStreamCalculator a generic type

2 files changed, 16 insertions, 33 deletions

diff --git a/crypto/src/cmp/ProtectedPkiMessage.cs b/crypto/src/cmp/ProtectedPkiMessage.cs
index 770fe5443..738fd8ab3 100644
--- a/crypto/src/cmp/ProtectedPkiMessage.cs
+++ b/crypto/src/cmp/ProtectedPkiMessage.cs
@@ -104,14 +104,14 @@ namespace Org.BouncyCastle.Cmp
         /// <returns>true if the provider is able to create a verifier that validates the signature, false otherwise.</returns>      
         public bool Verify(IVerifierFactory verifierFactory)
         {
-            IStreamCalculator streamCalculator = verifierFactory.CreateCalculator();
+            IStreamCalculator<IVerifier> streamCalculator = verifierFactory.CreateCalculator();
 
-            IVerifier result = (IVerifier)Process(streamCalculator);
+            IVerifier result = Process(streamCalculator);
 
             return result.IsVerified(pkiMessage.Protection.GetBytes());
         }
 
-        private object Process(IStreamCalculator streamCalculator)
+        private TResult Process<TResult>(IStreamCalculator<TResult> streamCalculator)
         {
             Asn1EncodableVector avec = new Asn1EncodableVector();
             avec.Add(pkiMessage.Header);
@@ -141,7 +141,7 @@ namespace Org.BouncyCastle.Cmp
 
             pkMacBuilder.SetParameters(parameter);
 
-            IBlockResult result = (IBlockResult)Process(pkMacBuilder.Build(password).CreateCalculator());
+            IBlockResult result = Process(pkMacBuilder.Build(password).CreateCalculator());
 
             return Arrays.ConstantTimeAreEqual(result.Collect(), this.pkiMessage.Protection.GetBytes());
         }
diff --git a/crypto/src/cmp/ProtectedPkiMessageBuilder.cs b/crypto/src/cmp/ProtectedPkiMessageBuilder.cs
index 6440c3f4e..8573d1fc8 100644
--- a/crypto/src/cmp/ProtectedPkiMessageBuilder.cs
+++ b/crypto/src/cmp/ProtectedPkiMessageBuilder.cs
@@ -92,7 +92,7 @@ namespace Org.BouncyCastle.Cmp
             if (null == body)
                 throw new InvalidOperationException("body must be set before building");
 
-            IStreamCalculator calculator = signatureFactory.CreateCalculator();
+            IStreamCalculator<IBlockResult> calculator = signatureFactory.CreateCalculator();
 
             if (!(signatureFactory.AlgorithmDetails is AlgorithmIdentifier))
             {
@@ -110,7 +110,7 @@ namespace Org.BouncyCastle.Cmp
             if (null == body)
                 throw new InvalidOperationException("body must be set before building");
 
-            IStreamCalculator calculator = factory.CreateCalculator();
+            IStreamCalculator<IBlockResult> calculator = factory.CreateCalculator();
             FinalizeHeader((AlgorithmIdentifier)factory.AlgorithmDetails);
             PkiHeader header = hdrBuilBuilder.Build();
             DerBitString protection = new DerBitString(CalculateSignature(calculator, header, body));
@@ -128,40 +128,23 @@ namespace Org.BouncyCastle.Cmp
 
         private ProtectedPkiMessage FinalizeMessage(PkiHeader header, DerBitString protection)
         {
-            if (extraCerts.Count > 0)
+            if (extraCerts.Count < 1)
+                return new ProtectedPkiMessage(new PkiMessage(header, body, protection));
+
+            CmpCertificate[] cmpCertificates = new CmpCertificate[extraCerts.Count];
+            for (int i = 0; i < cmpCertificates.Length; i++)
             {
-                CmpCertificate[] cmpCertificates = new CmpCertificate[extraCerts.Count];
-                for (int i = 0; i < cmpCertificates.Length; i++)
-                {
-                    byte[] cert = extraCerts[i].GetEncoded();
-                    cmpCertificates[i] = CmpCertificate.GetInstance(Asn1Object.FromByteArray(cert));
-                }
-
-                return new ProtectedPkiMessage(new PkiMessage(header, body, protection, cmpCertificates));
+                byte[] cert = extraCerts[i].GetEncoded();
+                cmpCertificates[i] = CmpCertificate.GetInstance(Asn1Object.FromByteArray(cert));
             }
 
-            return new ProtectedPkiMessage(new PkiMessage(header, body, protection));
+            return new ProtectedPkiMessage(new PkiMessage(header, body, protection, cmpCertificates));
         }
 
-        private byte[] CalculateSignature(IStreamCalculator signer, PkiHeader header, PkiBody body)
+        private byte[] CalculateSignature(IStreamCalculator<IBlockResult> signer, PkiHeader header, PkiBody body)
         {
             new DerSequence(header, body).EncodeTo(signer.Stream);
-            object result = signer.GetResult();
-
-            if (result is DefaultSignatureResult sigResult)
-            {
-                return sigResult.Collect();
-            }
-            else if (result is IBlockResult blockResult)
-            {
-                return blockResult.Collect();
-            }
-            else if (result is byte[] bytesResult)
-            {
-                return bytesResult;
-            }
-
-            throw new InvalidOperationException("result is not byte[] or DefaultSignatureResult");
+            return signer.GetResult().Collect();
         }
     }
 }