From c8326b073299f68c07e764f95ee922798089aa7f Mon Sep 17 00:00:00 2001
From: Raoul Wols <raoul.wols@keysight.com>
Date: Thu, 15 May 2025 11:24:08 +0200
Subject: [PATCH] INS-14641 Add a command to do just Dilithium NTT

For measuring leakage in the forward NTT stage of Dilithium. For example,
when the challenge is or the S1 private key is converted to the
frequency domain.
---
 src/dilithium/wrapper.c | 10 ++++++++++
 src/dilithium/wrapper.h | 11 +++++++++++
 src/main.c              | 11 +++++++++++
 src/main.h              |  9 +++++++++
 4 files changed, 41 insertions(+)
 mode change 100644 => 100755 src/dilithium/wrapper.c
 mode change 100644 => 100755 src/dilithium/wrapper.h
 mode change 100644 => 100755 src/main.c
 mode change 100644 => 100755 src/main.h

diff --git a/src/dilithium/wrapper.c b/src/dilithium/wrapper.c
old mode 100644
new mode 100755
index fa966f4..942ae78
--- a/src/dilithium/wrapper.c
+++ b/src/dilithium/wrapper.c
@@ -2,6 +2,7 @@
 #include "./params.h"
 #include "./api.h"
 #include "./sign.h"
+#include "./poly.h"
 #include <string.h>
 
 #if DILITHIUM_PUBLIC_KEY_SIZE != CRYPTO_PUBLICKEYBYTES
@@ -13,6 +14,9 @@
 #if DILITHIUM_SIGNATURE_SIZE != CRYPTO_BYTES
 #error invalid signature size, update me!
 #endif
+#if DILITHIUM_N != N
+#error invalid N, update me!
+#endif
 
 int getDilithiumAlgorithmVariant() {
 	return DILITHIUM_MODE;
@@ -39,3 +43,9 @@ int DilithiumState_sign(const DilithiumState* self, uint8_t* signature, const ui
 	size_t signatureSize = DILITHIUM_SIGNATURE_SIZE;
 	return crypto_sign_signature(signature, &signatureSize, message, DILITHIUM_MESSAGE_SIZE, self->m_sk);
 }
+
+int Dilithium_ntt(uint32_t* coefficients) {
+	poly* coeffs = (poly*)coefficients;
+	poly_ntt(coeffs);
+	return 0;
+}
diff --git a/src/dilithium/wrapper.h b/src/dilithium/wrapper.h
old mode 100644
new mode 100755
index e98daae..5952b0e
--- a/src/dilithium/wrapper.h
+++ b/src/dilithium/wrapper.h
@@ -8,6 +8,7 @@
 #define DILITHIUM_PRIVATE_KEY_SIZE 4016
 #define DILITHIUM_SIGNATURE_SIZE 3293
 #define DILITHIUM_MESSAGE_SIZE 16
+#define DILITHIUM_N 256
 #define DILITHIUM_SIGNED_MESSAGE_SIZE (DILITHIUM_SIGNATURE_SIZE + DILITHIUM_MESSAGE_SIZE)
 
 /**
@@ -79,4 +80,14 @@ int DilithiumState_verify(const DilithiumState* self, uint8_t *signedMessage);
  */
 int DilithiumState_sign(const DilithiumState* self, uint8_t* signature, const uint8_t* message);
 
+///
+/// @brief        Perform a forward NTT.
+///
+/// @param[inout] coefficients  Buffer of polynomial coefficients in integer
+///                             domain. The computation is done in-place, and
+///                             this array contains the coefficients in the
+///                             frequency domain after this function returns.
+///
+int Dilithium_ntt(uint32_t *coefficients);
+
 #endif // _DILITHIUM_WRAPPER_H_
diff --git a/src/main.c b/src/main.c
old mode 100644
new mode 100755
index 72db00f..86a4770
--- a/src/main.c
+++ b/src/main.c
@@ -298,6 +298,17 @@ int main(void) {
 				break;
 			}
 
+			case CMD_SW_DILITHIUM_NTT: {
+				int32_t polynomialBuffer[DILITHIUM_N];
+				// Receive the polynomial coefficients.
+				get_bytes(sizeof(int32_t)*DILITHIUM_N, (uint8_t*)polynomialBuffer);
+				BEGIN_INTERESTING_STUFF;
+				Dilithium_ntt(polynomialBuffer);
+				END_INTERESTING_STUFF;
+				// No reply is sent.
+				break;
+			}
+
 			case CMD_SW_KYBER512_SET_PUBLIC_AND_PRIVATE_KEY: {
 				// Receive the input parameters and handle the request.
 				get_bytes(KYBER512_PUBLIC_KEY_SIZE, Kyber512State_getPublicKey(&kyber512));
diff --git a/src/main.h b/src/main.h
old mode 100644
new mode 100755
index 5e7ddd6..baa1c1c
--- a/src/main.h
+++ b/src/main.h
@@ -216,6 +216,15 @@
 ///   16-bit unsigned integer in little endian order that contains the private key size
 #define CMD_SW_DILITHIUM_GET_KEY_SIZES 0x94
 
+/// Perform Dilithium NTT.
+///
+/// Expected Input:
+///   A total of DILITHIUM_N 32-bit integers in little endian order.
+///
+/// Output:
+///   No reply is sent back.
+#define CMD_SW_DILITHIUM_NTT 0x9A
+
 #define CMD_SWDES_ENC_MISALIGNED 0x14
 #define CMD_SWAES128_ENC_MISALIGNED 0x1E
 #define CMD_SWDES_ENC_DUMMYROUNDS 0x15