[keymgr] Minor keymgr fixes

- Obey valid/ready handshake even during error case when sending random data
- Re-organize each module to contain its own lfsr_en

Signed-off-by: Timothy Chen <timothytim@google.com>
diff --git a/hw/ip/keymgr/keymgr.core b/hw/ip/keymgr/keymgr.core
index fd7ae4d..e352190 100644
--- a/hw/ip/keymgr/keymgr.core
+++ b/hw/ip/keymgr/keymgr.core
@@ -15,6 +15,7 @@
     files:
       - rtl/keymgr_reg_pkg.sv
       - rtl/keymgr_reg_top.sv
+      - rtl/keymgr_sideload_key_ctrl.sv
       - rtl/keymgr_sideload_key.sv
       - rtl/keymgr_ctrl.sv
       - rtl/keymgr_cfg_en.sv
diff --git a/hw/ip/keymgr/lint/keymgr.waiver b/hw/ip/keymgr/lint/keymgr.waiver
index 53533ba..a6cadbc 100644
--- a/hw/ip/keymgr/lint/keymgr.waiver
+++ b/hw/ip/keymgr/lint/keymgr.waiver
@@ -6,3 +6,6 @@
 
 waive -rules {MISSING_STATE} -location {keymgr_ctrl.sv} -regexp {.*StDisabled.*} \
       -comment "Disabled state absorbed into default."
+
+waive -rules {TERMINAL_STATE} -location {keymgr_sideload_key_ctrl.sv} -regexp {.*StStop.*} \
+      -comment "Intentional terminal state."
diff --git a/hw/ip/keymgr/rtl/keymgr.sv b/hw/ip/keymgr/rtl/keymgr.sv
index 69f5f24..e87c10d 100644
--- a/hw/ip/keymgr/rtl/keymgr.sv
+++ b/hw/ip/keymgr/rtl/keymgr.sv
@@ -74,7 +74,7 @@
   // have more bits than the lfsr output, the lfsr value is simply replicated
 
   logic [63:0] lfsr;
-  logic lfsr_en;
+  logic ctrl_lfsr_en, data_lfsr_en, sideload_lfsr_en;
 
   prim_lfsr #(
     .LfsrDw(64),
