{-|
Copyright   : (c) Galois Inc, 2026
Maintainer  : Langston Barrett <langston@galois.com>

Read the @.llvm_jump_table_sizes@ section emitted by Clang/LLVM (>= 20)
when compiled with @-mllvm -emit-jump-table-sizes-section@.

The section consists of a sequence of fixed-size entries; each entry is two
target-pointer-sized words written in target endianness:

  * the address of the jump table, and
  * the number of basic-block targets in that jump table.

See @SHT_LLVM_JT_SIZES@ in the LLVM @ELF.h@ header and
<https://llvm.org/docs/Extensions.html>.
-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.Macaw.Memory.LLVMJumpTableSizes
  ( sectionName
  , JumpTableSize(..)
  , SectionSizeError(..)
  , UnresolvableAddress(..)
  , parseLLVMJumpTableSizes
  , llvmJumpTableSizesFromElf
  ) where

import qualified Data.ByteString as BS
import qualified Data.ElfEdit as Elf
import qualified Data.Foldable as F
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Word (Word64)
import           Numeric (showHex)
import qualified Prettyprinter as PP

import           Data.Macaw.Memory
import           Data.Macaw.Memory.ElfLoader (elfAddrWidth)

-- | The section byte count was not a multiple of the entry size.
data SectionSizeError
  = SectionSizeError
      Int  -- ^ actual section size in bytes
      Int  -- ^ required entry size in bytes
  deriving SectionSizeError -> SectionSizeError -> Bool
(SectionSizeError -> SectionSizeError -> Bool)
-> (SectionSizeError -> SectionSizeError -> Bool)
-> Eq SectionSizeError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SectionSizeError -> SectionSizeError -> Bool
== :: SectionSizeError -> SectionSizeError -> Bool
$c/= :: SectionSizeError -> SectionSizeError -> Bool
/= :: SectionSizeError -> SectionSizeError -> Bool
Eq

instance PP.Pretty SectionSizeError where
  pretty :: forall ann. SectionSizeError -> Doc ann
pretty = \case
    SectionSizeError Int
total Int
entry ->
      String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty (ByteString -> String
forall a. Show a => a -> String
show ByteString
sectionName) Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
PP.<+> Doc ann
"section size" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
PP.<+> Int -> Doc ann
forall ann. Int -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty Int
total
        Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
PP.<+> Doc ann
"is not a multiple of entry size" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
PP.<+> Int -> Doc ann
forall ann. Int -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty Int
entry

instance Show SectionSizeError where
  show :: SectionSizeError -> String
show = Doc (ZonkAny 1) -> String
forall a. Show a => a -> String
show (Doc (ZonkAny 1) -> String)
-> (SectionSizeError -> Doc (ZonkAny 1))
-> SectionSizeError
-> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SectionSizeError -> Doc (ZonkAny 1)
forall a ann. Pretty a => a -> Doc ann
forall ann. SectionSizeError -> Doc ann
PP.pretty

-- | An entry's jump-table address could not be resolved in the 'Memory'.
newtype UnresolvableAddress w = UnresolvableAddress (MemWord w)
  deriving UnresolvableAddress w -> UnresolvableAddress w -> Bool
(UnresolvableAddress w -> UnresolvableAddress w -> Bool)
-> (UnresolvableAddress w -> UnresolvableAddress w -> Bool)
-> Eq (UnresolvableAddress w)
forall (w :: Nat).
UnresolvableAddress w -> UnresolvableAddress w -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (w :: Nat).
UnresolvableAddress w -> UnresolvableAddress w -> Bool
== :: UnresolvableAddress w -> UnresolvableAddress w -> Bool
$c/= :: forall (w :: Nat).
UnresolvableAddress w -> UnresolvableAddress w -> Bool
/= :: UnresolvableAddress w -> UnresolvableAddress w -> Bool
Eq

instance MemWidth w => PP.Pretty (UnresolvableAddress w) where
  pretty :: forall ann. UnresolvableAddress w -> Doc ann
pretty = \case
    UnresolvableAddress MemWord w
addr ->
      String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty (ByteString -> String
forall a. Show a => a -> String
show ByteString
sectionName) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
PP.<> Doc ann
": could not resolve address 0x"
        Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
PP.<> String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty (Word64 -> ShowS
forall a. Integral a => a -> ShowS
showHex (MemWord w -> Word64
forall (w :: Nat). MemWord w -> Word64
memWordValue MemWord w
addr) String
"")

instance MemWidth w => Show (UnresolvableAddress w) where
  show :: UnresolvableAddress w -> String
show = Doc (ZonkAny 0) -> String
forall a. Show a => a -> String
show (Doc (ZonkAny 0) -> String)
-> (UnresolvableAddress w -> Doc (ZonkAny 0))
-> UnresolvableAddress w
-> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnresolvableAddress w -> Doc (ZonkAny 0)
forall a ann. Pretty a => a -> Doc ann
forall ann. UnresolvableAddress w -> Doc ann
PP.pretty

-- | The number of basic-block targets (entries) in a jump table, as recorded
-- in a Clang\/LLVM @.llvm_jump_table_sizes@ section.
--
-- The type parameter @w@ is the target address width.
newtype JumpTableSize w = JumpTableSize { forall (w :: Nat). JumpTableSize w -> MemWord w
getJumpTableSize :: MemWord w }
  deriving (JumpTableSize w -> JumpTableSize w -> Bool
(JumpTableSize w -> JumpTableSize w -> Bool)
-> (JumpTableSize w -> JumpTableSize w -> Bool)
-> Eq (JumpTableSize w)
forall (w :: Nat). JumpTableSize w -> JumpTableSize w -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (w :: Nat). JumpTableSize w -> JumpTableSize w -> Bool
== :: JumpTableSize w -> JumpTableSize w -> Bool
$c/= :: forall (w :: Nat). JumpTableSize w -> JumpTableSize w -> Bool
/= :: JumpTableSize w -> JumpTableSize w -> Bool
Eq, Int -> JumpTableSize w -> ShowS
[JumpTableSize w] -> ShowS
JumpTableSize w -> String
(Int -> JumpTableSize w -> ShowS)
-> (JumpTableSize w -> String)
-> ([JumpTableSize w] -> ShowS)
-> Show (JumpTableSize w)
forall (w :: Nat). Int -> JumpTableSize w -> ShowS
forall (w :: Nat). [JumpTableSize w] -> ShowS
forall (w :: Nat). JumpTableSize w -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall (w :: Nat). Int -> JumpTableSize w -> ShowS
showsPrec :: Int -> JumpTableSize w -> ShowS
$cshow :: forall (w :: Nat). JumpTableSize w -> String
show :: JumpTableSize w -> String
$cshowList :: forall (w :: Nat). [JumpTableSize w] -> ShowS
showList :: [JumpTableSize w] -> ShowS
Show)

-- | The ELF section name emitted by LLVM.
sectionName :: BS.ByteString
sectionName :: ByteString
sectionName = ByteString
".llvm_jump_table_sizes"

-- | Parse the raw bytes of a @.llvm_jump_table_sizes@ section.
--
-- Each entry is two pointer-sized words: @(address, entry_count)@. The
-- bytestring length must be an exact multiple of @2 * pointer_size@; otherwise
-- a 'Left' is returned.
parseLLVMJumpTableSizes
  :: forall w
  .  AddrWidthRepr w
  -> Endianness
  -> BS.ByteString
  -> Either SectionSizeError [(MemWord w, JumpTableSize w)]
parseLLVMJumpTableSizes :: forall (w :: Nat).
AddrWidthRepr w
-> Endianness
-> ByteString
-> Either SectionSizeError [(MemWord w, JumpTableSize w)]
parseLLVMJumpTableSizes AddrWidthRepr w
wRepr Endianness
end ByteString
bs0 = AddrWidthRepr w
-> (MemWidth w =>
    Either SectionSizeError [(MemWord w, JumpTableSize w)])
-> Either SectionSizeError [(MemWord w, JumpTableSize w)]
forall (w :: Nat) a. AddrWidthRepr w -> (MemWidth w => a) -> a
addrWidthClass AddrWidthRepr w
wRepr ((MemWidth w =>
  Either SectionSizeError [(MemWord w, JumpTableSize w)])
 -> Either SectionSizeError [(MemWord w, JumpTableSize w)])
-> (MemWidth w =>
    Either SectionSizeError [(MemWord w, JumpTableSize w)])
-> Either SectionSizeError [(MemWord w, JumpTableSize w)]
forall a b. (a -> b) -> a -> b
$
  let ptrBytes :: Int
ptrBytes = Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (AddrWidthRepr w -> Nat
forall (w :: Nat). AddrWidthRepr w -> Nat
addrWidthReprByteCount AddrWidthRepr w
wRepr) :: Int
      entryBytes :: Int
entryBytes = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
ptrBytes
      totalLen :: Int
totalLen = ByteString -> Int
BS.length ByteString
bs0
      readWord :: BS.ByteString -> Word64
      readWord :: ByteString -> Word64
readWord ByteString
bs = case AddrWidthRepr w
wRepr of
        AddrWidthRepr w
Addr32 -> Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Endianness -> ByteString -> Word32
bsWord32 Endianness
end ByteString
bs)
        AddrWidthRepr w
