commit 584dc047e8fdfb53d28b2190d5d2f755f1967b57
parent 278bc2d33847f190c2a5b10c1d0955db70428eb5
Author: Jared Tobin <jared@jtobin.io>
Date: Wed, 11 Feb 2026 20:41:13 +0400
feat: track stack slot kinds to preserve pointer/scalar across spills (IMPL14)
Extends TaintState with tsStackKind (IntMap RegKind) so pointer vs scalar
classification survives spills and reloads.
Changes:
- Add tsStackKind field to TaintState
- Update storeToStack/storePairToStack to record source register kinds
- Update setTaintLoadStack to restore kinds from stack map
- Update clearStackMap to clear tsStackKind on SP modification
- Update joinTaintState to join stack kinds
- Add tests for pointer spill/reload and scalar spill/reload
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Diffstat:
4 files changed, 138 insertions(+), 21 deletions(-)
diff --git a/lib/Audit/AArch64/Taint.hs b/lib/Audit/AArch64/Taint.hs
@@ -76,6 +76,8 @@ data TaintState = TaintState
-- ^ Stack slot provenance (keyed by SP offset)
, tsKind :: !(Map Reg RegKind)
-- ^ Register kinds (pointer vs scalar)
+ , tsStackKind :: !(IntMap RegKind)
+ -- ^ Stack slot kinds (keyed by SP offset)
} deriving (Eq, Show)
-- | GHC 9.10.3 AArch64 public root registers.
@@ -111,6 +113,7 @@ emptyTaintState = TaintState
, tsProv = Map.empty
, tsStackProv = IM.empty
, tsKind = Map.empty
+ , tsStackKind = IM.empty
}
-- | Initial taint state with public roots marked.
@@ -122,6 +125,7 @@ initTaintState = TaintState
, tsProv = Map.fromList [(r, ProvPublic) | r <- publicRoots]
, tsStackProv = IM.empty
, tsKind = Map.fromList [(r, KindPtr) | r <- publicRoots]
+ , tsStackKind = IM.empty
}
-- | Seed argument registers according to policy.
@@ -414,28 +418,32 @@ setTaintLoad dst st
isPublicRoot r = r `elem` publicRoots
-- | Set taint for a loaded value from a known stack slot.
--- If we have tracked taint/provenance at this offset, use it; else Unknown.
--- Loaded values are KindScalar unless they're public roots (which are Ptr).
+-- If we have tracked taint/provenance/kind at this offset, use it; else Unknown.
+-- Public roots always get Public/ProvPublic/KindPtr.
setTaintLoadStack :: Reg -> Int -> TaintState -> TaintState
setTaintLoadStack dst offset st
| isPublicRoot dst = setTaintProvKind dst Public ProvPublic KindPtr st
| otherwise =
let taint = IM.findWithDefault Unknown offset (tsStack st)
prov = IM.findWithDefault ProvUnknown offset (tsStackProv st)
- in setTaintProvKind dst taint prov KindScalar st
+ kind = IM.findWithDefault KindUnknown offset (tsStackKind st)
+ in setTaintProvKind dst taint prov kind st
where
isPublicRoot r = r `elem` publicRoots
--- | Store both taint and provenance to a stack slot.
-setStackTaintProv :: Int -> Taint -> Provenance -> TaintState -> TaintState
-setStackTaintProv offset t p st = st
+-- | Store taint, provenance, and kind to a stack slot.
+setStackTaintProvKind :: Int -> Taint -> Provenance -> RegKind -> TaintState
+ -> TaintState
+setStackTaintProvKind offset t p k st = st
{ tsStack = IM.insert offset t (tsStack st)
, tsStackProv = IM.insert offset p (tsStackProv st)
+ , tsStackKind = IM.insert offset k (tsStackKind st)
}
--- | Clear all stack slot taints and provenance (when SP is modified).
+-- | Clear all stack slot taints, provenance, and kinds (when SP is modified).
clearStackMap :: TaintState -> TaintState
-clearStackMap st = st { tsStack = IM.empty, tsStackProv = IM.empty }
+clearStackMap st = st
+ { tsStack = IM.empty, tsStackProv = IM.empty, tsStackKind = IM.empty }
-- | Update register taint, provenance, and kind, clearing stack map if SP.
updateWithSpCheckProvKind :: Reg -> Taint -> Provenance -> RegKind
@@ -450,15 +458,24 @@ storeToStack :: Reg -> AddrMode -> TaintState -> TaintState
storeToStack src addr st = case addr of
BaseImm SP imm ->
let off = fromInteger imm
- in setStackTaintProv off (getTaint src st) (getProvenance src st) st
+ t = getTaint src st
+ p = getProvenance src st
+ k = getKind src st
+ in setStackTaintProvKind off t p k st
PreIndex SP imm ->
-- Store at sp+imm, then SP changes
let off = fromInteger imm
- st' = setStackTaintProv off (getTaint src st) (getProvenance src st) st
+ t = getTaint src st
+ p = getProvenance src st
+ k = getKind src st
+ st' = setStackTaintProvKind off t p k st
in clearStackMap st'
PostIndex SP _ ->
-- Store at sp, then SP changes (offset 0)
- let st' = setStackTaintProv 0 (getTaint src st) (getProvenance src st) st
+ let t = getTaint src st
+ p = getProvenance src st
+ k = getKind src st
+ st' = setStackTaintProvKind 0 t p k st
in clearStackMap st'
_ -> st -- Non-SP stores don't affect stack tracking
@@ -469,22 +486,24 @@ storePairToStack :: Reg -> Reg -> AddrMode -> TaintState -> TaintState
storePairToStack src1 src2 addr st = case addr of
BaseImm SP imm ->
let off = fromInteger imm
- st' = setStackTaintProv off (getTaint src1 st) (getProvenance src1 st)
- (setStackTaintProv (off + 8) (getTaint src2 st)
- (getProvenance src2 st) st)
- in st'
+ t1 = getTaint src1 st; p1 = getProvenance src1 st; k1 = getKind src1 st
+ t2 = getTaint src2 st; p2 = getProvenance src2 st; k2 = getKind src2 st
+ in setStackTaintProvKind off t1 p1 k1
+ (setStackTaintProvKind (off + 8) t2 p2 k2 st)
PreIndex SP imm ->
-- Store at sp+imm and sp+imm+8, then SP changes
let off = fromInteger imm
- st' = setStackTaintProv off (getTaint src1 st) (getProvenance src1 st)
- (setStackTaintProv (off + 8) (getTaint src2 st)
- (getProvenance src2 st) st)
+ t1 = getTaint src1 st; p1 = getProvenance src1 st; k1 = getKind src1 st
+ t2 = getTaint src2 st; p2 = getProvenance src2 st; k2 = getKind src2 st
+ st' = setStackTaintProvKind off t1 p1 k1
+ (setStackTaintProvKind (off + 8) t2 p2 k2 st)
in clearStackMap st'
PostIndex SP _ ->
-- Store at sp and sp+8, then SP changes
- let st' = setStackTaintProv 0 (getTaint src1 st) (getProvenance src1 st)
- (setStackTaintProv 8 (getTaint src2 st)
- (getProvenance src2 st) st)
+ 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
+ st' = setStackTaintProvKind 0 t1 p1 k1
+ (setStackTaintProvKind 8 t2 p2 k2 st)
in clearStackMap st'
_ -> st
@@ -623,6 +642,7 @@ joinTaintState a b = TaintState
, tsProv = Map.unionWith joinProvenance (tsProv a) (tsProv b)
, tsStackProv = IM.unionWith joinProvenance (tsStackProv a) (tsStackProv b)
, tsKind = Map.unionWith joinKind (tsKind a) (tsKind b)
+ , tsStackKind = IM.unionWith joinKind (tsStackKind a) (tsStackKind b)
}
-- | Run forward dataflow analysis over a CFG.
@@ -683,6 +703,7 @@ initSummary = FuncSummary $ TaintState
, tsProv = Map.fromList [ (r, ProvUnknown) | r <- callerSavedRegs ]
, tsStackProv = IM.empty
, tsKind = Map.fromList [ (r, KindUnknown) | r <- callerSavedRegs ]
+ , tsStackKind = IM.empty
}
-- | Caller-saved registers per AArch64 ABI.
diff --git a/plans/ARCH14.md b/plans/ARCH14.md
@@ -0,0 +1,26 @@
+# ARCH14: Stack Spill Kind Tracking (Stage 2)
+
+## Goal
+
+Preserve pointer vs scalar kind information across spills and reloads by
+tracking kind for stack slots.
+
+## Scope
+
+- Extend taint state with `tsStackKind :: IntMap RegKind`.
+- Update stack store/load logic to propagate kind.
+
+## Kind Propagation Rules (Stage 2)
+
+- Store to `[sp, #imm]`: record source register kind in stack kind map.
+- Load from `[sp, #imm]`: restore kind from stack map; default Unknown.
+- Any SP adjustment clears stack kind map (same as taint/provenance).
+
+## Address Checks
+
+- Base/index checks remain pointer-aware using `RegKind` for registers.
+
+## Risks
+
+- Stack offset tracking is approximate; incorrect SP adjustments can
+ lead to stale kinds. The existing SP-clear heuristic helps.
diff --git a/plans/IMPL14.md b/plans/IMPL14.md
@@ -0,0 +1,33 @@
+# IMPL14: Implement Stack Kind Tracking (Stage 2)
+
+## Summary
+
+Extend taint state with stack slot kinds so pointer/scalar classification
+survives spills and reloads.
+
+## Steps
+
+1) Extend TaintState
+- Add `tsStackKind :: IntMap RegKind`.
+- Initialize empty; clear on SP modification.
+
+2) Stack store/load updates
+- `storeToStack`/`storePairToStack`: record source kinds.
+- `loadFromStack`/`loadPairFromStack`: restore kinds for destinations.
+
+3) Clear on SP changes
+- Ensure `clearStackMap` clears stack kinds too.
+
+4) Tests
+- Add a test where a pointer is spilled and reloaded; ensure kind is
+ preserved.
+- Add a test where a scalar spill does not upgrade to pointer.
+
+## Files to Touch
+
+- `lib/Audit/AArch64/Taint.hs`
+- `test/`
+
+## Validation
+
+- `cabal test`
diff --git a/test/Main.hs b/test/Main.hs
@@ -366,6 +366,43 @@ auditTests = testGroup "Audit" [
other -> assertFailure $ "wrong reason: " ++ show other
_ -> assertFailure "expected one violation"
+ , testCase "good: pointer spill/reload preserves kind" $ do
+ -- A pointer (from adrp) spilled and reloaded should stay KindPtr
+ -- and be usable as a base via provenance upgrade
+ let src = T.unlines
+ [ "foo:"
+ , " adrp x8, _const@PAGE" -- x8 = Public, ProvPublic, KindPtr
+ , " str x8, [sp, #16]" -- Spill pointer to stack
+ , " ldr x9, [sp, #16]" -- Reload: Unknown, ProvPublic, KindPtr
+ , " ldr x0, [x9]" -- Use as base - allowed via provenance
+ , " ret"
+ ]
+ case audit "test" src of
+ Left e -> assertFailure $ "parse failed: " ++ show e
+ Right ar -> assertEqual "no violations" 0 (length (arViolations ar))
+
+ , testCase "bad: scalar spill/reload doesn't upgrade kind" $ do
+ -- A scalar (loaded value) spilled and reloaded should stay KindScalar
+ -- even with public provenance, should not be usable as base
+ let src = T.unlines
+ [ "foo:"
+ , " ldr x8, [x20]" -- x8 = Unknown, ProvUnknown, KindScalar
+ , " add x8, x20, x8" -- x8 = Unknown, ProvPublic, KindScalar
+ , " str x8, [sp, #16]" -- Spill scalar to stack
+ , " ldr x9, [sp, #16]" -- Reload: Unknown, ProvPublic, KindScalar
+ , " ldr x0, [x9]" -- Use as base - violation!
+ , " ret"
+ ]
+ case audit "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
+ UnknownBase X9 -> pure ()
+ other -> assertFailure $ "wrong reason: " ++ show other
+ _ -> assertFailure "expected one violation"
+
, testCase "good: GOT/PLT pattern" $ do
-- GOTPAGEOFF pattern: adrp + ldr from GOT should produce public
let src = T.unlines