@@ -82,7 +82,7 @@
   ) u_lfsr (
     .clk_i,
     .rst_ni,
-    .lfsr_en_i(lfsr_en),
+    .lfsr_en_i(ctrl_lfsr_en | data_lfsr_en | sideload_lfsr_en),
     .seed_en_i(1'b0),
     .seed_i('0),
     .entropy_i('0), // TBD, this should be hooked up to an entropy distribution pkg
@@ -113,7 +113,7 @@
     .clk_i,
     .rst_ni,
     .keymgr_en_i(lc_i.keymgr_en),
-    .prng_en_o(lfsr_en),
+    .prng_en_o(ctrl_lfsr_en),
     .entropy_i(lfsr[63:32]),  // TBD, recommend directly interfacing with DRBG for keymgr_ctrl
     .init_i(reg2hw.control.init.q),
     .init_done_o(init_done),
@@ -257,6 +257,7 @@
   keymgr_kmac_if u_kmac_if (
     .clk_i,
     .rst_ni,
+    .prng_en_o(data_lfsr_en),
     .adv_data_i(adv_matrix[stage_sel]),
     .id_data_i(id_matrix[stage_sel]),
     .gen_data_i(gen_in),
@@ -279,45 +280,22 @@
   /////////////////////////////////////
   //  Side load key storage
   /////////////////////////////////////
-  keymgr_key_dest_e dest_sel;
-  logic aes_sel, hmac_sel, kmac_sel;
-
-  assign dest_sel = keymgr_key_dest_e'(reg2hw.control.dest_sel);
-  assign aes_sel  = dest_sel == Aes  & key_sel == HwKey;
-  assign hmac_sel = dest_sel == Hmac & key_sel == HwKey;
-  assign kmac_sel = dest_sel == Kmac & key_sel == HwKey;
-
-  keymgr_sideload_key u_aes_key (
+  keymgr_sideload_key_ctrl u_sideload_ctrl(
     .clk_i,
     .rst_ni,
-    .keymgr_en_i(lc_i.keymgr_en),
-    .set_i(data_valid & aes_sel),
-    .clr_i(lc_i.keymgr_en), // TBD, should add an option for software clear later
-    .entropy_i(lfsr[31:0]), //TBD, recommend directly interfacign with DRBG entropy for keys
-    .key_i(kmac_data),
-    .key_o(aes_key_o)
-  );
-
-  keymgr_sideload_key u_hmac_key (
-    .clk_i,
-    .rst_ni,
-    .keymgr_en_i(lc_i.keymgr_en),
-    .set_i(data_valid & hmac_sel),
-    .clr_i(lc_i.keymgr_en), // TBD, should add an option for software clear later
-    .entropy_i(lfsr[31:0]),
-    .key_i(kmac_data),
-    .key_o(hmac_key_o)
-  );
-
-  keymgr_sideload_key u_kmac_key (
-    .clk_i,
-    .rst_ni,
-    .keymgr_en_i(lc_i.keymgr_en),
-    .set_i(load_key | (data_valid & kmac_sel)),
-    .clr_i(~lc_i.keymgr_en), // TBD, should add an option for software clear later
-    .entropy_i(lfsr[31:0]),
-    .key_i(load_key ? kmac_key : kmac_data),
-    .key_o(kmac_key_o)
+    .en_i(lc_i.keymgr_en),
+    .init_i(init_done),
+    .entropy_i(lfsr[63:32]),
+    .dest_sel_i(keymgr_key_dest_e'(reg2hw.control.dest_sel)),
+    .key_sel_i(key_sel),
+    .load_key_i(load_key),
+    .data_valid_i(data_valid),
+    .key_i(kmac_key),
+    .data_i(kmac_data),
+    .prng_en_o(sideload_lfsr_en),
+    .aes_key_o(aes_key_o),
+    .hmac_key_o(hmac_key_o),
+    .kmac_key_o(kmac_key_o)
   );
 
   for (genvar i = 0; i < 8; i++) begin : gen_sw_assigns
@@ -435,13 +413,6 @@
     .alert_tx_o(alert_tx_o[1])
   );
 
-  /////////////////////////////////////
-  //  Assertions
-  /////////////////////////////////////
-
-  // Only 1 entity should be trying to use the secret kmac key input
-  `ASSERT(KmacKeyLoadExclusive_a, $onehot0({load_key, data_valid & kmac_sel}))
-
   // known asserts
   `ASSERT_KNOWN(TlDValidKnownO_A, tl_o.d_valid)
   `ASSERT_KNOWN(TlAReadyKnownO_A, tl_o.a_ready)
diff --git a/hw/ip/keymgr/rtl/keymgr_ctrl.sv b/hw/ip/keymgr/rtl/keymgr_ctrl.sv
index bbdfc4d..6021772 100644
--- a/hw/ip/keymgr/rtl/keymgr_ctrl.sv
+++ b/hw/ip/keymgr/rtl/keymgr_ctrl.sv
@@ -144,7 +144,7 @@
     stage_sel_o = Disable;
 
     // enable prng toggling
-    prng_en_o = 1'b1;
+    prng_en_o = 1'b0;
 
     op_done_o = 1'b0;
     init_done_o = 1'b0;
@@ -175,6 +175,8 @@
 
       // This state does not accept any command.
       StWipe: begin
+        prng_en_o = 1'b1;
+
         // populate both shares with the same entropy
         // This is the default mask
         if (cnt < EntropyRounds) begin
diff --git a/hw/ip/keymgr/rtl/keymgr_kmac_if.sv b/hw/ip/keymgr/rtl/keymgr_kmac_if.sv
index c26e9ab..4f357b7 100644
--- a/hw/ip/keymgr/rtl/keymgr_kmac_if.sv
+++ b/hw/ip/keymgr/rtl/keymgr_kmac_if.sv
@@ -30,6 +30,7 @@
   input kmac_data_rsp_t kmac_data_i,
 
   // entropy input
+  output logic prng_en_o,
   input [31:0] entropy_i,
 
   // error outputs
@@ -272,4 +273,7 @@
   // command error occurs if kmac errors or if the command itself is invalid
   assign cmd_error_o = |(enables & enables_sub);
 
+  // request entropy to churn whenever a transaction is accepted
+  assign prng_en_o = kmac_data_o.valid & kmac_data_i.ready;
+
 endmodule // keymgr_kmac_if
diff --git a/hw/ip/keymgr/rtl/keymgr_sideload_key.sv b/hw/ip/keymgr/rtl/keymgr_sideload_key.sv
index e9b8599..d69cb1a 100644
--- a/hw/ip/keymgr/rtl/keymgr_sideload_key.sv
+++ b/hw/ip/keymgr/rtl/keymgr_sideload_key.sv
@@ -9,7 +9,7 @@
 module keymgr_sideload_key import keymgr_pkg::*;(
   input clk_i,
   input rst_ni,
-  input keymgr_en_i,
+  input en_i,
   input set_i,
   input clr_i,
   input [31:0] entropy_i,
@@ -22,14 +22,14 @@
   logic valid_q;
   logic [Shares-1:0][KeyWidth-1:0] key_q;
 
-  assign key_o.valid = valid_q & keymgr_en_i;
+  assign key_o.valid = valid_q & en_i;
   assign key_o.key_share0 = key_q[0];
   assign key_o.key_share1 = key_q[1];
 
   always_ff @(posedge clk_i or negedge rst_ni) begin
     if (!rst_ni) begin
       valid_q <= 1'b0;
-    end else if (!keymgr_en_i || clr_i) begin
+    end else if (!en_i || clr_i) begin
       valid_q <= 1'b0;
     end else if (set_i) begin
       valid_q <= 1'b1;
@@ -37,7 +37,7 @@
   end
 
   always_ff @(posedge clk_i) begin
-    if (!keymgr_en_i || clr_i) begin
+    if (clr_i) begin
       for (int i = 0; i < Shares; i++) begin
         key_q[i] <= {EntropyCopies{entropy_i}};
       end
diff --git a/hw/ip/keymgr/rtl/keymgr_sideload_key_ctrl.sv b/hw/ip/keymgr/rtl/keymgr_sideload_key_ctrl.sv
new file mode 100644
index 0000000..c1e96f0
--- /dev/null
+++ b/hw/ip/keymgr/rtl/keymgr_sideload_key_ctrl.sv
@@ -0,0 +1,141 @@
+// Copyright lowRISC contributors.
+// Licensed under the Apache License, Version 2.0, see LICENSE for details.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Manage all sideload keys
+
+`include "prim_assert.sv"
+
+module keymgr_sideload_key_ctrl import keymgr_pkg::*;(
+  input clk_i,
+  input rst_ni,
+  input en_i,
+  input init_i,
+  input [31:0] entropy_i,
+  input keymgr_key_dest_e dest_sel_i,
+  input keymgr_gen_out_e key_sel_i,
+  input load_key_i,
+  input data_valid_i,
+  input [Shares-1:0][KeyWidth-1:0] key_i,
+  input [Shares-1:0][KeyWidth-1:0] data_i,
+  output logic prng_en_o,
+  output hw_key_req_t aes_key_o,
+  output hw_key_req_t hmac_key_o,
+  output hw_key_req_t kmac_key_o
+);
+
+  // Enumeration for working state
+  typedef enum logic [1:0] {
+    StReset,
+    StIdle,
+    StWipe,
+    StStop
+  } keymgr_sideload_e;
+
+  keymgr_sideload_e state_q, state_d;
+  logic [3:0] cnt_q, cnt_d;
+  logic cnt_end;
+  logic clr;
+  logic keys_en;
+
+
+  always_ff @(posedge clk_i or negedge rst_ni) begin
+    if (!rst_ni) begin
+      state_q <= StReset;
+      cnt_q <= '0;
+    end else begin
+      state_q <= state_d;
+      cnt_q <= cnt_d;
+    end
+  end
+
+  assign cnt_end = cnt_q[3];
+  assign cnt_d = cnt_end ? cnt_q :
+                 clr     ? cnt_q + 1'b1 : cnt_q;
+
+  always_comb begin
+
+    clr = 1'b0;
+    keys_en = 1'b0;
+    state_d = state_q;
+
+    unique case (state_q)
+      StReset: begin
+        if (init_i && en_i) begin
+          state_d = StIdle;
+        end
+      end
+
+      StIdle: begin
+        keys_en = 1'b1;
+        if (!en_i) begin
+          keys_en = 1'b0;
+          state_d = StWipe;
+        end
+      end
+
+      StWipe: begin
+        keys_en = 1'b0;
+        clr = 1'b1;
+        state_d = StStop;
+      end
+
+      // intentional terminal state
+      StStop: begin
+        keys_en = 1'b0;
+      end
+
+      default:;
+    endcase // unique case (state_q)
+  end
+
+  logic aes_sel, hmac_sel, kmac_sel;
+  assign aes_sel  = dest_sel_i == Aes  & key_sel_i == HwKey;
+  assign hmac_sel = dest_sel_i == Hmac & key_sel_i == HwKey;
+  assign kmac_sel = dest_sel_i == Kmac & key_sel_i == HwKey;
+
+  keymgr_sideload_key u_aes_key (
+    .clk_i,
+    .rst_ni,
+    .en_i(keys_en),
+    .set_i(data_valid_i & aes_sel),
+    .clr_i(clr), // TBD, should add an option for software clear later
+    .entropy_i(entropy_i),
+    .key_i(data_i),
+    .key_o(aes_key_o)
+  );
+
+  keymgr_sideload_key u_hmac_key (
+    .clk_i,
+    .rst_ni,
+    .en_i(keys_en),
+    .set_i(data_valid_i & hmac_sel),
+    .clr_i(clr), // TBD, should add an option for software clear later
+    .entropy_i(entropy_i),
+    .key_i(data_i),
+    .key_o(hmac_key_o)
+  );
+
+  keymgr_sideload_key u_kmac_key (
+    .clk_i,
+    .rst_ni,
+    .en_i(keys_en),
+    .set_i(load_key_i | (data_valid_i & kmac_sel)),
+    .clr_i(clr), // TBD, should add an option for software clear laterclr
+    .entropy_i(entropy_i),
+    .key_i(load_key_i ? key_i : data_i),
+    .key_o(kmac_key_o)
+  );
+
+  // when clearing, request prng
+  assign prng_en_o = clr;
+
+
+  /////////////////////////////////////
+  //  Assertions
+  /////////////////////////////////////
+
+  // Only 1 entity should be trying to use the secret kmac key input
+  `ASSERT(KmacKeyLoadExclusive_a, $onehot0({load_key_i, data_valid_i & kmac_sel}))
+
+endmodule // keymgr_sideload_key_ctrl