diff mbox series

[v3,15/20] target/riscv: shadow stack mmu index for shadow stack instructions

Message ID 20240807000652.1417776-16-debug@rivosinc.com
State New
Headers show
Series riscv support for control flow integrity extensions | expand

Commit Message

Deepak Gupta Aug. 7, 2024, 12:06 a.m. UTC
Shadow stack instructions shadow stack mmu index for load/stores.
`MMU_IDX_SS_ACCESS` at bit positon 3 is used as shadow stack index.
Shadow stack mmu index depend on privilege and SUM bit. If shadow stack
accesses happening in user mode, shadow stack mmu index = 0b1000. If
shaodw stack access happening in supervisor mode mmu index = 0b1001. If
shadow stack access happening in supervisor mode with SUM=1 then mmu
index = 0b1010

Signed-off-by: Deepak Gupta <debug@rivosinc.com>
---
 target/riscv/cpu.h                            | 13 ++++++++++
 target/riscv/cpu_helper.c                     |  3 +++
 target/riscv/insn_trans/trans_rva.c.inc       |  8 ++++++
 target/riscv/insn_trans/trans_rvzicfiss.c.inc |  6 +++++
 target/riscv/internals.h                      |  1 +
 target/riscv/translate.c                      | 25 +++++++++++++++++++
 6 files changed, 56 insertions(+)

Comments

Richard Henderson Aug. 7, 2024, 2:43 a.m. UTC | #1
On 8/7/24 10:06, Deepak Gupta wrote:
> Shadow stack instructions shadow stack mmu index for load/stores.
> `MMU_IDX_SS_ACCESS` at bit positon 3 is used as shadow stack index.
> Shadow stack mmu index depend on privilege and SUM bit. If shadow stack
> accesses happening in user mode, shadow stack mmu index = 0b1000. If
> shaodw stack access happening in supervisor mode mmu index = 0b1001. If
> shadow stack access happening in supervisor mode with SUM=1 then mmu
> index = 0b1010
> 
> Signed-off-by: Deepak Gupta <debug@rivosinc.com>
> ---
>   target/riscv/cpu.h                            | 13 ++++++++++
>   target/riscv/cpu_helper.c                     |  3 +++
>   target/riscv/insn_trans/trans_rva.c.inc       |  8 ++++++
>   target/riscv/insn_trans/trans_rvzicfiss.c.inc |  6 +++++
>   target/riscv/internals.h                      |  1 +
>   target/riscv/translate.c                      | 25 +++++++++++++++++++
>   6 files changed, 56 insertions(+)
> 
> diff --git a/target/riscv/cpu.h b/target/riscv/cpu.h
> index 6da94c417c..3ad220a9fe 100644
> --- a/target/riscv/cpu.h
> +++ b/target/riscv/cpu.h
> @@ -615,6 +615,19 @@ FIELD(TB_FLAGS, FCFI_ENABLED, 28, 1)
>   FIELD(TB_FLAGS, FCFI_LP_EXPECTED, 29, 1)
>   /* zicfiss needs a TB flag so that correct TB is located based on tb flags */
>   FIELD(TB_FLAGS, BCFI_ENABLED, 30, 1)
> +/*
> + * zicfiss shadow stack is special memory on which regular stores aren't
> + * allowed but shadow stack stores are allowed. Shadow stack stores can
> + * happen as `sspush` or `ssamoswap` instructions. `sspush` implicitly
> + * takes shadow stack address from CSR_SSP. But `ssamoswap` takes address
> + * from encoded input register and it will be used by supervisor software
> + * to access (read/write) user shadow stack for setting up rt_frame during
> + * signal delivery. Supervisor software will do so by setting SUM=1. Thus
> + * a TB flag is needed if SUM was 1 during TB generation to correctly
> + * reflect memory permissions to access shadow stack user memory from
> + * supervisor mode.
> + */
> +FIELD(TB_FLAGS, SUM, 31, 1)

