auditor

commit c71f0b539b1e57c0c618fceafd2a506dcbd86ec9
parent 9945a6a7a5d020dfd44f8ba4e830ec96e877ed4f
Author: Jared Tobin <jared@jtobin.io>
Date:   Fri, 27 Feb 2026 15:10:12 +0400

fix: address review feedback for tail call propagation

- Don't apply summaries for in-file tail calls; let CFG edges propagate
  state. This preserves x0-x7 argument taint across tail calls.
- Use enclosing function label for violation symbols instead of block
  label (fixes attribution for local labels like Lc*).
- Add blockFunction helper to CFG.hs for block->function lookup.
- Add test case verifying x0 taint preserved across in-file tail calls.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Diffstat:
M
lib/Audit/AArch64/CFG.hs
 | 
9
+++++++++
M
lib/Audit/AArch64/Check.hs
 | 
11
+++++++----
M
lib/Audit/AArch64/Taint.hs
 | 
6
++++--
M
test/Main.hs
 | 
18
++++++++++++++++++

4 files changed, 38 insertions(+), 6 deletions(-)
diff --git a/lib/Audit/AArch64/CFG.hs b/lib/Audit/AArch64/CFG.hs
@@ -19,6 +19,7 @@ module Audit.AArch64.CFG (
   , blockSuccessors
   , cfgBlockCount
   , indexBlock
+  , blockFunction
     -- * Function partitioning
   , isFunctionLabel
   , functionBlocks
@@ -34,6 +35,7 @@ import Audit.AArch64.Types
 import Control.DeepSeq (NFData)
 import Data.List (foldl')
 import Data.Map.Strict (Map)
+import Data.Maybe (listToMaybe)
 import qualified Data.Map.Strict as Map
 import Data.Primitive.Array (Array)
 import qualified Data.Primitive.Array as A
@@ -76,6 +78,13 @@ cfgBlockCount cfg = A.sizeofArray (cfgBlocks cfg)
 indexBlock :: CFG -> Int -> BasicBlock
 indexBlock cfg i = A.indexArray (cfgBlocks cfg) i
 
+-- | Find the enclosing function label for a block index.
+-- Returns Nothing if the block doesn't belong to any function.
+blockFunction :: CFG -> Int -> Maybe Text
+blockFunction cfg idx =
+  listToMaybe [ func | (func, idxs) <- Map.toList (cfgFuncBlocks cfg)
+                     , idx `elem` idxs ]
+
 -- | Build a CFG from parsed assembly lines.
 buildCFG :: [Line] -> CFG
 buildCFG lns = cfg
diff --git a/lib/Audit/AArch64/Check.hs b/lib/Audit/AArch64/Check.hs
@@ -26,7 +26,7 @@ module Audit.AArch64.Check (
   ) where
 
 import Audit.AArch64.CFG (BasicBlock(..), CFG(..), cfgBlockCount, indexBlock,
-                          functionLabels, functionBlocks)
+                          blockFunction)
 import Audit.AArch64.Taint
 import Audit.AArch64.Types
   ( Reg(..), Instr(..), Line(..), AddrMode(..)
@@ -214,7 +214,8 @@ checkCFGInterProc sym cfg =
         [ fst (checkBlockWithSummary blockSym summaries inState (bbLines bb))
         | idx <- [0..nBlocks-1]
         , let bb = indexBlock cfg idx
-              blockSym = fromMaybe sym (bbLabel bb)
+              -- Use enclosing function label, not block label
+              blockSym = fromMaybe sym (blockFunction cfg idx)
               inState = IM.findWithDefault initTaintState idx inStates
         ]
 
@@ -246,7 +247,8 @@ checkCFGWithConfig tcfg sym cfg =
         [ fst (checkBlock blockSym inState (bbLines bb))
         | idx <- [0..nBlocks-1]
         , let bb = indexBlock cfg idx
-              blockSym = fromMaybe sym (bbLabel bb)
+              -- Use enclosing function label, not block label
+              blockSym = fromMaybe sym (blockFunction cfg idx)
               inState = IM.findWithDefault initTaintState idx inStates
         ]
 
@@ -263,6 +265,7 @@ checkCFGInterProcWithConfig tcfg sym cfg =
         [ fst (checkBlockWithSummary blockSym summaries inState (bbLines bb))
         | idx <- [0..nBlocks-1]
         , let bb = indexBlock cfg idx
-              blockSym = fromMaybe sym (bbLabel bb)
+              -- Use enclosing function label, not block label
+              blockSym = fromMaybe sym (blockFunction cfg idx)
               inState = IM.findWithDefault initTaintState idx inStates
         ]
diff --git a/lib/Audit/AArch64/Taint.hs b/lib/Audit/AArch64/Taint.hs
@@ -1224,6 +1224,8 @@ analyzeBlockWithSummaries bb st0 summaries = foldl' go st0 (bbLines bb)
       Just instr -> transferWithSummary instr st summaries
 
 -- | Transfer with summary application for calls and tail calls.
+-- For in-file tail calls, state flows via CFG edges - no summary application.
+-- For external tail calls, we conservatively invalidate STG arg registers.
 transferWithSummary :: Instr -> TaintState -> Map Text FuncSummary -> TaintState
 transferWithSummary instr st summaries = case instr of
   Bl target ->
@@ -1235,8 +1237,8 @@ transferWithSummary instr st summaries = case instr of
   B target
     | isFunctionLabel target ->
         case Map.lookup target summaries of
-          Just summ -> applySummary summ st
-          Nothing   -> invalidateStgArgRegs st
+          Just _  -> st  -- In-file: CFG edge propagates state
+          Nothing -> invalidateStgArgRegs st  -- External: conservative
   -- Indirect jumps: conservative treatment
   Br _ -> invalidateStgArgRegs st
   _ -> transfer instr st
diff --git a/test/Main.hs b/test/Main.hs
@@ -1656,4 +1656,22 @@ tailCallTests = testGroup "TailCall" [
           assertEqual "secret violation" 1
             (length $ filter isSecretViolation (map vReason (arViolations ar)))
 
+  -- ABI argument registers preserved across in-file tail calls
+  , testCase "tail call preserves x0 argument taint" $ do
+      let src = T.unlines
+            [ "_caller:"
+            , "  b _callee_info"
+            , "_callee_info:"
+            , "  ldr x1, [x20, x0]"
+            , "  ret"
+            ]
+          cfg = TaintConfig (Map.singleton "_caller"
+                  (ArgPolicy (Set.singleton X0) Set.empty Set.empty
+                    Set.empty Set.empty)) True
+      case auditInterProcWithConfig cfg "test" src of
+        Left e -> assertFailure $ "parse failed: " ++ show e
+        Right ar ->
+          assertEqual "x0 secret preserved" 1
+            (length $ filter isSecretViolation (map vReason (arViolations ar)))
+
   ]