From d2df9f2a755ca299041da52900d5b09b40fd7831 Mon Sep 17 00:00:00 2001
From: David Hook <dgh@cryptoworkshop.com>
Date: Sat, 14 Sep 2024 21:05:32 +1000
Subject: [PATCH] refactoring of context setting

---
 .../pqc/crypto/slhdsa/HashSLHDSASigner.java   | 29 +++++++++---------
 .../pqc/crypto/slhdsa/SLHDSASigner.java       | 30 +++++++++++--------
 2 files changed, 32 insertions(+), 27 deletions(-)

diff --git a/core/src/main/java/org/bouncycastle/pqc/crypto/slhdsa/HashSLHDSASigner.java b/core/src/main/java/org/bouncycastle/pqc/crypto/slhdsa/HashSLHDSASigner.java
index 41afbf09e3..a7bace4e83 100644
--- a/core/src/main/java/org/bouncycastle/pqc/crypto/slhdsa/HashSLHDSASigner.java
+++ b/core/src/main/java/org/bouncycastle/pqc/crypto/slhdsa/HashSLHDSASigner.java
@@ -25,7 +25,7 @@ public class HashSLHDSASigner
 {
     private SLHDSAPrivateKeyParameters privKey;
     private SLHDSAPublicKeyParameters pubKey;
-
+    private byte[] ctx;
     private SecureRandom random;
     private Digest digest;
     private byte[] digestOidEncoding;
@@ -48,12 +48,26 @@ public void init(boolean forSigning, CipherParameters param)
                 privKey = (SLHDSAPrivateKeyParameters)param;
             }
 
+            ctx = privKey.getContext();
+
+            if (ctx.length > 255)
+            {
+                throw new IllegalArgumentException("context too long");
+            }
+
             initDigest(privKey);
         }
         else
         {
             pubKey = (SLHDSAPublicKeyParameters)param;
 
+            ctx = pubKey.getContext();
+
+            if (ctx.length > 255)
+            {
+                throw new IllegalArgumentException("context too long");
+            }
+
             initDigest(pubKey);
         }
 
@@ -93,12 +107,6 @@ public byte[] generateSignature() throws CryptoException, DataLengthException
         SLHDSAEngine engine = privKey.getParameters().getEngine();
 
         engine.init(privKey.pk.seed);
-        byte[] ctx = privKey.getContext();
-
-        if (ctx.length > 255)
-        {
-            throw new RuntimeException("Context too long");
-        }
 
         byte[] hash = new byte[digest.getDigestSize()];
         digest.doFinal(hash, 0);
@@ -118,13 +126,6 @@ public byte[] generateSignature() throws CryptoException, DataLengthException
     @Override
     public boolean verifySignature(byte[] signature)
     {
-        byte[] ctx = pubKey.getContext();
-
-        if (ctx.length > 255)
-        {
-            throw new RuntimeException("Context too long");
-        }
-
         byte[] hash = new byte[digest.getDigestSize()];
         digest.doFinal(hash, 0);
 
diff --git a/core/src/main/java/org/bouncycastle/pqc/crypto/slhdsa/SLHDSASigner.java b/core/src/main/java/org/bouncycastle/pqc/crypto/slhdsa/SLHDSASigner.java
index 686924f5c0..5264882eaf 100644
--- a/core/src/main/java/org/bouncycastle/pqc/crypto/slhdsa/SLHDSASigner.java
+++ b/core/src/main/java/org/bouncycastle/pqc/crypto/slhdsa/SLHDSASigner.java
@@ -22,7 +22,7 @@ public class SLHDSASigner
 {
     private SLHDSAPrivateKeyParameters privKey;
     private SLHDSAPublicKeyParameters pubKey;
-
+    private byte[] ctx;
     private SecureRandom random;
 
     /**
@@ -48,11 +48,26 @@ public void init(boolean forSigning, CipherParameters param)
                 privKey = (SLHDSAPrivateKeyParameters)param;
             }
 
+            ctx = privKey.getContext();
+
+            if (ctx.length > 255)
+            {
+                throw new IllegalArgumentException("context too long");
+            }
+
             isPreHash = privKey.parameters.isPreHash();
         }
         else
         {
             pubKey = (SLHDSAPublicKeyParameters)param;
+
+            ctx = pubKey.getContext();
+
+            if (ctx.length > 255)
+            {
+                throw new IllegalArgumentException("context too long");
+            }
+
             isPreHash = pubKey.parameters.isPreHash();
         }
 
@@ -67,12 +82,6 @@ public byte[] generateSignature(byte[] message)
         SLHDSAEngine engine = privKey.getParameters().getEngine();
 
         engine.init(privKey.pk.seed);
-        byte[] ctx = privKey.getContext();
-
-        if (ctx.length > 255)
-        {
-            throw new RuntimeException("Context too long");
-        }
 
         byte[] ds_message = new byte[1 + 1 + ctx.length + message.length];
         ds_message[0] = 0;
@@ -88,12 +97,6 @@ public byte[] generateSignature(byte[] message)
     // Equivalent to slh_verify_internal from specs
     public boolean verifySignature(byte[] message, byte[] signature)
     {
-        byte[] ctx = pubKey.getContext();
-        if (ctx.length > 255)
-        {
-            throw new RuntimeException("Context too long");
-        }
-
         byte[] ds_message = new byte[1 + 1 + ctx.length + message.length];
         ds_message[0] = 0;
         ds_message[1] = (byte)ctx.length;
@@ -102,6 +105,7 @@ public boolean verifySignature(byte[] message, byte[] signature)
 
         return internalVerifySignature(ds_message, signature);
     }
+
     public boolean internalVerifySignature(byte[] message, byte[] signature)
     {
         //# Input: Message M, signature SIG, public key PK