This is already encoded into the mmu_idx as MMUIdx_S_SUM.


r~
Deepak Gupta Aug. 7, 2024, 9:23 p.m. UTC | #2
On Wed, Aug 07, 2024 at 12:43:31PM +1000, Richard Henderson wrote:
>On 8/7/24 10:06, Deepak Gupta wrote:
>>Shadow stack instructions shadow stack mmu index for load/stores.
>>`MMU_IDX_SS_ACCESS` at bit positon 3 is used as shadow stack index.
>>Shadow stack mmu index depend on privilege and SUM bit. If shadow stack
>>accesses happening in user mode, shadow stack mmu index = 0b1000. If
>>shaodw stack access happening in supervisor mode mmu index = 0b1001. If
>>shadow stack access happening in supervisor mode with SUM=1 then mmu
>>index = 0b1010
>>
>>Signed-off-by: Deepak Gupta <debug@rivosinc.com>
>>---
>>  target/riscv/cpu.h                            | 13 ++++++++++
>>  target/riscv/cpu_helper.c                     |  3 +++
>>  target/riscv/insn_trans/trans_rva.c.inc       |  8 ++++++
>>  target/riscv/insn_trans/trans_rvzicfiss.c.inc |  6 +++++
>>  target/riscv/internals.h                      |  1 +
>>  target/riscv/translate.c                      | 25 +++++++++++++++++++
>>  6 files changed, 56 insertions(+)
>>
>>diff --git a/target/riscv/cpu.h b/target/riscv/cpu.h
>>index 6da94c417c..3ad220a9fe 100644
>>--- a/target/riscv/cpu.h
>>+++ b/target/riscv/cpu.h
>>@@ -615,6 +615,19 @@ FIELD(TB_FLAGS, FCFI_ENABLED, 28, 1)
>>  FIELD(TB_FLAGS, FCFI_LP_EXPECTED, 29, 1)
>>  /* zicfiss needs a TB flag so that correct TB is located based on tb flags */
>>  FIELD(TB_FLAGS, BCFI_ENABLED, 30, 1)
>>+/*
>>+ * zicfiss shadow stack is special memory on which regular stores aren't
>>+ * allowed but shadow stack stores are allowed. Shadow stack stores can
>>+ * happen as `sspush` or `ssamoswap` instructions. `sspush` implicitly
>>+ * takes shadow stack address from CSR_SSP. But `ssamoswap` takes address
>>+ * from encoded input register and it will be used by supervisor software
>>+ * to access (read/write) user shadow stack for setting up rt_frame during
>>+ * signal delivery. Supervisor software will do so by setting SUM=1. Thus
>>+ * a TB flag is needed if SUM was 1 during TB generation to correctly
>>+ * reflect memory permissions to access shadow stack user memory from
>>+ * supervisor mode.
>>+ */
>>+FIELD(TB_FLAGS, SUM, 31, 1)
>
>This is already encoded into the mmu_idx as MMUIdx_S_SUM.

This is where I need some help / suggestion and clarifications.

`riscv_env_mmu_index` is the which does mode --> mmu index translation and that's
where `MMUIdx_S_SUM` is set.

Although above function assumes following things
    -- Only loads ands stores are supposed to do read and write.
    -- Translates env/priv --> mmu index

In case of shadow stack, we need to hold following true:
Shadow stack are not writeable via regular stores but are allowed to be readable.
Shadow stack are writeable only via shadow stack instruction.
Shadow stack instructions can't operate on non-shadow stack memory.

