[dif/kmac] add DIF to poll for status flag

This exposes a KMAC DIF to allow polling for a specific status flag.
This will be used in the power virus test to manually actuate the KMAC
block to better align all crypto operations in time.

Signed-off-by: Timothy Trippel <ttrippel@google.com>
diff --git a/sw/device/lib/dif/dif_kmac.c b/sw/device/lib/dif/dif_kmac.c
index 866abd3..50d02a7 100644
--- a/sw/device/lib/dif/dif_kmac.c
+++ b/sw/device/lib/dif/dif_kmac.c
@@ -31,11 +31,6 @@
    * The lowest security level is 128 (e.g. SHAKE128).
    */
   kDifKmacMaximumBitRate = 1600 - (2 * 128),
-
-  /**
-   * The offset of the second share within the output state register.
-   */
-  kDifKmacStateShareOffset = 0x100,
 };
 
 dif_result_t dif_kmac_customization_string_init(
@@ -151,14 +146,7 @@
   return bitfield_bit32_read(reg, KMAC_INTR_STATE_KMAC_ERR_BIT);
 }
 
-/**
- * Poll until the status register is in the 'absorb' state or int state register
- * has indicated an error.
- *
- * @param kmac
- * @return dif_result
- */
-static dif_result_t poll_state(const dif_kmac_t *kmac, uint32_t flag) {
+dif_result_t dif_kmac_poll_status(const dif_kmac_t *kmac, uint32_t flag) {
   while (true) {
     uint32_t reg = mmio_region_read32(kmac->base_addr, KMAC_STATUS_REG_OFFSET);
     if (bitfield_bit32_read(reg, flag)) {
@@ -335,7 +323,7 @@
   mmio_region_write32(kmac->base_addr, KMAC_CMD_REG_OFFSET, cmd_reg);
 
   // Poll until the status register is in the 'absorb' state.
-  return poll_state(kmac, KMAC_STATUS_SHA3_ABSORB_BIT);
+  return dif_kmac_poll_status(kmac, KMAC_STATUS_SHA3_ABSORB_BIT);
 }
 
 dif_result_t dif_kmac_mode_shake_start(
@@ -384,7 +372,7 @@
       bitfield_field32_write(0, KMAC_CMD_CMD_FIELD, KMAC_CMD_CMD_VALUE_START);
   mmio_region_write32(kmac->base_addr, KMAC_CMD_REG_OFFSET, cmd_reg);
 
-  return poll_state(kmac, KMAC_STATUS_SHA3_ABSORB_BIT);
+  return dif_kmac_poll_status(kmac, KMAC_STATUS_SHA3_ABSORB_BIT);
 }
 
 dif_result_t dif_kmac_mode_cshake_start(
@@ -484,7 +472,7 @@
       bitfield_field32_write(0, KMAC_CMD_CMD_FIELD, KMAC_CMD_CMD_VALUE_START);
   mmio_region_write32(kmac->base_addr, KMAC_CMD_REG_OFFSET, cmd_reg);
 
-  return poll_state(kmac, KMAC_STATUS_SHA3_ABSORB_BIT);
+  return dif_kmac_poll_status(kmac, KMAC_STATUS_SHA3_ABSORB_BIT);
 }
 
 dif_result_t dif_kmac_mode_kmac_start(
@@ -599,7 +587,7 @@
       bitfield_field32_write(0, KMAC_CMD_CMD_FIELD, KMAC_CMD_CMD_VALUE_START);
   mmio_region_write32(kmac->base_addr, KMAC_CMD_REG_OFFSET, cmd_reg);
 
-  return poll_state(kmac, KMAC_STATUS_SHA3_ABSORB_BIT);
+  return dif_kmac_poll_status(kmac, KMAC_STATUS_SHA3_ABSORB_BIT);
 }
 
 static void msg_fifo_write(const dif_kmac_t *kmac, const unsigned char *data,
@@ -751,7 +739,8 @@
     }
 
     // Poll the status register until in the 'squeeze' state.
-    DIF_RETURN_IF_ERROR(poll_state(kmac, KMAC_STATUS_SHA3_SQUEEZE_BIT));
+    DIF_RETURN_IF_ERROR(
+        dif_kmac_poll_status(kmac, KMAC_STATUS_SHA3_SQUEEZE_BIT));
 
     uint32_t offset =
         KMAC_STATE_REG_OFFSET + operation_state->offset * sizeof(uint32_t);
diff --git a/sw/device/lib/dif/dif_kmac.h b/sw/device/lib/dif/dif_kmac.h
index 1c33cad..6bca289 100644
--- a/sw/device/lib/dif/dif_kmac.h
+++ b/sw/device/lib/dif/dif_kmac.h
@@ -135,6 +135,10 @@
    * The length is in 32-bit words.
    */
   kDifKmacEntropySeedWords = 5,
+  /**
+   * The offset of the second share within the output state register.
+   */
+  kDifKmacStateShareOffset = 0x100,
 };
 
 /**
@@ -740,6 +744,16 @@
 OT_WARN_UNUSED_RESULT
 dif_result_t dif_kmac_config_is_locked(const dif_kmac_t *kmac, bool *is_locked);
 
+/**
+ * Poll until a given flag in the status register is set.
+ *
+ * @param kmac A KMAC handle.
+ * @param flag the
+ * @return The result of the operation.
+ */
+OT_WARN_UNUSED_RESULT
+dif_result_t dif_kmac_poll_status(const dif_kmac_t *kmac, uint32_t flag);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif  // __cplusplus
diff --git a/sw/device/sca/kmac_serial.c b/sw/device/sca/kmac_serial.c
index e0d11ff..b429820 100644
--- a/sw/device/sca/kmac_serial.c
+++ b/sw/device/sca/kmac_serial.c
@@ -47,10 +47,6 @@
    */
   kDigestLength = 8,
   /**
-   * The offset of the second share within the output state register.
-   */
-  kDifKmacStateShareOffset = 0x100,
-  /**
    * Number of cycles (at `kClockFreqCpuHz`) that Ibex should sleep to minimize
    * noise during SHA3 operations. Caution: This number should be chosen to
    * provide enough time. Otherwise, Ibex might wake up while SHA3 is still busy
diff --git a/sw/device/sca/sha3_serial.c b/sw/device/sca/sha3_serial.c
index 5fff9d9..96aa414 100644
--- a/sw/device/sca/sha3_serial.c
+++ b/sw/device/sca/sha3_serial.c
@@ -47,10 +47,6 @@
    */
   kDigestLength = 8,
   /**
-   * The offset of the second share within the output state register.
-   */
-  kDifKmacStateShareOffset = 0x100,
-  /**
    * Number of cycles (at `kClockFreqCpuHz`) that Ibex should sleep to minimize
    * noise during SHA3 operations. Caution: This number should be chosen to
    * provide enough time. Otherwise, Ibex might wake up while SHA3 is still busy