commit 8153ab942cbfd3b0df7cfb4851f91d5ac4ab8a22
parent 7ab24b4c9867f765cc2ff3a52dc42fb177750a0f
Author: Jared Tobin <jared@jtobin.io>
Date: Wed, 11 Feb 2026 22:38:05 +0400
feat: preserve STG stack slots across constant x20 adjustments (IMPL19)
Track constant deltas to x20 by shifting STG stack slot maps instead of
clearing them. This preserves seeded stg_secret/stg_public slots across
common prologue patterns like `sub x20, x20, #16`.
Changes:
- Add shiftStgStackMap to shift slot offsets by a signed delta
- Update Add/Sub/Adds/Subs transfer to shift when x20 += imm or x20 -= imm
- Update pre/post-indexed [x20, #imm] handling to shift instead of clear
- Remove unused updateWithSpCheckProvKind helper
- Add tests for delta tracking across add/sub x20
Semantics:
- sub x20, x20, #imm: shift slots by +imm (old offset 8 -> new offset 24)
- add x20, x20, #imm: shift slots by -imm (old offset 24 -> new offset 8)
- Non-constant x20 updates still clear the map
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Diffstat:
4 files changed, 216 insertions(+), 27 deletions(-)
diff --git a/lib/Audit/AArch64/Taint.hs b/lib/Audit/AArch64/Taint.hs
@@ -233,27 +233,28 @@ transfer instr st = case instr of
-- Arithmetic: result is join of operand taints/provenances
-- Clear stack map if SP is modified (offsets become stale)
+ -- For X20 with immediate operand, shift STG stack map instead of clearing
-- Pointer arithmetic (ptr + imm) preserves pointer kind
Add dst r1 op ->
let t = join2 (getTaint r1 st) (operandTaint op st)
p = provJoin2 (getProvenance r1 st) (operandProv op st)
k = pointerArithKind r1 op st
- in updateWithSpCheckProvKind dst t p k st
+ in updateWithX20ShiftAdd dst r1 op t p k st
Sub dst r1 op ->
let t = join2 (getTaint r1 st) (operandTaint op st)
p = provJoin2 (getProvenance r1 st) (operandProv op st)
k = pointerArithKind r1 op st
- in updateWithSpCheckProvKind dst t p k st
+ in updateWithX20ShiftSub dst r1 op t p k st
Adds dst r1 op ->
let t = join2 (getTaint r1 st) (operandTaint op st)
p = provJoin2 (getProvenance r1 st) (operandProv op st)
k = pointerArithKind r1 op st
- in updateWithSpCheckProvKind dst t p k st
+ in updateWithX20ShiftAdd dst r1 op t p k st
Subs dst r1 op ->
let t = join2 (getTaint r1 st) (operandTaint op st)
p = provJoin2 (getProvenance r1 st) (operandProv op st)
k = pointerArithKind r1 op st
- in updateWithSpCheckProvKind dst t p k st
+ in updateWithX20ShiftSub dst r1 op t p k st
Adc dst r1 r2 ->
let t = join2 (getTaint r1 st) (getTaint r2 st)
p = provJoin2 (getProvenance r1 st) (getProvenance r2 st)
@@ -520,6 +521,22 @@ clearStgStackMap :: TaintState -> TaintState
clearStgStackMap st = st
{ tsStgStack = IM.empty, tsStgStackProv = IM.empty, tsStgStackKind = IM.empty }
+-- | Shift all STG stack slot offsets by a signed delta.
+-- Used when x20 is adjusted by a constant: after `sub x20, x20, #16`, the old
+-- slot at offset 8 is now at offset 24 relative to the new x20.
+-- For `add x20, x20, #imm`, delta should be `-imm`.
+-- For `sub x20, x20, #imm`, delta should be `+imm`.
+shiftStgStackMap :: Int -> TaintState -> TaintState
+shiftStgStackMap delta st
+ | delta == 0 = st
+ | otherwise = st
+ { tsStgStack = shiftIM (tsStgStack st)
+ , tsStgStackProv = shiftIM (tsStgStackProv st)
+ , tsStgStackKind = shiftIM (tsStgStackKind st)
+ }
+ where
+ shiftIM im = IM.fromList [(k + delta, v) | (k, v) <- IM.toList im]
+
-- | Join taint, provenance, and kind into the heap bucket.
-- Used for non-stack stores to propagate secrets through heap memory.
joinHeapBucket :: Taint -> Provenance -> RegKind -> TaintState -> TaintState
@@ -559,12 +576,27 @@ setTaintLoadStgStack dst offset st
where
isPublicRoot r = r `elem` publicRoots
--- | Update register taint, provenance, and kind.
--- Clears hardware stack map if SP, STG stack map if X20.
-updateWithSpCheckProvKind :: Reg -> Taint -> Provenance -> RegKind
- -> TaintState -> TaintState
-updateWithSpCheckProvKind dst t p k st
+-- | Update for Add instruction, shifting STG stack map when x20 += imm.
+-- After `add x20, x20, #imm`, old offset O maps to new offset O - imm.
+updateWithX20ShiftAdd :: Reg -> Reg -> Operand -> Taint -> Provenance
+ -> RegKind -> TaintState -> TaintState
+updateWithX20ShiftAdd dst r1 op t p k st
| dst == SP = clearStackMap (setTaintProvKind dst t p k st)
+ | dst == X20, r1 == X20, OpImm imm <- op =
+ let delta = negate (fromInteger imm)
+ in shiftStgStackMap delta (setTaintProvKind dst t p k st)
+ | dst == X20 = clearStgStackMap (setTaintProvKind dst t p k st)
+ | otherwise = setTaintProvKind dst t p k st
+
+-- | Update for Sub instruction, shifting STG stack map when x20 -= imm.
+-- After `sub x20, x20, #imm`, old offset O maps to new offset O + imm.
+updateWithX20ShiftSub :: Reg -> Reg -> Operand -> Taint -> Provenance
+ -> RegKind -> TaintState -> TaintState
+updateWithX20ShiftSub dst r1 op t p k st
+ | dst == SP = clearStackMap (setTaintProvKind dst t p k st)
+ | dst == X20, r1 == X20, OpImm imm <- op =
+ let delta = fromInteger imm
+ in shiftStgStackMap delta (setTaintProvKind dst t p k st)
| dst == X20 = clearStgStackMap (setTaintProvKind dst t p k st)
| otherwise = setTaintProvKind dst t p k st
@@ -602,20 +634,22 @@ storeToStack src addr st = case addr of
k = getKind src st
in setStgStackTaintProvKind off t p k st
PreIndex X20 imm ->
- -- Store at x20+imm, then X20 changes
+ -- Store at x20+imm, then x20 = x20 + imm
let off = fromInteger imm
t = getTaint src st
p = getProvenance src st
k = getKind src st
+ delta = negate (fromInteger imm)
st' = setStgStackTaintProvKind off t p k st
- in clearStgStackMap st'
- PostIndex X20 _ ->
- -- Store at x20, then X20 changes (offset 0)
+ in shiftStgStackMap delta st'
+ PostIndex X20 imm ->
+ -- Store at x20, then x20 = x20 + imm
let t = getTaint src st
p = getProvenance src st
k = getKind src st
+ delta = negate (fromInteger imm)
st' = setStgStackTaintProvKind 0 t p k st
- in clearStgStackMap st'
+ in shiftStgStackMap delta st'
-- Refined heap slot: [base, #imm] where base is a public pointer
BaseImm base imm
| isPublicPointer base st ->
@@ -666,20 +700,22 @@ storePairToStack src1 src2 addr st = case addr of
in setStgStackTaintProvKind off t1 p1 k1
(setStgStackTaintProvKind (off + 8) t2 p2 k2 st)
PreIndex X20 imm ->
- -- Store at x20+imm and x20+imm+8, then X20 changes
+ -- Store at x20+imm and x20+imm+8, then x20 = x20 + imm
let off = fromInteger imm
t1 = getTaint src1 st; p1 = getProvenance src1 st; k1 = getKind src1 st
t2 = getTaint src2 st; p2 = getProvenance src2 st; k2 = getKind src2 st
+ delta = negate (fromInteger imm)
st' = setStgStackTaintProvKind off t1 p1 k1
(setStgStackTaintProvKind (off + 8) t2 p2 k2 st)
- in clearStgStackMap st'
- PostIndex X20 _ ->
- -- Store at x20 and x20+8, then X20 changes
+ in shiftStgStackMap delta st'
+ PostIndex X20 imm ->
+ -- Store at x20 and x20+8, then x20 = x20 + imm
let t1 = getTaint src1 st; p1 = getProvenance src1 st; k1 = getKind src1 st
t2 = getTaint src2 st; p2 = getProvenance src2 st; k2 = getKind src2 st
+ delta = negate (fromInteger imm)
st' = setStgStackTaintProvKind 0 t1 p1 k1
(setStgStackTaintProvKind 8 t2 p2 k2 st)
- in clearStgStackMap st'
+ in shiftStgStackMap delta st'
-- Refined heap slot: [base, #imm] where base is a public pointer
BaseImm base imm
| isPublicPointer base st ->
@@ -714,11 +750,13 @@ loadFromStack dst addr st = case addr of
-- STG stack (X20)
BaseImm X20 imm -> setTaintLoadStgStack dst (fromInteger imm) st
PreIndex X20 imm ->
- -- Load from x20+imm, then x20 changes
- clearStgStackMap (setTaintLoadStgStack dst (fromInteger imm) st)
+ -- Load from x20+imm, then x20 = x20 + imm
+ let delta = negate (fromInteger imm)
+ in shiftStgStackMap delta (setTaintLoadStgStack dst (fromInteger imm) st)
PostIndex X20 imm ->
- -- Load first, then clear (X20 changes after access)
- clearStgStackMap (setTaintLoadStgStack dst (fromInteger imm) st)
+ -- Load first, then x20 = x20 + imm
+ let delta = negate (fromInteger imm)
+ in shiftStgStackMap delta (setTaintLoadStgStack dst (fromInteger imm) st)
-- Symbol/literal loads
BaseSymbol _ _ ->
setTaintProv dst Public ProvPublic st -- GOT/PLT entry -> address
@@ -755,15 +793,17 @@ loadPairFromStack dst1 dst2 addr st = case addr of
let off = fromInteger imm
in setTaintLoadStgStack dst1 off (setTaintLoadStgStack dst2 (off + 8) st)
PreIndex X20 imm ->
- -- Load from x20+imm and x20+imm+8, then x20 changes
+ -- Load from x20+imm and x20+imm+8, then x20 = x20 + imm
let off = fromInteger imm
+ delta = negate (fromInteger imm)
st' = setTaintLoadStgStack dst1 off (setTaintLoadStgStack dst2 (off+8) st)
- in clearStgStackMap st'
+ in shiftStgStackMap delta st'
PostIndex X20 imm ->
- -- Load first, then clear (X20 changes after access)
+ -- Load first, then x20 = x20 + imm
let off = fromInteger imm
+ delta = negate (fromInteger imm)
st' = setTaintLoadStgStack dst1 off (setTaintLoadStgStack dst2 (off+8) st)
- in clearStgStackMap st'
+ in shiftStgStackMap delta st'
-- Refined heap slot: [base, #imm] where base is a public pointer
BaseImm base imm
| isPublicPointer base st ->
diff --git a/plans/ARCH19.md b/plans/ARCH19.md
@@ -0,0 +1,47 @@
+# ARCH19: STG Stack Delta Tracking
+
+## Goal
+
+Preserve STG stack slot taint across constant adjustments of `x20`
+so seeding by offset remains valid after small prologues/epilogues.
+
+## Context
+
+We currently clear all `tsStgStack*` maps when `x20` changes. This is
+safe but loses information for common prologue patterns like:
+
+ sub x20, x20, #16
+
+This wipes seeded `stg_secret` slots before later loads.
+
+## Approach
+
+Track constant deltas to `x20` by shifting STG stack slot maps:
+
+- On `add x20, x20, #imm`: shift all offsets by `-imm`.
+- On `sub x20, x20, #imm`: shift all offsets by `+imm`.
+
+This matches the address mapping: after `sub x20, x20, #16`, the old
+slot at offset `8` is now at offset `24` relative to the new `x20`.
+
+For any non-constant update to `x20`, continue to clear all maps.
+
+## Scope
+
+- `tsStgStack`, `tsStgStackProv`, `tsStgStackKind` only.
+- Do not change `tsStack*` (SP) logic.
+- Keep taint/provenance/kind semantics unchanged for all other regs.
+
+## Integration Points
+
+- Add a helper to shift STG stack maps by a signed delta.
+- Update transfer cases for `Add/Sub` where `dst == X20` and operand
+ is an immediate.
+- Update pre/post-indexed `[x20, #imm]` load/store handling to shift
+ maps by `imm` instead of clearing after access.
+
+## Risks
+
+- Incorrect sign or shift direction causes false positives/negatives.
+- Large deltas could overflow `Int`; treat as best-effort and allow
+ the existing `Int` semantics.
diff --git a/plans/IMPL19.md b/plans/IMPL19.md
@@ -0,0 +1,38 @@
+# IMPL19: STG Stack Delta Tracking
+
+## Changes
+
+1. Add `shiftStgStackMap :: Int -> TaintState -> TaintState` in
+ `lib/Audit/AArch64/Taint.hs`:
+ - If delta is 0, return input.
+ - Otherwise rebuild `tsStgStack`, `tsStgStackProv`, `tsStgStackKind`
+ by shifting keys by `delta`.
+ - Preserve values; drop any collisions by `IM.insert` semantics
+ (last wins) or use `IM.fromList` (stable order is fine).
+
+2. Update transfer for `Add/Sub`:
+ - In `Add dst r1 op` and `Sub dst r1 op`, if `dst == X20`, `r1 == X20`,
+ and `op` is `OpImm imm`, use `shiftStgStackMap` instead of
+ `clearStgStackMap` when updating `x20`.
+ - Preserve existing taint/prov/kind updates for `x20`.
+ - For any other `x20` update, keep `clearStgStackMap` behavior.
+
+3. Update pre/post-index handling in load/store helpers:
+ - In `storeToStack` and `storePairToStack`, for `PreIndex X20 imm`
+ and `PostIndex X20 imm`, replace `clearStgStackMap` with
+ `shiftStgStackMap` using `imm` (post-index applies after access,
+ so shift after storing).
+ - In `loadFromStack` and `loadPairFromStack`, do the same for
+ `PreIndex X20 imm` and `PostIndex X20 imm` after loading.
+
+## Tests
+
+- Add a small fixture and unit test that seeds `stg_secret` at offset 8,
+ then processes `sub x20, x20, #16` followed by `ldr x1, [x20, #24]` and
+ asserts the load is tainted.
+- Add a symmetric test for `add x20, x20, #16` shifting the other way.
+
+## Notes
+
+- Keep lines under 80 chars.
+- Avoid new dependencies.
diff --git a/test/Main.hs b/test/Main.hs
@@ -914,4 +914,68 @@ taintConfigTests = testGroup "TaintConfig" [
SecretIndex X8 -> pure ()
other -> assertFailure $ "wrong reason: " ++ show other
_ -> assertFailure "expected one violation"
+
+ -- STG stack delta tracking tests (IMPL19)
+ , testCase "stg_secret preserved after sub x20" $ do
+ -- Seed secret at offset 8, then sub x20, x20, #16
+ -- After sub, old offset 8 maps to new offset 24
+ let src = T.unlines
+ [ "_foo:"
+ , " sub x20, x20, #16" -- x20 -= 16, shift slots by +16
+ , " ldr x8, [x20, #24]" -- Load from new offset 24 (was 8)
+ , " ldr x9, [x21, x8]" -- Use as index -> violation
+ , " ret"
+ ]
+ cfg = TaintConfig $ Map.singleton "_foo"
+ (ArgPolicy Set.empty Set.empty (Set.singleton 8) Set.empty)
+ case auditWithConfig cfg "test" src of
+ Left e -> assertFailure $ "parse failed: " ++ show e
+ Right ar -> do
+ assertEqual "one violation" 1 (length (arViolations ar))
+ case arViolations ar of
+ [v] -> case vReason v of
+ SecretIndex X8 -> pure ()
+ other -> assertFailure $ "wrong reason: " ++ show other
+ _ -> assertFailure "expected one violation"
+
+ , testCase "stg_secret preserved after add x20" $ do
+ -- Seed secret at offset 24, then add x20, x20, #16
+ -- After add, old offset 24 maps to new offset 8
+ let src = T.unlines
+ [ "_foo:"
+ , " add x20, x20, #16" -- x20 += 16, shift slots by -16
+ , " ldr x8, [x20, #8]" -- Load from new offset 8 (was 24)
+ , " ldr x9, [x21, x8]" -- Use as index -> violation
+ , " ret"
+ ]
+ cfg = TaintConfig $ Map.singleton "_foo"
+ (ArgPolicy Set.empty Set.empty (Set.singleton 24) Set.empty)
+ case auditWithConfig cfg "test" src of
+ Left e -> assertFailure $ "parse failed: " ++ show e
+ Right ar -> do
+ assertEqual "one violation" 1 (length (arViolations ar))
+ case arViolations ar of
+ [v] -> case vReason v of
+ SecretIndex X8 -> pure ()
+ other -> assertFailure $ "wrong reason: " ++ show other
+ _ -> assertFailure "expected one violation"
+
+ , testCase "stg slot cleared on non-constant x20 update" $ do
+ -- Secret at offset 8, then x20 updated non-constantly
+ -- Should clear slot map, load yields unknown (not secret)
+ -- Note: x20 itself becomes unknown, causing a base violation too
+ let src = T.unlines
+ [ "_foo:"
+ , " add x20, x20, x1" -- Non-constant update clears map, x20 unknown
+ , " ldr x8, [x20, #8]" -- Unknown base x20 + load yields unknown
+ , " ldr x9, [x21, x8]" -- Use as index -> unknown violation
+ , " ret"
+ ]
+ cfg = TaintConfig $ Map.singleton "_foo"
+ (ArgPolicy Set.empty Set.empty (Set.singleton 8) Set.empty)
+ case auditWithConfig cfg "test" src of
+ Left e -> assertFailure $ "parse failed: " ++ show e
+ Right ar -> do
+ -- Two violations: unknown base x20, unknown index x8
+ assertEqual "two violations" 2 (length (arViolations ar))
]