This let me to create a new mmu index (as you saw in patches). This mmu index is only
setup by shadow stack instruction and thus has to be known at translation time
(and that's why SUM TB flag)

There is no way of telling in `riscv_env_mmu_index` about whether mmu index is requested
for regular load/store or some other instruction (in this case shadow stack instruction).
If that is available then I think I can use `riscv_env_mmu_index`.

Question:
I see that `riscv_env_mmu_index` could be called from a bunch of places in (like
`accel/tcg/ldst_common.c.inc` as well. Does it exclude loads, stores which calculate mmu
indexes during translation (like shadow stack load, stores) ?

>
>
>r~
Richard Henderson Aug. 7, 2024, 10:57 p.m. UTC | #3
On 8/8/24 07:23, Deepak Gupta wrote:
> On Wed, Aug 07, 2024 at 12:43:31PM +1000, Richard Henderson wrote:
>> On 8/7/24 10:06, Deepak Gupta wrote:
>>> Shadow stack instructions shadow stack mmu index for load/stores.
>>> `MMU_IDX_SS_ACCESS` at bit positon 3 is used as shadow stack index.
>>> Shadow stack mmu index depend on privilege and SUM bit. If shadow stack
>>> accesses happening in user mode, shadow stack mmu index = 0b1000. If
>>> shaodw stack access happening in supervisor mode mmu index = 0b1001. If
>>> shadow stack access happening in supervisor mode with SUM=1 then mmu
>>> index = 0b1010
>>>
>>> Signed-off-by: Deepak Gupta <debug@rivosinc.com>
>>> ---
>>>  target/riscv/cpu.h                            | 13 ++++++++++
>>>  target/riscv/cpu_helper.c                     |  3 +++
>>>  target/riscv/insn_trans/trans_rva.c.inc       |  8 ++++++
>>>  target/riscv/insn_trans/trans_rvzicfiss.c.inc |  6 +++++
>>>  target/riscv/internals.h                      |  1 +
>>>  target/riscv/translate.c                      | 25 +++++++++++++++++++
>>>  6 files changed, 56 insertions(+)
>>>
>>> diff --git a/target/riscv/cpu.h b/target/riscv/cpu.h
>>> index 6da94c417c..3ad220a9fe 100644
>>> --- a/target/riscv/cpu.h
>>> +++ b/target/riscv/cpu.h
>>> @@ -615,6 +615,19 @@ FIELD(TB_FLAGS, FCFI_ENABLED, 28, 1)
>>>  FIELD(TB_FLAGS, FCFI_LP_EXPECTED, 29, 1)
>>>  /* zicfiss needs a TB flag so that correct TB is located based on tb flags */
>>>  FIELD(TB_FLAGS, BCFI_ENABLED, 30, 1)
>>> +/*
>>> + * zicfiss shadow stack is special memory on which regular stores aren't
>>> + * allowed but shadow stack stores are allowed. Shadow stack stores can
>>> + * happen as `sspush` or `ssamoswap` instructions. `sspush` implicitly
>>> + * takes shadow stack address from CSR_SSP. But `ssamoswap` takes address
>>> + * from encoded input register and it will be used by supervisor software
>>> + * to access (read/write) user shadow stack for setting up rt_frame during
>>> + * signal delivery. Supervisor software will do so by setting SUM=1. Thus
>>> + * a TB flag is needed if SUM was 1 during TB generation to correctly
>>> + * reflect memory permissions to access shadow stack user memory from
>>> + * supervisor mode.
>>> + */
>>> +FIELD(TB_FLAGS, SUM, 31, 1)
>>
>> This is already encoded into the mmu_idx as MMUIdx_S_SUM.
> 
> This is where I need some help / suggestion and clarifications.
> 
> `riscv_env_mmu_index` is the which does mode --> mmu index translation and that's
> where `MMUIdx_S_SUM` is set.
> 
> Although above function assumes following things
>     -- Only loads ands stores are supposed to do read and write.
>     -- Translates env/priv --> mmu index
> 
> In case of shadow stack, we need to hold following true:
> Shadow stack are not writeable via regular stores but are allowed to be readable.
> Shadow stack are writeable only via shadow stack instruction.
> Shadow stack instructions can't operate on non-shadow stack memory.
> 
> This let me to create a new mmu index (as you saw in patches). This mmu index is only
> setup by shadow stack instruction and thus has to be known at translation time

All good so far.

> There is no way of telling in `riscv_env_mmu_index` about whether mmu index is requested
> for regular load/store or some other instruction (in this case shadow stack instruction).
> If that is available then I think I can use `riscv_env_mmu_index`.

What you miss is that the result of riscv_env_mmu_index is stored

   ctx->mem_idx

So that takes care of U, S, SUM, M, VS, VU, etc.  All you need at
this point is to or in your shadow stack bit:

   ctx->mem_idx | MMU_IDX_SS_ACCESS

(Perhaps SS_WRITE is a better name, since read access is never prohibited?)

Note that you can do this without ifdefs -- user-only will happily accept and ignore the 
mmu index.  Also note that user-only will *not* be able to restrict access to the shadow 
stack pages in the way the spec describes.  We rely on the host mmu for read/write 
permission for user-only.  For now -- changing that is a long term goal.


> Question:
> I see that `riscv_env_mmu_index` could be called from a bunch of places in (like
> `accel/tcg/ldst_common.c.inc` as well. Does it exclude loads, stores which calculate mmu
> indexes during translation (like shadow stack load, stores) ?

It means you cannot use the legacy interfaces for the shadow stack.
The current interfaces:

  *        cpu_ld{sign}{size}{end}_mmuidx_ra(env, ptr, mmu_idx, retaddr)
  *        cpu_ld{sign}{size}{end}_mmu(env, ptr, oi, retaddr)
...
  *        cpu_st{size}{end}_mmuidx_ra(env, ptr, val, mmu_idx, retaddr)
  *        cpu_st{size}{end}_mmu(env, ptr, val, oi, retaddr)

take the mmu_idx as a parameter.

But as it happens, the shadow stack instructions are simple enough to implement all 
inline, so you won't need to call the out-of-line load/store functions from cpu helpers.


r~
Deepak Gupta Aug. 7, 2024, 11:13 p.m. UTC | #4
On Thu, Aug 08, 2024 at 08:57:47AM +1000, Richard Henderson wrote:
>On 8/8/24 07:23, Deepak Gupta wrote:
>>On Wed, Aug 07, 2024 at 12:43:31PM +1000, Richard Henderson wrote:
>>>On 8/7/24 10:06, Deepak Gupta wrote:
>>>>Shadow stack instructions shadow stack mmu index for load/stores.
>>>>`MMU_IDX_SS_ACCESS` at bit positon 3 is used as shadow stack index.
>>>>Shadow stack mmu index depend on privilege and SUM bit. If shadow stack
>>>>accesses happening in user mode, shadow stack mmu index = 0b1000. If
>>>>shaodw stack access happening in supervisor mode mmu index = 0b1001. If
>>>>shadow stack access happening in supervisor mode with SUM=1 then mmu
>>>>index = 0b1010
>>>>
>>>>Signed-off-by: Deepak Gupta <debug@rivosinc.com>
>>>>---
>>>> target/riscv/cpu.h                            | 13 ++++++++++
>>>> target/riscv/cpu_helper.c                     |  3 +++
>>>> target/riscv/insn_trans/trans_rva.c.inc       |  8 ++++++
>>>> target/riscv/insn_trans/trans_rvzicfiss.c.inc |  6 +++++
>>>> target/riscv/internals.h                      |  1 +
>>>> target/riscv/translate.c                      | 25 +++++++++++++++++++
>>>> 6 files changed, 56 insertions(+)
>>>>
>>>>diff --git a/target/riscv/cpu.h b/target/riscv/cpu.h
>>>>index 6da94c417c..3ad220a9fe 100644
>>>>--- a/target/riscv/cpu.h
>>>>+++ b/target/riscv/cpu.h
>>>>@@ -615,6 +615,19 @@ FIELD(TB_FLAGS, FCFI_ENABLED, 28, 1)
>>>> FIELD(TB_FLAGS, FCFI_LP_EXPECTED, 29, 1)
>>>> /* zicfiss needs a TB flag so that correct TB is located based on tb flags */
>>>> FIELD(TB_FLAGS, BCFI_ENABLED, 30, 1)
>>>>+/*
>>>>+ * zicfiss shadow stack is special memory on which regular stores aren't
>>>>+ * allowed but shadow stack stores are allowed. Shadow stack stores can
>>>>+ * happen as `sspush` or `ssamoswap` instructions. `sspush` implicitly
>>>>+ * takes shadow stack address from CSR_SSP. But `ssamoswap` takes address
>>>>+ * from encoded input register and it will be used by supervisor software
>>>>+ * to access (read/write) user shadow stack for setting up rt_frame during
>>>>+ * signal delivery. Supervisor software will do so by setting SUM=1. Thus
>>>>+ * a TB flag is needed if SUM was 1 during TB generation to correctly
>>>>+ * reflect memory permissions to access shadow stack user memory from
>>>>+ * supervisor mode.
>>>>+ */
>>>>+FIELD(TB_FLAGS, SUM, 31, 1)
>>>
>>>This is already encoded into the mmu_idx as MMUIdx_S_SUM.
>>
>>This is where I need some help / suggestion and clarifications.
>>
>>`riscv_env_mmu_index` is the which does mode --> mmu index translation and that's
>>where `MMUIdx_S_SUM` is set.
>>
>>Although above function assumes following things
>>    -- Only loads ands stores are supposed to do read and write.
>>    -- Translates env/priv --> mmu index
>>
>>In case of shadow stack, we need to hold following true:
>>Shadow stack are not writeable via regular stores but are allowed to be readable.
>>Shadow stack are writeable only via shadow stack instruction.
>>Shadow stack instructions can't operate on non-shadow stack memory.
>>
>>This let me to create a new mmu index (as you saw in patches). This mmu index is only
>>setup by shadow stack instruction and thus has to be known at translation time
>
>All good so far.
>
>>There is no way of telling in `riscv_env_mmu_index` about whether mmu index is requested
>>for regular load/store or some other instruction (in this case shadow stack instruction).
>>If that is available then I think I can use `riscv_env_mmu_index`.
>
>What you miss is that the result of riscv_env_mmu_index is stored
>
>  ctx->mem_idx

Yeah, that's right. thanks.

>
>So that takes care of U, S, SUM, M, VS, VU, etc.  All you need at
>this point is to or in your shadow stack bit:
>
>  ctx->mem_idx | MMU_IDX_SS_ACCESS


>
>(Perhaps SS_WRITE is a better name, since read access is never prohibited?)

Yeah MMU_IDX_SS_WRITE is probably more self-explainatory.

>
>Note that you can do this without ifdefs -- user-only will happily 
>accept and ignore the mmu index.  Also note that user-only will *not* 
>be able to restrict access to the shadow stack pages in the way the 
>spec describes.  We rely on the host mmu for read/write permission for 
>user-only.  For now -- changing that is a long term goal.

Ok, didn't know user-only will ignore.
I'll remove ifdef and simply do

ctx->mem_idx =| MMU_IDX_SS_WRITE

>
>
>>Question:
>>I see that `riscv_env_mmu_index` could be called from a bunch of places in (like
>>`accel/tcg/ldst_common.c.inc` as well. Does it exclude loads, stores which calculate mmu
>>indexes during translation (like shadow stack load, stores) ?
>
>It means you cannot use the legacy interfaces for the shadow stack.
>The current interfaces:
>
> *        cpu_ld{sign}{size}{end}_mmuidx_ra(env, ptr, mmu_idx, retaddr)
> *        cpu_ld{sign}{size}{end}_mmu(env, ptr, oi, retaddr)
>...
> *        cpu_st{size}{end}_mmuidx_ra(env, ptr, val, mmu_idx, retaddr)
> *        cpu_st{size}{end}_mmu(env, ptr, val, oi, retaddr)
>
>take the mmu_idx as a parameter.
>
>But as it happens, the shadow stack instructions are simple enough to 
>implement all inline, so you won't need to call the out-of-line 
>load/store functions from cpu helpers.

Thanks for explaining this.

>
>
>r~
diff mbox series

Patch

diff --git a/target/riscv/cpu.h b/target/riscv/cpu.h
index 6da94c417c..3ad220a9fe 100644
--- a/target/riscv/cpu.h
+++ b/target/riscv/cpu.h
@@ -615,6 +615,19 @@  FIELD(TB_FLAGS, FCFI_ENABLED, 28, 1)
 FIELD(TB_FLAGS, FCFI_LP_EXPECTED, 29, 1)
 /* zicfiss needs a TB flag so that correct TB is located based on tb flags */
 FIELD(TB_FLAGS, BCFI_ENABLED, 30, 1)
+/*
+ * zicfiss shadow stack is special memory on which regular stores aren't
+ * allowed but shadow stack stores are allowed. Shadow stack stores can
+ * happen as `sspush` or `ssamoswap` instructions. `sspush` implicitly
+ * takes shadow stack address from CSR_SSP. But `ssamoswap` takes address
+ * from encoded input register and it will be used by supervisor software
+ * to access (read/write) user shadow stack for setting up rt_frame during
+ * signal delivery. Supervisor software will do so by setting SUM=1. Thus
+ * a TB flag is needed if SUM was 1 during TB generation to correctly
+ * reflect memory permissions to access shadow stack user memory from
+ * supervisor mode.
+ */
+FIELD(TB_FLAGS, SUM, 31, 1)
 
 #ifdef TARGET_RISCV32
 #define riscv_cpu_mxl(env)  ((void)(env), MXL_RV32)
diff --git a/target/riscv/cpu_helper.c b/target/riscv/cpu_helper.c
index 5d5da8dce1..ad40b10e74 100644
--- a/target/riscv/cpu_helper.c
+++ b/target/riscv/cpu_helper.c
@@ -181,6 +181,9 @@  void cpu_get_tb_cpu_state(CPURISCVState *env, vaddr *pc,
     fs = EXT_STATUS_DIRTY;
     vs = EXT_STATUS_DIRTY;
 #else
+   flags = FIELD_DP32(flags, TB_FLAGS, SUM,
+                    ((env->mstatus & MSTATUS_SUM) == MSTATUS_SUM));
+
     flags = FIELD_DP32(flags, TB_FLAGS, PRIV, env->priv);
 
     flags |= riscv_env_mmu_index(env, 0);
diff --git a/target/riscv/insn_trans/trans_rva.c.inc b/target/riscv/insn_trans/trans_rva.c.inc
index db6c03f6a8..68b71339a3 100644
--- a/target/riscv/insn_trans/trans_rva.c.inc
+++ b/target/riscv/insn_trans/trans_rva.c.inc
@@ -132,6 +132,10 @@  static bool trans_ssamoswap_w(DisasContext *ctx, arg_amoswap_w *a)
 
     decode_save_opc(ctx);
     src1 = get_address(ctx, a->rs1, 0);
+#ifndef CONFIG_USER_ONLY
+    /* Shadow stack access and thus index is SS TLB index */
+    ss_mmu_idx = get_ss_index(ctx);
+#endif
 
     tcg_gen_atomic_xchg_tl(dest, src1, src2, ss_mmu_idx, (MO_ALIGN | MO_TESL));
     gen_set_gpr(ctx, a->rd, dest);
@@ -224,6 +228,10 @@  static bool trans_ssamoswap_d(DisasContext *ctx, arg_amoswap_w *a)
 
     decode_save_opc(ctx);
     src1 = get_address(ctx, a->rs1, 0);
+#ifndef CONFIG_USER_ONLY
+    /* Shadow stack access and thus index is SS TLB index */
+    ss_mmu_idx = get_ss_index(ctx);
+#endif
 
     tcg_gen_atomic_xchg_tl(dest, src1, src2, ss_mmu_idx, (MO_ALIGN | MO_TESQ));
     gen_set_gpr(ctx, a->rd, dest);
diff --git a/target/riscv/insn_trans/trans_rvzicfiss.c.inc b/target/riscv/insn_trans/trans_rvzicfiss.c.inc
index c538b7ad99..4e741c061d 100644
--- a/target/riscv/insn_trans/trans_rvzicfiss.c.inc
+++ b/target/riscv/insn_trans/trans_rvzicfiss.c.inc
@@ -70,6 +70,9 @@  static bool trans_sspopchk(DisasContext *ctx, arg_sspopchk *a)
     TCGv_i32 ssp_csr = tcg_constant_i32(CSR_SSP);
     TCGv data = tcg_temp_new();
     gen_helper_csrr(addr, tcg_env, ssp_csr);
+#ifndef CONFIG_USER_ONLY
+    ss_mmu_idx = get_ss_index(ctx);
+#endif
 
     tcg_gen_qemu_ld_tl(data, addr, ss_mmu_idx,
                        mxl_memop(ctx) | MO_ALIGN);
@@ -118,6 +121,9 @@  static bool trans_sspush(DisasContext *ctx, arg_sspush *a)
     TCGv_i32 ssp_csr = tcg_constant_i32(CSR_SSP);
     TCGv data = get_gpr(ctx, a->rs2, EXT_NONE);
     gen_helper_csrr(addr, tcg_env, ssp_csr);
+#ifndef CONFIG_USER_ONLY
+    ss_mmu_idx = get_ss_index(ctx);
+#endif
 
     tcg_gen_addi_tl(addr, addr, tmp);
 
diff --git a/target/riscv/internals.h b/target/riscv/internals.h
index dad0657c80..5147d6bf90 100644
--- a/target/riscv/internals.h
+++ b/target/riscv/internals.h
@@ -32,6 +32,7 @@ 
  *  - S+SUM+2STAGE      0b110
  *  - Shadow stack+U   0b1000
  *  - Shadow stack+S   0b1001
+ *  - Shadow stack+SUM 0b1010
  */
 #define MMUIdx_U            0
 #define MMUIdx_S            1
diff --git a/target/riscv/translate.c b/target/riscv/translate.c
index de375c32a1..4772191bd8 100644
--- a/target/riscv/translate.c
+++ b/target/riscv/translate.c
@@ -122,6 +122,8 @@  typedef struct DisasContext {
     bool fcfi_lp_expected;
     /* zicfiss extension, if shadow stack was enabled during TB gen */
     bool bcfi_enabled;
+    /* SUM was on during tb translation? */
+    bool sum;
 } DisasContext;
 
 static inline bool has_ext(DisasContext *ctx, uint32_t ext)
@@ -1127,6 +1129,29 @@  static uint32_t opcode_at(DisasContextBase *dcbase, target_ulong pc)
     return translator_ldl(env, &ctx->base, pc);
 }
 
+#ifndef CONFIG_USER_ONLY
+static unsigned int get_ss_index(DisasContext *ctx)
+{
+    int ss_mmu_idx = MMU_IDX_SS_ACCESS;
+
+    /*
+     * If priv mode is S then a separate index for supervisor
+     * shadow stack accesses
+     */
+    if (ctx->priv == PRV_S) {
+        ss_mmu_idx |= MMUIdx_S;
+    }
+
+    /* If SUM was set, SS index should have S cleared */
+    if (ctx->sum) {
+        ss_mmu_idx &= ~(MMUIdx_S);
+        ss_mmu_idx |= MMUIdx_S_SUM;
+    }
+
+    return ss_mmu_idx;
+}
+#endif
+
 /* Include insn module translation function */
 #include "insn_trans/trans_rvi.c.inc"
 #include "insn_trans/trans_rvm.c.inc"