Addr64 -> Endianness -> ByteString -> Word64
bsWord64 Endianness
end ByteString
bs
      go :: MemWidth w => BS.ByteString -> [(MemWord w, JumpTableSize w)]
      go :: MemWidth w => ByteString -> [(MemWord w, JumpTableSize w)]
go ByteString
bs
        | ByteString -> Bool
BS.null ByteString
bs = []
        | Bool
otherwise =
            let (ByteString
addrBs, ByteString
rest1) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
ptrBytes ByteString
bs
                (ByteString
cntBs, ByteString
rest2) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
ptrBytes ByteString
rest1
            in (Word64 -> MemWord w
forall (w :: Nat). MemWidth w => Word64 -> MemWord w
memWord (ByteString -> Word64
readWord ByteString
addrBs), MemWord w -> JumpTableSize w
forall (w :: Nat). MemWord w -> JumpTableSize w
JumpTableSize (Word64 -> MemWord w
forall (w :: Nat). MemWidth w => Word64 -> MemWord w
memWord (ByteString -> Word64
readWord ByteString
cntBs))) (MemWord w, JumpTableSize w)
-> [(MemWord w, JumpTableSize w)] -> [(MemWord w, JumpTableSize w)]
forall a. a -> [a] -> [a]
: ByteString -> [(MemWord w, JumpTableSize w)]
MemWidth w => ByteString -> [(MemWord w, JumpTableSize w)]
go ByteString
rest2
  in if Int
