Taint.hs (57262B)
1 {-# OPTIONS_HADDOCK prune #-} 2 {-# LANGUAGE OverloadedStrings #-} 3 4 -- | 5 -- Module: Audit.AArch64.Taint 6 -- Copyright: (c) 2025 Jared Tobin 7 -- License: MIT 8 -- Maintainer: jared@ppad.tech 9 -- 10 -- Taint analysis for AArch64 registers. 11 -- 12 -- Implements forward dataflow analysis to track register "publicness". 13 -- Public registers are those derived from known-safe sources (stack 14 -- pointers, heap pointers, constants). Registers with unknown or 15 -- secret-derived values are flagged when used in memory addressing. 16 17 module Audit.AArch64.Taint ( 18 TaintState(..) 19 , initTaintState 20 , initTaintStateWith 21 , analyzeLine 22 , analyzeBlock 23 , getTaint 24 , setTaint 25 , getProvenance 26 , getKind 27 , joinTaintState 28 , runDataflow 29 , runDataflowWithConfig 30 , runDataflowWithConfigAndSummaries 31 -- * Function summaries 32 , FuncSummary(..) 33 , initSummary 34 , joinSummary 35 , applySummary 36 , runFunctionDataflow 37 , runFunctionBlocks 38 , runInterProc 39 , runInterProcWithConfig 40 , runFunctionBlocksWithEntry 41 -- * Entry taint seeding 42 , seedArgs 43 ) where 44 45 import Audit.AArch64.CFG 46 ( BasicBlock(..), CFG(..) 47 , blockSuccessors, functionLabels, functionBlocks 48 , buildFunctionBlocksMap, buildCallerMap 49 , cfgBlockCount, indexBlock 50 , isFunctionLabel 51 ) 52 import Audit.AArch64.Runtime 53 (RuntimeConfig(..), SecondaryStack(..)) 54 import Audit.AArch64.Types 55 ( Reg(..), Instr(..), Line(..), Operand(..), AddrMode(..) 56 , Taint(..), joinTaint, Provenance(..), joinProvenance 57 , RegKind(..), joinKind 58 , TaintConfig(..), ArgPolicy(..) 59 ) 60 import Data.IntMap.Strict (IntMap) 61 import qualified Data.IntMap.Strict as IM 62 import Data.IntSet (IntSet) 63 import qualified Data.IntSet as IS 64 import Data.Map.Strict (Map) 65 import qualified Data.Map.Strict as Map 66 import qualified Data.Set as Set 67 import Data.Text (Text) 68 69 -- | Taint state: maps registers to their publicness, plus stack slots. 70 -- Also tracks provenance for upgrading Unknown to Public when provable, 71 -- and register kinds (pointer vs scalar) for safe provenance upgrades. 72 -- Tracks both hardware stack (SP) and secondary stack slots separately. 73 -- A coarse heap bucket captures taint for non-stack memory accesses. 74 -- A refined heap slot map tracks [base, #imm] accesses for public pointers. 75 data TaintState = TaintState 76 { tsRegs :: !(Map Reg Taint) 77 -- ^ Register taints 78 , tsStack :: !(IntMap Taint) 79 -- ^ Stack slot taints (keyed by SP offset) 80 , tsProv :: !(Map Reg Provenance) 81 -- ^ Register provenance 82 , tsStackProv :: !(IntMap Provenance) 83 -- ^ Stack slot provenance (keyed by SP offset) 84 , tsKind :: !(Map Reg RegKind) 85 -- ^ Register kinds (pointer vs scalar) 86 , tsStackKind :: !(IntMap RegKind) 87 -- ^ Stack slot kinds (keyed by SP offset) 88 , tsStgStack :: !(IntMap Taint) 89 -- ^ Secondary stack slot taints (keyed by base reg offset) 90 , tsStgStackProv :: !(IntMap Provenance) 91 -- ^ Secondary stack slot provenance 92 , tsStgStackKind :: !(IntMap RegKind) 93 -- ^ Secondary stack slot kinds 94 , tsHeapTaint :: !Taint 95 -- ^ Coarse heap taint bucket (joined from all non-stack stores) 96 , tsHeapProv :: !Provenance 97 -- ^ Coarse heap provenance bucket 98 , tsHeapKind :: !RegKind 99 -- ^ Coarse heap kind bucket 100 , tsHeapSlots :: !(Map (Reg, Int) (Taint, Provenance, RegKind)) 101 -- ^ Refined heap slots keyed by (base register, offset) 102 , tsAssumeStgPublic :: !Bool 103 -- ^ Assume untracked secondary stack slots are public pointers 104 } deriving (Eq, Show) 105 106 -- | Check if a register is a secondary stack base register. 107 isSecStackBase :: RuntimeConfig -> Reg -> Bool 108 isSecStackBase rt r = case rtSecondaryStack rt of 109 Nothing -> False 110 Just ss -> ssBaseReg ss == r 111 112 -- | Check if a register is in the public roots. 113 isPublicRoot :: RuntimeConfig -> Reg -> Bool 114 isPublicRoot rt r = r `elem` rtPublicRoots rt 115 116 -- | Empty taint state with configurable secondary stack assumption. 117 emptyTaintStateWith :: Bool -> TaintState 118 emptyTaintStateWith assumeStgPublic = TaintState 119 { tsRegs = Map.empty 120 , tsStack = IM.empty 121 , tsProv = Map.empty 122 , tsStackProv = IM.empty 123 , tsKind = Map.empty 124 , tsStackKind = IM.empty 125 , tsStgStack = IM.empty 126 , tsStgStackProv = IM.empty 127 , tsStgStackKind = IM.empty 128 , tsHeapTaint = Unknown 129 , tsHeapProv = ProvUnknown 130 , tsHeapKind = KindUnknown 131 , tsHeapSlots = Map.empty 132 , tsAssumeStgPublic = assumeStgPublic 133 } 134 135 -- | Initial taint state with public roots marked. 136 -- Public roots are marked as Ptr kind (they are pointers). 137 -- Heap bucket starts Unknown (conservative for untracked memory). 138 initTaintState :: RuntimeConfig -> TaintState 139 initTaintState rt = initTaintStateWith rt assumeStg 140 where 141 assumeStg = case rtSecondaryStack rt of 142 Nothing -> False 143 Just ss -> ssAssumePublic ss 144 145 -- | Initial taint state with configurable secondary stack assumption. 146 initTaintStateWith :: RuntimeConfig -> Bool -> TaintState 147 initTaintStateWith rt assumeStgPublic = 148 let roots = rtPublicRoots rt 149 in TaintState 150 { tsRegs = Map.fromList 151 [(r, Public) | r <- roots] 152 , tsStack = IM.empty 153 , tsProv = Map.fromList 154 [(r, ProvPublic) | r <- roots] 155 , tsStackProv = IM.empty 156 , tsKind = Map.fromList 157 [(r, KindPtr) | r <- roots] 158 , tsStackKind = IM.empty 159 , tsStgStack = IM.empty 160 , tsStgStackProv = IM.empty 161 , tsStgStackKind = IM.empty 162 , tsHeapTaint = Unknown 163 , tsHeapProv = ProvUnknown 164 , tsHeapKind = KindUnknown 165 , tsHeapSlots = Map.empty 166 , tsAssumeStgPublic = assumeStgPublic 167 } 168 169 -- | Seed argument registers and secondary stack slots 170 -- according to policy. 171 seedArgs :: RuntimeConfig -> ArgPolicy -> TaintState 172 -> TaintState 173 seedArgs rt policy st = 174 let -- Apply public first, then secret pointee, then secret 175 st1 = Set.foldr markPublic st (apPublic policy) 176 st2 = Set.foldr markSecretPointee st1 177 (apSecretPointee policy) 178 st3 = Set.foldr markSecret st2 (apSecret policy) 179 -- Seed secondary stack slots (conditional) 180 st4 = case rtSecondaryStack rt of 181 Nothing -> st3 182 Just _ -> 183 let st4' = Set.foldr markStgPublic st3 184 (apStgPublic policy) 185 in Set.foldr markStgSecret st4' 186 (apStgSecret policy) 187 in st4 188 where 189 markPublic r s = s 190 { tsRegs = Map.insert r Public (tsRegs s) 191 , tsProv = Map.insert r ProvPublic (tsProv s) 192 } 193 markSecretPointee r s = s 194 { tsRegs = Map.insert r Public (tsRegs s) 195 , tsProv = Map.insert r ProvSecretData (tsProv s) 196 , tsKind = Map.insert r KindPtr (tsKind s) 197 } 198 markSecret r s = s 199 { tsRegs = Map.insert r Secret (tsRegs s) 200 , tsProv = Map.insert r ProvUnknown (tsProv s) 201 } 202 markStgPublic off s = s 203 { tsStgStack = 204 IM.insert off Public (tsStgStack s) 205 , tsStgStackProv = 206 IM.insert off ProvPublic (tsStgStackProv s) 207 , tsStgStackKind = 208 IM.insert off KindScalar (tsStgStackKind s) 209 } 210 markStgSecret off s = s 211 { tsStgStack = 212 IM.insert off Secret (tsStgStack s) 213 , tsStgStackProv = 214 IM.insert off ProvUnknown (tsStgStackProv s) 215 , tsStgStackKind = 216 IM.insert off KindScalar (tsStgStackKind s) 217 } 218 219 -- | Get the taint of a register. 220 getTaint :: Reg -> TaintState -> Taint 221 getTaint r st = Map.findWithDefault Unknown r (tsRegs st) 222 223 -- | Get the provenance of a register. 224 getProvenance :: Reg -> TaintState -> Provenance 225 getProvenance r st = 226 Map.findWithDefault ProvUnknown r (tsProv st) 227 228 -- | Get the kind of a register. 229 getKind :: Reg -> TaintState -> RegKind 230 getKind r st = 231 Map.findWithDefault KindUnknown r (tsKind st) 232 233 -- | Analyze a single line, updating taint state. 234 analyzeLine :: RuntimeConfig -> Line -> TaintState 235 -> TaintState 236 analyzeLine rt l st = case lineInstr l of 237 Nothing -> st 238 Just instr -> transfer rt instr st 239 240 -- | Analyze a basic block, threading taint state through. 241 analyzeBlock :: RuntimeConfig -> [Line] -> TaintState 242 -> TaintState 243 analyzeBlock rt lns st = 244 foldl' (flip (analyzeLine rt)) st lns 245 246 -- | Transfer function for taint analysis. 247 transfer :: RuntimeConfig -> Instr -> TaintState 248 -> TaintState 249 transfer rt instr st = case instr of 250 -- Move: destination gets source taint, provenance, and kind 251 Mov dst op -> 252 setTaintProvKind dst (operandTaint op st) 253 (operandProv op st) (operandKind op st) st 254 Movz dst _ _ -> 255 setTaintProvKind dst Public ProvPublic KindScalar st 256 Movk _ _ _ -> st 257 Movn dst _ _ -> 258 setTaintProvKind dst Public ProvPublic KindScalar st 259 260 -- Arithmetic 261 Add dst r1 op -> 262 let t = join2 (getTaint r1 st) (operandTaint op st) 263 p = provJoin2 (getProvenance r1 st) 264 (operandProv op st) 265 k = pointerArithKind r1 op st 266 in updateWithSecStackShiftAdd rt dst r1 op t p k st 267 Sub dst r1 op -> 268 let t = join2 (getTaint r1 st) (operandTaint op st) 269 p = provJoin2 (getProvenance r1 st) 270 (operandProv op st) 271 k = pointerArithKind r1 op st 272 in updateWithSecStackShiftSub rt dst r1 op t p k st 273 Adds dst r1 op -> 274 let t = join2 (getTaint r1 st) (operandTaint op st) 275 p = provJoin2 (getProvenance r1 st) 276 (operandProv op st) 277 k = pointerArithKind r1 op st 278 in updateWithSecStackShiftAdd rt dst r1 op t p k st 279 Subs dst r1 op -> 280 let t = join2 (getTaint r1 st) (operandTaint op st) 281 p = provJoin2 (getProvenance r1 st) 282 (operandProv op st) 283 k = pointerArithKind r1 op st 284 in updateWithSecStackShiftSub rt dst r1 op t p k st 285 Adc dst r1 r2 -> 286 let t = join2 (getTaint r1 st) (getTaint r2 st) 287 p = provJoin2 (getProvenance r1 st) 288 (getProvenance r2 st) 289 in setTaintProvKind dst t p KindScalar st 290 Adcs dst r1 r2 -> 291 let t = join2 (getTaint r1 st) (getTaint r2 st) 292 p = provJoin2 (getProvenance r1 st) 293 (getProvenance r2 st) 294 in setTaintProvKind dst t p KindScalar st 295 Sbc dst r1 r2 -> 296 let t = join2 (getTaint r1 st) (getTaint r2 st) 297 p = provJoin2 (getProvenance r1 st) 298 (getProvenance r2 st) 299 in setTaintProvKind dst t p KindScalar st 300 Neg dst op -> 301 setTaintProvKind dst (operandTaint op st) 302 (operandProv op st) KindScalar st 303 Negs dst op -> 304 setTaintProvKind dst (operandTaint op st) 305 (operandProv op st) KindScalar st 306 Mul dst r1 r2 -> 307 let t = join2 (getTaint r1 st) (getTaint r2 st) 308 p = provJoin2 (getProvenance r1 st) 309 (getProvenance r2 st) 310 in setTaintProvKind dst t p KindScalar st 311 Mneg dst r1 r2 -> 312 let t = join2 (getTaint r1 st) (getTaint r2 st) 313 p = provJoin2 (getProvenance r1 st) 314 (getProvenance r2 st) 315 in setTaintProvKind dst t p KindScalar st 316 Madd dst r1 r2 r3 -> 317 let t = join3 (getTaint r1 st) (getTaint r2 st) 318 (getTaint r3 st) 319 p = provJoin3 (getProvenance r1 st) 320 (getProvenance r2 st) 321 (getProvenance r3 st) 322 in setTaintProvKind dst t p KindScalar st 323 Msub dst r1 r2 r3 -> 324 let t = join3 (getTaint r1 st) (getTaint r2 st) 325 (getTaint r3 st) 326 p = provJoin3 (getProvenance r1 st) 327 (getProvenance r2 st) 328 (getProvenance r3 st) 329 in setTaintProvKind dst t p KindScalar st 330 Umulh dst r1 r2 -> 331 let t = join2 (getTaint r1 st) (getTaint r2 st) 332 p = provJoin2 (getProvenance r1 st) 333 (getProvenance r2 st) 334 in setTaintProvKind dst t p KindScalar st 335 Smulh dst r1 r2 -> 336 let t = join2 (getTaint r1 st) (getTaint r2 st) 337 p = provJoin2 (getProvenance r1 st) 338 (getProvenance r2 st) 339 in setTaintProvKind dst t p KindScalar st 340 Udiv dst r1 r2 -> 341 let t = join2 (getTaint r1 st) (getTaint r2 st) 342 p = provJoin2 (getProvenance r1 st) 343 (getProvenance r2 st) 344 in setTaintProvKind dst t p KindScalar st 345 Sdiv dst r1 r2 -> 346 let t = join2 (getTaint r1 st) (getTaint r2 st) 347 p = provJoin2 (getProvenance r1 st) 348 (getProvenance r2 st) 349 in setTaintProvKind dst t p KindScalar st 350 351 -- Logical: special case for pointer untagging 352 And dst r1 op -> 353 let t = join2 (getTaint r1 st) (operandTaint op st) 354 p = provJoin2 (getProvenance r1 st) 355 (operandProv op st) 356 srcProv = getProvenance r1 st 357 srcKind = getKind r1 st 358 isUntag = isPointerUntagMask rt op 359 p' = if isUntag && srcProv == ProvPublic 360 then ProvPublic else p 361 k = if isUntag && srcKind == KindPtr 362 then KindPtr else KindScalar 363 in setTaintProvKind dst t p' k st 364 Orr dst r1 op -> 365 let t = join2 (getTaint r1 st) (operandTaint op st) 366 p = provJoin2 (getProvenance r1 st) 367 (operandProv op st) 368 in setTaintProvKind dst t p KindScalar st 369 Eor dst r1 op -> 370 let t = join2 (getTaint r1 st) (operandTaint op st) 371 p = provJoin2 (getProvenance r1 st) 372 (operandProv op st) 373 in setTaintProvKind dst t p KindScalar st 374 Bic dst r1 op -> 375 let t = join2 (getTaint r1 st) (operandTaint op st) 376 p = provJoin2 (getProvenance r1 st) 377 (operandProv op st) 378 in setTaintProvKind dst t p KindScalar st 379 Mvn dst op -> 380 setTaintProvKind dst (operandTaint op st) 381 (operandProv op st) KindScalar st 382 Tst _ _ -> st 383 384 -- Shifts: result is scalar 385 Lsl dst r1 op -> 386 let t = join2 (getTaint r1 st) (operandTaint op st) 387 p = provJoin2 (getProvenance r1 st) 388 (operandProv op st) 389 in setTaintProvKind dst t p KindScalar st 390 Lsr dst r1 op -> 391 let t = join2 (getTaint r1 st) (operandTaint op st) 392 p = provJoin2 (getProvenance r1 st) 393 (operandProv op st) 394 in setTaintProvKind dst t p KindScalar st 395 Asr dst r1 op -> 396 let t = join2 (getTaint r1 st) (operandTaint op st) 397 p = provJoin2 (getProvenance r1 st) 398 (operandProv op st) 399 in setTaintProvKind dst t p KindScalar st 400 Ror dst r1 op -> 401 let t = join2 (getTaint r1 st) (operandTaint op st) 402 p = provJoin2 (getProvenance r1 st) 403 (operandProv op st) 404 in setTaintProvKind dst t p KindScalar st 405 406 -- Bit manipulation: result is scalar 407 Ubfx dst r1 _ _ -> 408 setTaintProvKind dst (getTaint r1 st) 409 (getProvenance r1 st) KindScalar st 410 Sbfx dst r1 _ _ -> 411 setTaintProvKind dst (getTaint r1 st) 412 (getProvenance r1 st) KindScalar st 413 Bfi dst r1 _ _ -> 414 let t = join2 (getTaint dst st) (getTaint r1 st) 415 p = provJoin2 (getProvenance dst st) 416 (getProvenance r1 st) 417 in setTaintProvKind dst t p KindScalar st 418 Extr dst r1 r2 _ -> 419 let t = join2 (getTaint r1 st) (getTaint r2 st) 420 p = provJoin2 (getProvenance r1 st) 421 (getProvenance r2 st) 422 in setTaintProvKind dst t p KindScalar st 423 424 -- Address generation 425 Adr dst _ -> 426 setTaintProvKind dst Public ProvPublic KindPtr st 427 Adrp dst _ -> 428 setTaintProvKind dst Public ProvPublic KindPtr st 429 430 -- Loads 431 Ldr dst addr -> loadFromStack rt dst addr st 432 Ldrb dst addr -> loadFromStack rt dst addr st 433 Ldrh dst addr -> loadFromStack rt dst addr st 434 Ldrsb dst addr -> loadFromStack rt dst addr st 435 Ldrsh dst addr -> loadFromStack rt dst addr st 436 Ldrsw dst addr -> loadFromStack rt dst addr st 437 Ldur dst addr -> loadFromStack rt dst addr st 438 Ldp dst1 dst2 addr -> 439 loadPairFromStack rt dst1 dst2 addr st 440 441 -- Stores 442 Str src addr -> storeToStack rt src addr st 443 Strb src addr -> storeToStack rt src addr st 444 Strh src addr -> storeToStack rt src addr st 445 Stur src addr -> storeToStack rt src addr st 446 Stp src1 src2 addr -> 447 storePairToStack rt src1 src2 addr st 448 449 -- Acquire/release loads 450 Ldar dst addr -> loadFromStack rt dst addr st 451 Ldarb dst addr -> loadFromStack rt dst addr st 452 Ldarh dst addr -> loadFromStack rt dst addr st 453 454 -- Release stores 455 Stlr src addr -> storeToStack rt src addr st 456 Stlrb src addr -> storeToStack rt src addr st 457 Stlrh src addr -> storeToStack rt src addr st 458 459 -- Exclusive loads 460 Ldxr dst addr -> loadFromStack rt dst addr st 461 Ldxrb dst addr -> loadFromStack rt dst addr st 462 Ldxrh dst addr -> loadFromStack rt dst addr st 463 464 -- Exclusive stores 465 Stxr status src addr -> 466 let st' = storeToStack rt src addr st 467 in setTaintProvKind status Public ProvPublic 468 KindScalar st' 469 Stxrb status src addr -> 470 let st' = storeToStack rt src addr st 471 in setTaintProvKind status Public ProvPublic 472 KindScalar st' 473 Stxrh status src addr -> 474 let st' = storeToStack rt src addr st 475 in setTaintProvKind status Public ProvPublic 476 KindScalar st' 477 478 -- Acquire-exclusive loads 479 Ldaxr dst addr -> loadFromStack rt dst addr st 480 Ldaxrb dst addr -> loadFromStack rt dst addr st 481 Ldaxrh dst addr -> loadFromStack rt dst addr st 482 483 -- Release-exclusive stores 484 Stlxr status src addr -> 485 let st' = storeToStack rt src addr st 486 in setTaintProvKind status Public ProvPublic 487 KindScalar st' 488 Stlxrb status src addr -> 489 let st' = storeToStack rt src addr st 490 in setTaintProvKind status Public ProvPublic 491 KindScalar st' 492 Stlxrh status src addr -> 493 let st' = storeToStack rt src addr st 494 in setTaintProvKind status Public ProvPublic 495 KindScalar st' 496 497 -- Conditionals 498 Cmp _ _ -> st 499 Cmn _ _ -> st 500 Csel dst r1 r2 _ -> 501 let t = join2 (getTaint r1 st) (getTaint r2 st) 502 p = provJoin2 (getProvenance r1 st) 503 (getProvenance r2 st) 504 k = joinKind (getKind r1 st) (getKind r2 st) 505 in setTaintProvKind dst t p k st 506 Csinc dst r1 r2 _ -> 507 let t = join2 (getTaint r1 st) (getTaint r2 st) 508 p = provJoin2 (getProvenance r1 st) 509 (getProvenance r2 st) 510 in setTaintProvKind dst t p KindScalar st 511 Csinv dst r1 r2 _ -> 512 let t = join2 (getTaint r1 st) (getTaint r2 st) 513 p = provJoin2 (getProvenance r1 st) 514 (getProvenance r2 st) 515 in setTaintProvKind dst t p KindScalar st 516 Csneg dst r1 r2 _ -> 517 let t = join2 (getTaint r1 st) (getTaint r2 st) 518 p = provJoin2 (getProvenance r1 st) 519 (getProvenance r2 st) 520 in setTaintProvKind dst t p KindScalar st 521 Cset dst _ -> 522 setTaintProvKind dst Public ProvPublic KindScalar st 523 Cinc dst r1 _ -> 524 setTaintProvKind dst (getTaint r1 st) 525 (getProvenance r1 st) KindScalar st 526 527 -- Branches 528 B _ -> st 529 BCond _ _ -> st 530 Bl _ -> invalidateCallerSaved st 531 Blr _ -> invalidateCallerSaved st 532 Br _ -> st 533 Ret _ -> st 534 Cbz _ _ -> st 535 Cbnz _ _ -> st 536 Tbz _ _ _ -> st 537 Tbnz _ _ _ -> st 538 539 -- System 540 Nop -> st 541 Svc _ -> invalidateCallerSaved st 542 543 -- Unknown instruction 544 Other _ _ -> st 545 546 -- | Set taint for a register. 547 setTaint :: Reg -> Taint -> TaintState -> TaintState 548 setTaint r t st = st { tsRegs = Map.insert r t (tsRegs st) } 549 550 -- | Set both taint and provenance for a register. 551 setTaintProv :: Reg -> Taint -> Provenance -> TaintState 552 -> TaintState 553 setTaintProv r t p st = st 554 { tsRegs = Map.insert r t (tsRegs st) 555 , tsProv = Map.insert r p (tsProv st) 556 } 557 558 -- | Set taint, provenance, and kind for a register. 559 setTaintProvKind :: Reg -> Taint -> Provenance -> RegKind 560 -> TaintState -> TaintState 561 setTaintProvKind r t p k st = st 562 { tsRegs = Map.insert r t (tsRegs st) 563 , tsProv = Map.insert r p (tsProv st) 564 , tsKind = Map.insert r k (tsKind st) 565 , tsHeapSlots = Map.filterWithKey 566 (\(base, _) _ -> base /= r) (tsHeapSlots st) 567 } 568 569 -- | Set taint for a loaded value from the heap bucket. 570 setTaintLoadHeap :: RuntimeConfig -> Reg -> TaintState 571 -> TaintState 572 setTaintLoadHeap rt dst st 573 | isPublicRoot rt dst = 574 setTaintProvKind dst Public ProvPublic KindPtr st 575 | otherwise = 576 setTaintProvKind dst (tsHeapTaint st) (tsHeapProv st) 577 (tsHeapKind st) st 578 579 -- | Set taint for a loaded value from a refined heap slot. 580 setTaintLoadHeapSlot :: RuntimeConfig -> Reg -> Reg -> Int 581 -> TaintState -> TaintState 582 setTaintLoadHeapSlot rt dst base off st 583 | isPublicRoot rt dst = 584 setTaintProvKind dst Public ProvPublic KindPtr st 585 | otherwise = case getHeapSlot base off st of 586 Just (t, p, k) -> setTaintProvKind dst t p k st 587 Nothing -> setTaintProvKind dst (tsHeapTaint st) 588 (tsHeapProv st) (tsHeapKind st) st 589 590 -- | Set taint for a loaded value from a known stack slot. 591 setTaintLoadStack :: RuntimeConfig -> Reg -> Int 592 -> TaintState -> TaintState 593 setTaintLoadStack rt dst offset st 594 | isPublicRoot rt dst = 595 setTaintProvKind dst Public ProvPublic KindPtr st 596 | otherwise = 597 let taint = IM.findWithDefault Unknown offset 598 (tsStack st) 599 prov = IM.findWithDefault ProvUnknown offset 600 (tsStackProv st) 601 kind = IM.findWithDefault KindUnknown offset 602 (tsStackKind st) 603 in setTaintProvKind dst taint prov kind st 604 605 -- | Store taint, provenance, and kind to a stack slot. 606 setStackTaintProvKind :: Int -> Taint -> Provenance 607 -> RegKind -> TaintState -> TaintState 608 setStackTaintProvKind offset t p k st = st 609 { tsStack = IM.insert offset t (tsStack st) 610 , tsStackProv = IM.insert offset p (tsStackProv st) 611 , tsStackKind = IM.insert offset k (tsStackKind st) 612 } 613 614 -- | Clear all stack slot taints (when SP is modified). 615 clearStackMap :: TaintState -> TaintState 616 clearStackMap st = st 617 { tsStack = IM.empty 618 , tsStackProv = IM.empty 619 , tsStackKind = IM.empty 620 } 621 622 -- | Store taint to a secondary stack slot. 623 setSecStackTaintProvKind :: Int -> Taint -> Provenance 624 -> RegKind -> TaintState 625 -> TaintState 626 setSecStackTaintProvKind offset t p k st = st 627 { tsStgStack = IM.insert offset t (tsStgStack st) 628 , tsStgStackProv = IM.insert offset p 629 (tsStgStackProv st) 630 , tsStgStackKind = IM.insert offset k 631 (tsStgStackKind st) 632 } 633 634 -- | Clear all secondary stack slot taints. 635 clearSecStackMap :: TaintState -> TaintState 636 clearSecStackMap st = st 637 { tsStgStack = IM.empty 638 , tsStgStackProv = IM.empty 639 , tsStgStackKind = IM.empty 640 } 641 642 -- | Shift all secondary stack slot offsets by a signed delta. 643 shiftSecStackMap :: Int -> TaintState -> TaintState 644 shiftSecStackMap delta st 645 | delta == 0 = st 646 | otherwise = st 647 { tsStgStack = shiftIM (tsStgStack st) 648 , tsStgStackProv = shiftIM (tsStgStackProv st) 649 , tsStgStackKind = shiftIM (tsStgStackKind st) 650 } 651 where 652 shiftIM im = 653 IM.fromList [(k + delta, v) | (k, v) <- IM.toList im] 654 655 -- | Join taint into the heap bucket. 656 joinHeapBucket :: Taint -> Provenance -> RegKind 657 -> TaintState -> TaintState 658 joinHeapBucket t p k st = st 659 { tsHeapTaint = joinTaint (tsHeapTaint st) t 660 , tsHeapProv = joinProvenance (tsHeapProv st) p 661 , tsHeapKind = joinKind (tsHeapKind st) k 662 } 663 664 -- | Check if a register is a public pointer for refined heap. 665 isPublicPointer :: Reg -> TaintState -> Bool 666 isPublicPointer r st = 667 getKind r st == KindPtr && getProvenance r st == ProvPublic 668 669 -- | Check if a register is a pointer to secret data. 670 isSecretDataPointer :: Reg -> TaintState -> Bool 671 isSecretDataPointer r st = 672 getTaint r st == Public && 673 getKind r st == KindPtr && 674 getProvenance r st == ProvSecretData 675 676 -- | Set a heap slot for a (base, offset) pair. 677 setHeapSlot :: Reg -> Int -> Taint -> Provenance 678 -> RegKind -> TaintState -> TaintState 679 setHeapSlot base off t p k st = st 680 { tsHeapSlots = Map.insert (base, off) (t, p, k) 681 (tsHeapSlots st) } 682 683 -- | Get a heap slot for a (base, offset) pair. 684 getHeapSlot :: Reg -> Int -> TaintState 685 -> Maybe (Taint, Provenance, RegKind) 686 getHeapSlot base off st = 687 Map.lookup (base, off) (tsHeapSlots st) 688 689 -- | Set taint for a loaded value from a secondary stack slot. 690 setTaintLoadSecStack :: RuntimeConfig -> Reg -> Int 691 -> TaintState -> TaintState 692 setTaintLoadSecStack rt dst offset st 693 | isPublicRoot rt dst = 694 setTaintProvKind dst Public ProvPublic KindPtr st 695 | otherwise = 696 let (defT, defP, defK) = if tsAssumeStgPublic st 697 then (Public, ProvPublic, KindPtr) 698 else (Unknown, ProvUnknown, KindUnknown) 699 taint = IM.findWithDefault defT offset 700 (tsStgStack st) 701 prov = IM.findWithDefault defP offset 702 (tsStgStackProv st) 703 kind = IM.findWithDefault defK offset 704 (tsStgStackKind st) 705 in setTaintProvKind dst taint prov kind st 706 707 -- | Update for Add, shifting secondary stack map when 708 -- base += imm. 709 updateWithSecStackShiftAdd 710 :: RuntimeConfig -> Reg -> Reg -> Operand 711 -> Taint -> Provenance -> RegKind -> TaintState 712 -> TaintState 713 updateWithSecStackShiftAdd rt dst r1 op t p k st 714 | dst == SP = 715 clearStackMap (setTaintProvKind dst t p k st) 716 | isSecStackBase rt dst, r1 == dst, OpImm imm <- op = 717 let delta = negate (fromInteger imm) 718 in shiftSecStackMap delta 719 (setTaintProvKind dst t p k st) 720 | isSecStackBase rt dst = 721 clearSecStackMap (setTaintProvKind dst t p k st) 722 | otherwise = setTaintProvKind dst t p k st 723 724 -- | Update for Sub, shifting secondary stack map when 725 -- base -= imm. 726 updateWithSecStackShiftSub 727 :: RuntimeConfig -> Reg -> Reg -> Operand 728 -> Taint -> Provenance -> RegKind -> TaintState 729 -> TaintState 730 updateWithSecStackShiftSub rt dst r1 op t p k st 731 | dst == SP = 732 clearStackMap (setTaintProvKind dst t p k st) 733 | isSecStackBase rt dst, r1 == dst, OpImm imm <- op = 734 let delta = fromInteger imm 735 in shiftSecStackMap delta 736 (setTaintProvKind dst t p k st) 737 | isSecStackBase rt dst = 738 clearSecStackMap (setTaintProvKind dst t p k st) 739 | otherwise = setTaintProvKind dst t p k st 740 741 -- | Track a store to stack if address matches SP or 742 -- secondary stack base. 743 storeToStack :: RuntimeConfig -> Reg -> AddrMode 744 -> TaintState -> TaintState 745 storeToStack rt src addr st = case addr of 746 -- Hardware stack (SP) 747 BaseImm SP imm -> 748 let off = fromInteger imm 749 t = getTaint src st 750 p = getProvenance src st 751 k = getKind src st 752 in setStackTaintProvKind off t p k st 753 PreIndex SP imm -> 754 let off = fromInteger imm 755 t = getTaint src st 756 p = getProvenance src st 757 k = getKind src st 758 st' = setStackTaintProvKind off t p k st 759 in clearStackMap st' 760 PostIndex SP _ -> 761 let t = getTaint src st 762 p = getProvenance src st 763 k = getKind src st 764 st' = setStackTaintProvKind 0 t p k st 765 in clearStackMap st' 766 -- Secondary stack 767 BaseImm base imm | isSecStackBase rt base -> 768 let off = fromInteger imm 769 t = getTaint src st 770 p = getProvenance src st 771 k = getKind src st 772 in setSecStackTaintProvKind off t p k st 773 PreIndex base imm | isSecStackBase rt base -> 774 let off = fromInteger imm 775 t = getTaint src st 776 p = getProvenance src st 777 k = getKind src st 778 delta = negate (fromInteger imm) 779 st' = setSecStackTaintProvKind off t p k st 780 in shiftSecStackMap delta st' 781 PostIndex base imm | isSecStackBase rt base -> 782 let t = getTaint src st 783 p = getProvenance src st 784 k = getKind src st 785 delta = negate (fromInteger imm) 786 st' = setSecStackTaintProvKind 0 t p k st 787 in shiftSecStackMap delta st' 788 -- Refined heap slot 789 BaseImm base imm | isPublicPointer base st -> 790 let off = fromInteger imm 791 t = getTaint src st 792 p = getProvenance src st 793 k = getKind src st 794 in setHeapSlot base off t p k 795 (joinHeapBucket t p k st) 796 -- Other non-stack stores 797 _ -> let t = getTaint src st 798 p = getProvenance src st 799 k = getKind src st 800 in joinHeapBucket t p k st 801 802 -- | Track a store pair to stack. 803 storePairToStack :: RuntimeConfig -> Reg -> Reg -> AddrMode 804 -> TaintState -> TaintState 805 storePairToStack rt src1 src2 addr st = case addr of 806 -- Hardware stack (SP) 807 BaseImm SP imm -> 808 let off = fromInteger imm 809 t1 = getTaint src1 st; p1 = getProvenance src1 st 810 k1 = getKind src1 st 811 t2 = getTaint src2 st; p2 = getProvenance src2 st 812 k2 = getKind src2 st 813 in setStackTaintProvKind off t1 p1 k1 814 (setStackTaintProvKind (off + 8) t2 p2 k2 st) 815 PreIndex SP imm -> 816 let off = fromInteger imm 817 t1 = getTaint src1 st; p1 = getProvenance src1 st 818 k1 = getKind src1 st 819 t2 = getTaint src2 st; p2 = getProvenance src2 st 820 k2 = getKind src2 st 821 st' = setStackTaintProvKind off t1 p1 k1 822 (setStackTaintProvKind (off+8) t2 p2 k2 st) 823 in clearStackMap st' 824 PostIndex SP _ -> 825 let t1 = getTaint src1 st; p1 = getProvenance src1 st 826 k1 = getKind src1 st 827 t2 = getTaint src2 st; p2 = getProvenance src2 st 828 k2 = getKind src2 st 829 st' = setStackTaintProvKind 0 t1 p1 k1 830 (setStackTaintProvKind 8 t2 p2 k2 st) 831 in clearStackMap st' 832 -- Secondary stack 833 BaseImm base imm | isSecStackBase rt base -> 834 let off = fromInteger imm 835 t1 = getTaint src1 st; p1 = getProvenance src1 st 836 k1 = getKind src1 st 837 t2 = getTaint src2 st; p2 = getProvenance src2 st 838 k2 = getKind src2 st 839 in setSecStackTaintProvKind off t1 p1 k1 840 (setSecStackTaintProvKind (off+8) t2 p2 k2 st) 841 PreIndex base imm | isSecStackBase rt base -> 842 let off = fromInteger imm 843 t1 = getTaint src1 st; p1 = getProvenance src1 st 844 k1 = getKind src1 st 845 t2 = getTaint src2 st; p2 = getProvenance src2 st 846 k2 = getKind src2 st 847 delta = negate (fromInteger imm) 848 st' = setSecStackTaintProvKind off t1 p1 k1 849 (setSecStackTaintProvKind (off+8) t2 p2 k2 st) 850 in shiftSecStackMap delta st' 851 PostIndex base imm | isSecStackBase rt base -> 852 let t1 = getTaint src1 st; p1 = getProvenance src1 st 853 k1 = getKind src1 st 854 t2 = getTaint src2 st; p2 = getProvenance src2 st 855 k2 = getKind src2 st 856 delta = negate (fromInteger imm) 857 st' = setSecStackTaintProvKind 0 t1 p1 k1 858 (setSecStackTaintProvKind 8 t2 p2 k2 st) 859 in shiftSecStackMap delta st' 860 -- Refined heap slot 861 BaseImm base imm | isPublicPointer base st -> 862 let off = fromInteger imm 863 t1 = getTaint src1 st; p1 = getProvenance src1 st 864 k1 = getKind src1 st 865 t2 = getTaint src2 st; p2 = getProvenance src2 st 866 k2 = getKind src2 st 867 st1 = setHeapSlot base off t1 p1 k1 st 868 st2 = setHeapSlot base (off+8) t2 p2 k2 st1 869 in joinHeapBucket t1 p1 k1 870 (joinHeapBucket t2 p2 k2 st2) 871 -- Other non-stack stores 872 _ -> let t1 = getTaint src1 st; p1 = getProvenance src1 st 873 k1 = getKind src1 st 874 t2 = getTaint src2 st; p2 = getProvenance src2 st 875 k2 = getKind src2 st 876 in joinHeapBucket t1 p1 k1 877 (joinHeapBucket t2 p2 k2 st) 878 879 -- | Load from memory, handling SP, secondary stack, 880 -- GOT entries, secret pointers, and heap. 881 loadFromStack :: RuntimeConfig -> Reg -> AddrMode 882 -> TaintState -> TaintState 883 loadFromStack rt dst addr st = case addr of 884 -- Hardware stack (SP) 885 BaseImm SP imm -> 886 setTaintLoadStack rt dst (fromInteger imm) st 887 PreIndex SP imm -> 888 clearStackMap 889 (setTaintLoadStack rt dst (fromInteger imm) st) 890 PostIndex SP imm -> 891 clearStackMap 892 (setTaintLoadStack rt dst (fromInteger imm) st) 893 -- Secondary stack 894 BaseImm base imm | isSecStackBase rt base -> 895 setTaintLoadSecStack rt dst (fromInteger imm) st 896 PreIndex base imm | isSecStackBase rt base -> 897 let delta = negate (fromInteger imm) 898 in shiftSecStackMap delta 899 (setTaintLoadSecStack rt dst (fromInteger imm) st) 900 PostIndex base imm | isSecStackBase rt base -> 901 let delta = negate (fromInteger imm) 902 in shiftSecStackMap delta 903 (setTaintLoadSecStack rt dst (fromInteger imm) st) 904 -- Symbol/literal loads 905 BaseSymbol _ _ -> 906 setTaintProv dst Public ProvPublic st 907 Literal _ -> 908 setTaintProv dst Public ProvPublic st 909 -- Secret data pointer: [base, #imm] 910 BaseImm base _imm | isSecretDataPointer base st -> 911 setTaintProvKind dst Secret ProvUnknown KindScalar st 912 -- Refined heap slot 913 BaseImm base imm | isPublicPointer base st -> 914 setTaintLoadHeapSlot rt dst base (fromInteger imm) st 915 -- Secret data pointer with register offset 916 BaseReg base _idx | isSecretDataPointer base st -> 917 setTaintProvKind dst Secret ProvUnknown KindScalar st 918 BaseRegShift base _idx _shift 919 | isSecretDataPointer base st -> 920 setTaintProvKind dst Secret ProvUnknown KindScalar st 921 BaseRegExtend base _idx _ext 922 | isSecretDataPointer base st -> 923 setTaintProvKind dst Secret ProvUnknown KindScalar st 924 -- Other loads: coarse heap bucket 925 _ -> setTaintLoadHeap rt dst st 926 927 -- | Load pair from stack. 928 loadPairFromStack :: RuntimeConfig -> Reg -> Reg 929 -> AddrMode -> TaintState -> TaintState 930 loadPairFromStack rt dst1 dst2 addr st = case addr of 931 -- Hardware stack (SP) 932 BaseImm SP imm -> 933 let off = fromInteger imm 934 in setTaintLoadStack rt dst1 off 935 (setTaintLoadStack rt dst2 (off + 8) st) 936 PreIndex SP imm -> 937 let off = fromInteger imm 938 st' = setTaintLoadStack rt dst1 off 939 (setTaintLoadStack rt dst2 (off + 8) st) 940 in clearStackMap st' 941 PostIndex SP imm -> 942 let off = fromInteger imm 943 st' = setTaintLoadStack rt dst1 off 944 (setTaintLoadStack rt dst2 (off + 8) st) 945 in clearStackMap st' 946 -- Secondary stack 947 BaseImm base imm | isSecStackBase rt base -> 948 let off = fromInteger imm 949 in setTaintLoadSecStack rt dst1 off 950 (setTaintLoadSecStack rt dst2 (off + 8) st) 951 PreIndex base imm | isSecStackBase rt base -> 952 let off = fromInteger imm 953 delta = negate (fromInteger imm) 954 st' = setTaintLoadSecStack rt dst1 off 955 (setTaintLoadSecStack rt dst2 (off+8) st) 956 in shiftSecStackMap delta st' 957 PostIndex base imm | isSecStackBase rt base -> 958 let off = fromInteger imm 959 delta = negate (fromInteger imm) 960 st' = setTaintLoadSecStack rt dst1 off 961 (setTaintLoadSecStack rt dst2 (off+8) st) 962 in shiftSecStackMap delta st' 963 -- Secret data pointer 964 BaseImm base _imm | isSecretDataPointer base st -> 965 setTaintProvKind dst1 Secret ProvUnknown KindScalar 966 (setTaintProvKind dst2 Secret ProvUnknown 967 KindScalar st) 968 -- Refined heap slot 969 BaseImm base imm | isPublicPointer base st -> 970 let off = fromInteger imm 971 in setTaintLoadHeapSlot rt dst1 base off 972 (setTaintLoadHeapSlot rt dst2 base (off+8) st) 973 -- Other loads: coarse heap bucket 974 _ -> setTaintLoadHeap rt dst1 975 (setTaintLoadHeap rt dst2 st) 976 977 -- | Get taint of an operand. 978 operandTaint :: Operand -> TaintState -> Taint 979 operandTaint op st = case op of 980 OpReg r -> getTaint r st 981 OpImm _ -> Public 982 OpShiftedReg r _ -> getTaint r st 983 OpExtendedReg r _ -> getTaint r st 984 OpLabel _ -> Public 985 OpAddr addr -> addrBaseTaint addr st 986 987 -- | Get provenance of an operand. 988 operandProv :: Operand -> TaintState -> Provenance 989 operandProv op st = case op of 990 OpReg r -> getProvenance r st 991 OpImm _ -> ProvPublic 992 OpShiftedReg r _ -> getProvenance r st 993 OpExtendedReg r _ -> getProvenance r st 994 OpLabel _ -> ProvPublic 995 OpAddr _ -> ProvUnknown 996 997 -- | Get kind of an operand. 998 operandKind :: Operand -> TaintState -> RegKind 999 operandKind op st = case op of 1000 OpReg r -> getKind r st 1001 OpImm _ -> KindScalar 1002 OpShiftedReg r _ -> getKind r st 1003 OpExtendedReg r _ -> getKind r st 1004 OpLabel _ -> KindPtr 1005 OpAddr _ -> KindUnknown 1006 1007 -- | Compute kind for pointer arithmetic. 1008 pointerArithKind :: Reg -> Operand -> TaintState -> RegKind 1009 pointerArithKind base op st = 1010 case op of 1011 OpImm _ | getKind base st == KindPtr -> KindPtr 1012 _ -> KindScalar 1013 1014 -- | Check if operand is a pointer-untagging mask. 1015 isPointerUntagMask :: RuntimeConfig -> Operand -> Bool 1016 isPointerUntagMask rt (OpImm imm) = 1017 imm `elem` rtUntagMasks rt 1018 isPointerUntagMask _ _ = False 1019 1020 -- | Get taint of address base register. 1021 addrBaseTaint :: AddrMode -> TaintState -> Taint 1022 addrBaseTaint addr st = case addr of 1023 BaseImm r _ -> getTaint r st 1024 BaseReg r _ -> getTaint r st 1025 BaseRegShift r _ _ -> getTaint r st 1026 BaseRegExtend r _ _ -> getTaint r st 1027 BaseSymbol r _ -> getTaint r st 1028 PreIndex r _ -> getTaint r st 1029 PostIndex r _ -> getTaint r st 1030 Literal _ -> Public 1031 1032 -- | Join two taints. 1033 join2 :: Taint -> Taint -> Taint 1034 join2 = joinTaint 1035 1036 -- | Join three taints. 1037 join3 :: Taint -> Taint -> Taint -> Taint 1038 join3 a b c = joinTaint a (joinTaint b c) 1039 1040 -- | Join two provenances. 1041 provJoin2 :: Provenance -> Provenance -> Provenance 1042 provJoin2 = joinProvenance 1043 1044 -- | Join three provenances. 1045 provJoin3 :: Provenance -> Provenance -> Provenance 1046 -> Provenance 1047 provJoin3 a b c = joinProvenance a (joinProvenance b c) 1048 1049 -- | Invalidate caller-saved registers after a call. 1050 invalidateCallerSaved :: TaintState -> TaintState 1051 invalidateCallerSaved st = st 1052 { tsRegs = foldr (\r -> Map.insert r Unknown) 1053 (tsRegs st) callerSaved 1054 , tsProv = foldr (\r -> Map.insert r ProvUnknown) 1055 (tsProv st) callerSaved 1056 , tsKind = foldr (\r -> Map.insert r KindUnknown) 1057 (tsKind st) callerSaved 1058 } 1059 where 1060 callerSaved = 1061 [ X0, X1, X2, X3, X4, X5, X6, X7 1062 , X8, X9, X10, X11, X12, X13, X14, X15 1063 , X16, X17 1064 ] 1065 1066 -- | Join two taint states (element-wise join). 1067 joinTaintState :: TaintState -> TaintState -> TaintState 1068 joinTaintState a b = TaintState 1069 { tsRegs = Map.unionWith joinTaint 1070 (tsRegs a) (tsRegs b) 1071 , tsStack = IM.unionWith joinTaint 1072 (tsStack a) (tsStack b) 1073 , tsProv = Map.unionWith joinProvenance 1074 (tsProv a) (tsProv b) 1075 , tsStackProv = IM.unionWith joinProvenance 1076 (tsStackProv a) (tsStackProv b) 1077 , tsKind = Map.unionWith joinKind 1078 (tsKind a) (tsKind b) 1079 , tsStackKind = IM.unionWith joinKind 1080 (tsStackKind a) (tsStackKind b) 1081 , tsStgStack = IM.unionWith joinTaint 1082 (tsStgStack a) (tsStgStack b) 1083 , tsStgStackProv = IM.unionWith joinProvenance 1084 (tsStgStackProv a) (tsStgStackProv b) 1085 , tsStgStackKind = IM.unionWith joinKind 1086 (tsStgStackKind a) (tsStgStackKind b) 1087 , tsHeapTaint = joinTaint 1088 (tsHeapTaint a) (tsHeapTaint b) 1089 , tsHeapProv = joinProvenance 1090 (tsHeapProv a) (tsHeapProv b) 1091 , tsHeapKind = joinKind 1092 (tsHeapKind a) (tsHeapKind b) 1093 , tsHeapSlots = Map.unionWith joinSlot 1094 (tsHeapSlots a) (tsHeapSlots b) 1095 , tsAssumeStgPublic = 1096 let l = tsAssumeStgPublic a 1097 r = tsAssumeStgPublic b 1098 in if l == r then l 1099 else error 1100 "joinTaintState: mismatched tsAssumeStgPublic" 1101 } 1102 where 1103 joinSlot (t1, p1, k1) (t2, p2, k2) = 1104 (joinTaint t1 t2, joinProvenance p1 p2, 1105 joinKind k1 k2) 1106 1107 -- | Run forward dataflow analysis over a CFG. 1108 runDataflow :: RuntimeConfig -> CFG -> IntMap TaintState 1109 runDataflow rt cfg 1110 | nBlocks == 0 = IM.empty 1111 | otherwise = go initWorklist initIn initOut 1112 where 1113 nBlocks = cfgBlockCount cfg 1114 baseState = initTaintState rt 1115 emptyState = emptyTaintStateWith 1116 (tsAssumeStgPublic baseState) 1117 1118 initIn = IM.fromList 1119 [(i, baseState) | i <- [0..nBlocks-1]] 1120 initOut = IM.empty 1121 initWorklist = IS.fromList [0..nBlocks-1] 1122 1123 go :: IntSet -> IntMap TaintState 1124 -> IntMap TaintState -> IntMap TaintState 1125 go worklist inStates outStates 1126 | IS.null worklist = inStates 1127 | otherwise = 1128 let (idx, worklist') = IS.deleteFindMin worklist 1129 bb = indexBlock cfg idx 1130 inState = IM.findWithDefault baseState 1131 idx inStates 1132 outState = analyzeBlock rt (bbLines bb) 1133 inState 1134 oldOut = IM.lookup idx outStates 1135 changed = oldOut /= Just outState 1136 outStates' = IM.insert idx outState 1137 outStates 1138 succs = blockSuccessors cfg idx 1139 (worklist'', inStates') = if changed 1140 then propagateToSuccs succs outState 1141 worklist' inStates 1142 else (worklist', inStates) 1143 in go worklist'' inStates' outStates' 1144 1145 propagateToSuccs [] _ wl ins = (wl, ins) 1146 propagateToSuccs (s:ss) out wl ins = 1147 let oldIn = IM.findWithDefault emptyState s ins 1148 newIn = joinTaintState oldIn out 1149 changed = oldIn /= newIn 1150 ins' = IM.insert s newIn ins 1151 wl' = if changed then IS.insert s wl else wl 1152 in propagateToSuccs ss out wl' ins' 1153 1154 -- ------------------------------------------------------------- 1155 -- Function summaries for inter-procedural analysis 1156 1157 -- | Function summary: taint state at return. 1158 newtype FuncSummary = FuncSummary 1159 { summaryState :: TaintState } 1160 deriving (Eq, Show) 1161 1162 -- | Initial summary with configurable secondary stack assumption. 1163 initSummaryWith :: Bool -> FuncSummary 1164 initSummaryWith assumeStgPublic = FuncSummary $ TaintState 1165 { tsRegs = Map.fromList 1166 [(r, Unknown) | r <- callerSavedRegs] 1167 , tsStack = IM.empty 1168 , tsProv = Map.fromList 1169 [(r, ProvUnknown) | r <- callerSavedRegs] 1170 , tsStackProv = IM.empty 1171 , tsKind = Map.fromList 1172 [(r, KindUnknown) | r <- callerSavedRegs] 1173 , tsStackKind = IM.empty 1174 , tsStgStack = IM.empty 1175 , tsStgStackProv = IM.empty 1176 , tsStgStackKind = IM.empty 1177 , tsHeapTaint = Unknown 1178 , tsHeapProv = ProvUnknown 1179 , tsHeapKind = KindUnknown 1180 , tsHeapSlots = Map.empty 1181 , tsAssumeStgPublic = assumeStgPublic 1182 } 1183 1184 -- | Initial conservative summary. 1185 initSummary :: RuntimeConfig -> FuncSummary 1186 initSummary rt = initSummaryWith assumeStg 1187 where 1188 assumeStg = case rtSecondaryStack rt of 1189 Nothing -> False 1190 Just ss -> ssAssumePublic ss 1191 1192 -- | Caller-saved registers per AArch64 ABI. 1193 callerSavedRegs :: [Reg] 1194 callerSavedRegs = 1195 [ X0, X1, X2, X3, X4, X5, X6, X7 1196 , X8, X9, X10, X11, X12, X13, X14, X15 1197 , X16, X17 1198 ] 1199 1200 -- | Join two summaries. 1201 joinSummary :: FuncSummary -> FuncSummary -> FuncSummary 1202 joinSummary (FuncSummary a) (FuncSummary b) = 1203 FuncSummary (joinTaintState a b) 1204 1205 -- | Apply a function summary at a call site. 1206 applySummary :: FuncSummary -> TaintState -> TaintState 1207 applySummary (FuncSummary summ) st = st 1208 { tsRegs = foldr applyReg (tsRegs st) callerSavedRegs 1209 , tsProv = foldr applyProv (tsProv st) callerSavedRegs 1210 , tsKind = foldr applyKind (tsKind st) callerSavedRegs 1211 } 1212 where 1213 summRegs = tsRegs summ 1214 summProv = tsProv summ 1215 summKind = tsKind summ 1216 applyReg r s = 1217 Map.insert r 1218 (Map.findWithDefault Unknown r summRegs) s 1219 applyProv r s = 1220 Map.insert r 1221 (Map.findWithDefault ProvUnknown r summProv) s 1222 applyKind r s = 1223 Map.insert r 1224 (Map.findWithDefault KindUnknown r summKind) s 1225 1226 -- | Run dataflow for a single function. 1227 runFunctionDataflow :: RuntimeConfig -> CFG -> [Int] 1228 -> Map Text FuncSummary -> TaintState 1229 runFunctionDataflow rt cfg blockIndices summaries = 1230 let inStates = runFunctionBlocks rt cfg blockIndices 1231 summaries 1232 baseState = initTaintState rt 1233 returnOuts = 1234 [ analyzeBlockWithSummaries rt bb inState 1235 summaries 1236 | i <- blockIndices 1237 , let bb = indexBlock cfg i 1238 inState = IM.findWithDefault baseState i 1239 inStates 1240 , endsWithRet bb 1241 ] 1242 in case returnOuts of 1243 [] -> baseState 1244 (o:os) -> foldl' joinTaintState o os 1245 1246 -- | Check if block ends with a return. 1247 endsWithRet :: BasicBlock -> Bool 1248 endsWithRet bb = case bbLines bb of 1249 [] -> False 1250 ls -> case lineInstr (last ls) of 1251 Just (Ret _) -> True 1252 _ -> False 1253 1254 -- | Run dataflow on a subset of blocks (one function). 1255 runFunctionBlocks :: RuntimeConfig -> CFG -> [Int] 1256 -> Map Text FuncSummary 1257 -> IntMap TaintState 1258 runFunctionBlocks _ _ [] _ = IM.empty 1259 runFunctionBlocks rt cfg (entryIdx:rest) summaries = 1260 go initWorklist initIn IM.empty 1261 where 1262 blockSet = IS.fromList (entryIdx:rest) 1263 baseState = initTaintState rt 1264 emptyState = emptyTaintStateWith 1265 (tsAssumeStgPublic baseState) 1266 1267 initIn = IM.singleton entryIdx baseState 1268 initWorklist = IS.singleton entryIdx 1269 1270 go wl inStates outStates 1271 | IS.null wl = inStates 1272 | otherwise = 1273 let (idx, wl') = IS.deleteFindMin wl 1274 bb = indexBlock cfg idx 1275 inState = IM.findWithDefault baseState 1276 idx inStates 1277 outState = analyzeBlockWithSummaries rt bb 1278 inState summaries 1279 oldOut = IM.lookup idx outStates 1280 changed = oldOut /= Just outState 1281 outStates' = IM.insert idx outState 1282 outStates 1283 succs = filter (`IS.member` blockSet) 1284 (blockSuccessors cfg idx) 1285 (wl'', inStates') = if changed 1286 then propagate succs outState wl' inStates 1287 else (wl', inStates) 1288 in go wl'' inStates' outStates' 1289 1290 propagate [] _ wl ins = (wl, ins) 1291 propagate (s:ss) out wl ins = 1292 let oldIn = IM.findWithDefault emptyState s ins 1293 newIn = joinTaintState oldIn out 1294 ins' = IM.insert s newIn ins 1295 wl' = if oldIn /= newIn 1296 then IS.insert s wl else wl 1297 in propagate ss out wl' ins' 1298 1299 -- | Analyze a block with call summaries. 1300 analyzeBlockWithSummaries :: RuntimeConfig -> BasicBlock 1301 -> TaintState 1302 -> Map Text FuncSummary 1303 -> TaintState 1304 analyzeBlockWithSummaries rt bb st0 summaries = 1305 foldl' go st0 (bbLines bb) 1306 where 1307 go st l = case lineInstr l of 1308 Nothing -> st 1309 Just instr -> 1310 transferWithSummary rt instr st summaries 1311 1312 -- | Transfer with summary application for calls. 1313 transferWithSummary :: RuntimeConfig -> Instr -> TaintState 1314 -> Map Text FuncSummary -> TaintState 1315 transferWithSummary rt instr st summaries = case instr of 1316 Bl target -> 1317 case Map.lookup target summaries of 1318 Just summ -> applySummary summ st 1319 Nothing -> invalidateCallerSaved st 1320 Blr _ -> invalidateCallerSaved st 1321 B target 1322 | isFunctionLabel rt target -> 1323 case Map.lookup target summaries of 1324 Just _ -> st -- In-file: CFG edge propagates 1325 Nothing -> invalidateCallerSaved st 1326 Br _ -> invalidateCallerSaved st 1327 _ -> transfer rt instr st 1328 1329 -- | Run inter-procedural fixpoint analysis. 1330 runInterProc :: RuntimeConfig -> CFG 1331 -> Map Text FuncSummary 1332 runInterProc rt cfg = go initSummaries (Set.fromList funcs) 1333 where 1334 funcs = functionLabels cfg 1335 baseSummary = initSummary rt 1336 initSummaries = Map.fromList 1337 [(f, baseSummary) | f <- funcs] 1338 1339 funcBlocksCache = buildFunctionBlocksMap cfg 1340 callerMapCache = buildCallerMap rt cfg funcBlocksCache 1341 1342 lookupBlocks f = 1343 Map.findWithDefault [] f funcBlocksCache 1344 lookupCallers f = 1345 Map.findWithDefault [] f callerMapCache 1346 1347 go summaries worklist 1348 | Set.null worklist = summaries 1349 | otherwise = 1350 let (func, worklist') = 1351 Set.deleteFindMin worklist 1352 blockIdxs = lookupBlocks func 1353 outState = runFunctionDataflow rt cfg 1354 blockIdxs summaries 1355 newSumm = FuncSummary outState 1356 oldSumm = Map.findWithDefault baseSummary 1357 func summaries 1358 changed = newSumm /= oldSumm 1359 summaries' = Map.insert func newSumm 1360 summaries 1361 callers = lookupCallers func 1362 worklist'' = if changed 1363 then foldr Set.insert worklist' callers 1364 else worklist' 1365 in go summaries' worklist'' 1366 1367 -- ------------------------------------------------------------- 1368 -- Config-aware dataflow analysis 1369 1370 -- | Run forward dataflow with taint config. 1371 runDataflowWithConfig :: RuntimeConfig -> TaintConfig 1372 -> CFG -> IntMap TaintState 1373 runDataflowWithConfig rt tcfg cfg 1374 | nBlocks == 0 = IM.empty 1375 | otherwise = go initWorklist initIn IM.empty 1376 where 1377 nBlocks = cfgBlockCount cfg 1378 assumeStg = tcAssumeStgPublic tcfg 1379 baseState = initTaintStateWith rt assumeStg 1380 emptyState = emptyTaintStateWith assumeStg 1381 1382 initIn = IM.fromList 1383 [ (i, entryState i (indexBlock cfg i)) 1384 | i <- [0..nBlocks-1] 1385 ] 1386 1387 entryState idx bb = 1388 case bbLabel bb of 1389 Nothing -> baseState 1390 Just lbl -> 1391 case functionBlocks cfg lbl of 1392 (entry:_) | entry == idx -> 1393 case Map.lookup lbl (tcPolicies tcfg) of 1394 Nothing -> baseState 1395 Just policy -> 1396 seedArgs rt policy baseState 1397 _ -> baseState 1398 1399 initWorklist = IS.fromList [0..nBlocks-1] 1400 1401 go worklist inStates outStates 1402 | IS.null worklist = inStates 1403 | otherwise = 1404 let (idx, worklist') = IS.deleteFindMin worklist 1405 bb = indexBlock cfg idx 1406 inState = IM.findWithDefault baseState 1407 idx inStates 1408 outState = analyzeBlock rt (bbLines bb) 1409 inState 1410 oldOut = IM.lookup idx outStates 1411 changed = oldOut /= Just outState 1412 outStates' = IM.insert idx outState 1413 outStates 1414 succs = blockSuccessors cfg idx 1415 (worklist'', inStates') = if changed 1416 then propagate succs outState 1417 worklist' inStates 1418 else (worklist', inStates) 1419 in go worklist'' inStates' outStates' 1420 1421 propagate [] _ wl ins = (wl, ins) 1422 propagate (s:ss) out wl ins = 1423 let oldIn = IM.findWithDefault emptyState s ins 1424 newIn = joinTaintState oldIn out 1425 ins' = IM.insert s newIn ins 1426 wl' = if oldIn /= newIn 1427 then IS.insert s wl else wl 1428 in propagate ss out wl' ins' 1429 1430 -- | Run forward dataflow with taint config and summaries. 1431 runDataflowWithConfigAndSummaries 1432 :: RuntimeConfig -> TaintConfig 1433 -> Map Text FuncSummary -> CFG -> IntMap TaintState 1434 runDataflowWithConfigAndSummaries rt tcfg summaries cfg 1435 | nBlocks == 0 = IM.empty 1436 | otherwise = go initWorklist initIn IM.empty 1437 where 1438 nBlocks = cfgBlockCount cfg 1439 assumeStg = tcAssumeStgPublic tcfg 1440 baseState = initTaintStateWith rt assumeStg 1441 emptyState = emptyTaintStateWith assumeStg 1442 1443 initIn = IM.fromList 1444 [ (i, entryState i (indexBlock cfg i)) 1445 | i <- [0..nBlocks-1] 1446 ] 1447 1448 entryState idx bb = 1449 case bbLabel bb of 1450 Nothing -> baseState 1451 Just lbl -> 1452 case functionBlocks cfg lbl of 1453 (entry:_) | entry == idx -> 1454 case Map.lookup lbl (tcPolicies tcfg) of 1455 Nothing -> baseState 1456 Just policy -> 1457 seedArgs rt policy baseState 1458 _ -> baseState 1459 1460 initWorklist = IS.fromList [0..nBlocks-1] 1461 1462 go worklist inStates outStates 1463 | IS.null worklist = inStates 1464 | otherwise = 1465 let (idx, worklist') = IS.deleteFindMin worklist 1466 bb = indexBlock cfg idx 1467 inState = IM.findWithDefault baseState 1468 idx inStates 1469 outState = analyzeBlockWithSummaries rt bb 1470 inState summaries 1471 oldOut = IM.lookup idx outStates 1472 changed = oldOut /= Just outState 1473 outStates' = IM.insert idx outState 1474 outStates 1475 succs = blockSuccessors cfg idx 1476 (worklist'', inStates') = if changed 1477 then propagateS succs outState 1478 worklist' inStates 1479 else (worklist', inStates) 1480 in go worklist'' inStates' outStates' 1481 1482 propagateS [] _ wl ins = (wl, ins) 1483 propagateS (s:ss) out wl ins = 1484 let oldIn = IM.findWithDefault emptyState s ins 1485 newIn = joinTaintState oldIn out 1486 ins' = IM.insert s newIn ins 1487 wl' = if oldIn /= newIn 1488 then IS.insert s wl else wl 1489 in propagateS ss out wl' ins' 1490 1491 -- | Run inter-procedural analysis with taint config. 1492 runInterProcWithConfig :: RuntimeConfig -> TaintConfig 1493 -> CFG -> Map Text FuncSummary 1494 runInterProcWithConfig rt tcfg cfg = 1495 go initSummaries (Set.fromList funcs) 1496 where 1497 funcs = functionLabels cfg 1498 assumeStg = tcAssumeStgPublic tcfg 1499 baseSummary = initSummaryWith assumeStg 1500 initSummaries = Map.fromList 1501 [(f, baseSummary) | f <- funcs] 1502 1503 funcBlocksCache = buildFunctionBlocksMap cfg 1504 callerMapCache = 1505 buildCallerMap rt cfg funcBlocksCache 1506 1507 lookupBlocks f = 1508 Map.findWithDefault [] f funcBlocksCache 1509 lookupCallers f = 1510 Map.findWithDefault [] f callerMapCache 1511 1512 go summaries worklist 1513 | Set.null worklist = summaries 1514 | otherwise = 1515 let (func, worklist') = 1516 Set.deleteFindMin worklist 1517 blockIdxs = lookupBlocks func 1518 outState = 1519 runFunctionDataflowWithConfig rt tcfg cfg 1520 func blockIdxs summaries 1521 newSumm = FuncSummary outState 1522 oldSumm = Map.findWithDefault baseSummary 1523 func summaries 1524 changed = newSumm /= oldSumm 1525 summaries' = Map.insert func newSumm 1526 summaries 1527 callers = lookupCallers func 1528 worklist'' = if changed 1529 then foldr Set.insert worklist' callers 1530 else worklist' 1531 in go summaries' worklist'' 1532 1533 -- | Run function dataflow with config-based entry seeding. 1534 runFunctionDataflowWithConfig 1535 :: RuntimeConfig -> TaintConfig -> CFG -> Text -> [Int] 1536 -> Map Text FuncSummary -> TaintState 1537 runFunctionDataflowWithConfig rt tcfg cfg funcName 1538 blockIndices summaries = 1539 let assumeStg = tcAssumeStgPublic tcfg 1540 baseEntry = initTaintStateWith rt assumeStg 1541 entryState = 1542 case Map.lookup funcName (tcPolicies tcfg) of 1543 Nothing -> baseEntry 1544 Just policy -> seedArgs rt policy baseEntry 1545 inStates = runFunctionBlocksWithEntry rt cfg 1546 blockIndices summaries entryState 1547 returnOuts = 1548 [ analyzeBlockWithSummaries rt bb inState 1549 summaries 1550 | i <- blockIndices 1551 , let bb = indexBlock cfg i 1552 inState = IM.findWithDefault entryState i 1553 inStates 1554 , endsWithRet bb 1555 ] 1556 in case returnOuts of 1557 [] -> baseEntry 1558 (o:os) -> foldl' joinTaintState o os 1559 1560 -- | Run function blocks with a custom entry state. 1561 runFunctionBlocksWithEntry :: RuntimeConfig -> CFG -> [Int] 1562 -> Map Text FuncSummary 1563 -> TaintState 1564 -> IntMap TaintState 1565 runFunctionBlocksWithEntry _ _ [] _ _ = IM.empty 1566 runFunctionBlocksWithEntry rt cfg (entryIdx:rest) 1567 summaries entryState = 1568 go initWorklist initIn IM.empty 1569 where 1570 blockSet = IS.fromList (entryIdx:rest) 1571 emptyState = emptyTaintStateWith 1572 (tsAssumeStgPublic entryState) 1573 1574 initIn = IM.singleton entryIdx entryState 1575 initWorklist = IS.singleton entryIdx 1576 1577 go wl inStates outStates 1578 | IS.null wl = inStates 1579 | otherwise = 1580 let (idx, wl') = IS.deleteFindMin wl 1581 bb = indexBlock cfg idx 1582 inState = IM.findWithDefault entryState 1583 idx inStates 1584 outState = analyzeBlockWithSummaries rt bb 1585 inState summaries 1586 oldOut = IM.lookup idx outStates 1587 changed = oldOut /= Just outState 1588 outStates' = IM.insert idx outState 1589 outStates 1590 succs = filter (`IS.member` blockSet) 1591 (blockSuccessors cfg idx) 1592 (wl'', inStates') = if changed 1593 then propagate succs outState wl' inStates 1594 else (wl', inStates) 1595 in go wl'' inStates' outStates' 1596 1597 propagate [] _ wl ins = (wl, ins) 1598 propagate (s:ss) out wl ins = 1599 let oldIn = IM.findWithDefault emptyState s ins 1600 newIn = joinTaintState oldIn out 1601 ins' = IM.insert s newIn ins 1602 wl' = if oldIn /= newIn 1603 then IS.insert s wl else wl 1604 in propagate ss out wl' ins'