totalLen Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
entryBytes Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
       then SectionSizeError
-> Either SectionSizeError [(MemWord w, JumpTableSize w)]
forall a b. a -> Either a b
Left (Int -> Int -> SectionSizeError
SectionSizeError Int
totalLen Int
entryBytes)
       else [(MemWord w, JumpTableSize w)]
-> Either SectionSizeError [(MemWord w, JumpTableSize w)]
forall a b. b -> Either a b
Right (ByteString -> [(MemWord w, JumpTableSize w)]
MemWidth w => ByteString -> [(MemWord w, JumpTableSize w)]
go ByteString
bs0)

-- | Read the @.llvm_jump_table_sizes@ section out of an ELF binary, parse
-- it, and resolve each address against the provided 'Memory'.
--
-- Returns a 'SectionSizeError' if the section is malformed, otherwise a list
-- of per-entry warnings (one per unresolvable address) along with a map from
-- resolved jump-table base addresses to entry counts.
--
-- If the section is absent the result is @Right ([], Map.empty)@.
llvmJumpTableSizesFromElf
  :: forall w
  .  Elf.ElfHeaderInfo w
  -> Memory w
  -> Either SectionSizeError ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
llvmJumpTableSizesFromElf :: forall (w :: Nat).
ElfHeaderInfo w
-> Memory w
-> Either
     SectionSizeError
     ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
llvmJumpTableSizesFromElf ElfHeaderInfo w
ehi Memory w
mem =
  let end :: Endianness
end = case ElfHeader w -> ElfData
forall (w :: Nat). ElfHeader w -> ElfData
Elf.headerData (ElfHeaderInfo w -> ElfHeader w
forall (w :: Nat). ElfHeaderInfo w -> ElfHeader w
Elf.header ElfHeaderInfo w
ehi) of
              ElfData
Elf.ELFDATA2LSB -> Endianness
LittleEndian
              ElfData
Elf.ELFDATA2MSB -> Endianness
BigEndian
      sectionBytes :: [ByteString]
sectionBytes =
        [ ElfSection (ElfWordType w) -> ByteString
forall w. ElfSection w -> ByteString
Elf.elfSectionData ElfSection (ElfWordType w)
sec
        | (FileRange (ElfWordType w)
_, ElfSection (ElfWordType w)
sec) <- Vector (FileRange (ElfWordType w), ElfSection (ElfWordType w))
-> [(FileRange (ElfWordType w), ElfSection (ElfWordType w))]
forall a. Vector a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList (ElfHeaderInfo w
-> Vector (FileRange (ElfWordType w), ElfSection (ElfWordType w))
forall (w :: Nat).
ElfHeaderInfo w
-> Vector (FileRange (ElfWordType w), ElfSection (ElfWordType w))
Elf.headerSections ElfHeaderInfo w
ehi)
        , ElfSection (ElfWordType w) -> ByteString
forall w. ElfSection w -> ByteString
Elf.elfSectionName ElfSection (ElfWordType w)
sec ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
sectionName
        ]
      wRepr :: AddrWidthRepr w
      wRepr :: AddrWidthRepr w
wRepr = ElfClass w -> AddrWidthRepr w
forall (w :: Nat). ElfClass w -> AddrWidthRepr w
elfAddrWidth (ElfHeader w -> ElfClass w
forall (w :: Nat). ElfHeader w -> ElfClass w
Elf.headerClass (ElfHeaderInfo w -> ElfHeader w
forall (w :: Nat). ElfHeaderInfo w -> ElfHeader w
Elf.header ElfHeaderInfo w
ehi))
  in AddrWidthRepr w
-> (MemWidth w =>
    Either
      SectionSizeError
      ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w)))
-> Either
     SectionSizeError
     ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
forall (w :: Nat) a. AddrWidthRepr w -> (MemWidth w => a) -> a
addrWidthClass AddrWidthRepr w
wRepr ((MemWidth w =>
  Either
    SectionSizeError
    ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w)))
 -> Either
      SectionSizeError
      ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w)))
-> (MemWidth w =>
    Either
      SectionSizeError
      ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w)))
-> Either
     SectionSizeError
     ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
forall a b. (a -> b) -> a -> b
$ case [ByteString]
sectionBytes of
       [] -> ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
-> Either
     SectionSizeError
     ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
forall a b. b -> Either a b
Right ([], Map (MemSegmentOff w) (JumpTableSize w)
forall k a. Map k a
Map.empty)
       ByteString
bs : [ByteString]
_ -> Memory w
-> [(MemWord w, JumpTableSize w)]
-> ([UnresolvableAddress w],
    Map (MemSegmentOff w) (JumpTableSize w))
forall (w :: Nat).
MemWidth w =>
Memory w
-> [(MemWord w, JumpTableSize w)]
-> ([UnresolvableAddress w],
    Map (MemSegmentOff w) (JumpTableSize w))
resolveEntries Memory w
mem ([(MemWord w, JumpTableSize w)]
 -> ([UnresolvableAddress w],
     Map (MemSegmentOff w) (JumpTableSize w)))
-> Either SectionSizeError [(MemWord w, JumpTableSize w)]
-> Either
     SectionSizeError
     ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AddrWidthRepr w
-> Endianness
-> ByteString
-> Either SectionSizeError [(MemWord w, JumpTableSize w)]
forall (w :: Nat).
AddrWidthRepr w
-> Endianness
-> ByteString
-> Either SectionSizeError [(MemWord w, JumpTableSize w)]
parseLLVMJumpTableSizes AddrWidthRepr w
wRepr Endianness
end ByteString
bs

resolveEntries
  :: MemWidth w
  => Memory w
  -> [(MemWord w, JumpTableSize w)]
  -> ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
resolveEntries :: forall (w :: Nat).
MemWidth w =>
Memory w
-> [(MemWord w, JumpTableSize w)]
-> ([UnresolvableAddress w],
    Map (MemSegmentOff w) (JumpTableSize w))
resolveEntries Memory w
mem = (([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
 -> (MemWord w, JumpTableSize w)
 -> ([UnresolvableAddress w],
     Map (MemSegmentOff w) (JumpTableSize w)))
-> ([UnresolvableAddress w],
    Map (MemSegmentOff w) (JumpTableSize w))
-> [(MemWord w, JumpTableSize w)]
-> ([UnresolvableAddress w],
    Map (MemSegmentOff w) (JumpTableSize w))
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
-> (MemWord w, JumpTableSize w)
-> ([UnresolvableAddress w],
    Map (MemSegmentOff w) (JumpTableSize w))
step ([], Map (MemSegmentOff w) (JumpTableSize w)
forall k a. Map k a
Map.empty)
  where
    step :: ([UnresolvableAddress w], Map (MemSegmentOff w) (JumpTableSize w))
-> (MemWord w, JumpTableSize w)
-> ([UnresolvableAddress w],
    Map (MemSegmentOff w) (JumpTableSize w))
step ([UnresolvableAddress w]
errs, Map (MemSegmentOff w) (JumpTableSize w)
m) (MemWord w
addr, JumpTableSize w
cnt) =
      case Memory w -> MemWord w -> Maybe (MemSegmentOff w)
forall (w :: Nat). Memory w -> MemWord w -> Maybe (MemSegmentOff w)
resolveAbsoluteAddr Memory w
mem MemWord w
addr of
        Maybe (MemSegmentOff w)
Nothing -> (MemWord w -> UnresolvableAddress w
forall (w :: Nat). MemWord w -> UnresolvableAddress w
UnresolvableAddress MemWord w
addr UnresolvableAddress w
-> [UnresolvableAddress w] -> [UnresolvableAddress w]
forall a. a -> [a] -> [a]
: [UnresolvableAddress w]
errs, Map (MemSegmentOff w) (JumpTableSize w)
m)
        Just MemSegmentOff w
segOff -> ([UnresolvableAddress w]
errs, MemSegmentOff w
-> JumpTableSize w
-> Map (MemSegmentOff w) (JumpTableSize w)
-> Map (MemSegmentOff w) (JumpTableSize w)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert MemSegmentOff w
segOff JumpTableSize w
cnt Map (MemSegmentOff w) (JumpTableSize w)
m)