{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

-- | An implementation of an /Executable Loadable Format/ (ELF) loader. The module
-- is responsible for loading instructions into a provided memory implementation
-- and obtaining the entry point for the executable.
module LibRISCV.Loader (readElf, LoadFunc, loadElf, startAddr) where

import Control.Monad.Catch (MonadCatch)
import Data.Bits ()
import qualified Data.ByteString.Lazy as BSL
import Data.Elf (
    Elf (..),
    ElfListXX (..),
    ElfNodeType (..),
    ElfSectionData (ElfSectionData),
    ElfXX (
        ElfSection,
        ElfSegment,
        ehEntry,
        epAddMemSize,
        epData,
        epType,
        esAddr,
        esData
    ),
    elfFindHeader,
    parseElf,
 )
import Data.Elf.Constants
import Data.Elf.Headers (
    SingElfClass (SELFCLASS32, SELFCLASS64),
    SingElfClassI,
    withSingElfClassI,
 )
import Data.Elf.PrettyPrint (readFileLazy)
import Data.Int (Int64)
import Data.Word (Word32)
import LibRISCV
import System.FilePath ()

-- Filter all ELF segments with type PT_LOAD.
loadableSegments :: ElfListXX a -> [ElfXX 'Segment a]
loadableSegments :: forall (a :: ElfClass). ElfListXX a -> [ElfXX 'Segment a]
loadableSegments (ElfListCons v :: ElfXX t a
v@(ElfSegment{ElfSegmentType
WordXX a
ElfListXX a
epAddMemSize :: forall (c :: ElfClass). ElfXX 'Segment c -> WordXX c
epData :: forall (c :: ElfClass). ElfXX 'Segment c -> ElfListXX c
epType :: forall (c :: ElfClass). ElfXX 'Segment c -> ElfSegmentType
epType :: ElfSegmentType
epAddMemSize :: WordXX a
epData :: ElfListXX a
..}) ElfListXX a
l) =
    if ElfSegmentType
epType ElfSegmentType -> ElfSegmentType -> Bool
forall a. Eq a => a -> a -> Bool
== ElfSegmentType
PT_LOAD
        then ElfXX t a
ElfXX 'Segment a
v ElfXX 'Segment a -> [ElfXX 'Segment a] -> [ElfXX 'Segment a]
forall a. a -> [a] -> [a]
: ElfListXX a -> [ElfXX 'Segment a]
forall (a :: ElfClass). ElfListXX a -> [ElfXX 'Segment a]
loadableSegments ElfListXX a
l
        else ElfListXX a -> [ElfXX 'Segment a]
forall (a :: ElfClass). ElfListXX a -> [ElfXX 'Segment a]
loadableSegments ElfListXX a
l
loadableSegments (ElfListCons ElfXX t a
_ ElfListXX a
l) = ElfListXX a -> [ElfXX 'Segment a]
forall (a :: ElfClass). ElfListXX a -> [ElfXX 'Segment a]
loadableSegments ElfListXX a
l
loadableSegments ElfListXX a
ElfListNull = []

-- Copy data from ElfSection to memory at the given absolute address.
copyData :: (Monad m, SingElfClassI a) => ElfListXX a -> Int64 -> LoadFunc m -> m ()
copyData :: forall (m :: * -> *) (a :: ElfClass).
(Monad m, SingElfClassI a) =>
ElfListXX a -> Int64 -> LoadFunc m -> m ()
copyData ElfListXX a
ElfListNull Int64
_ LoadFunc m
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
copyData (ElfListCons (ElfSection{esData :: forall (c :: ElfClass). ElfXX 'Section c -> ElfSectionData c
esData = ElfSectionData ByteString
textData, WordXX a
esAddr :: forall (c :: ElfClass). ElfXX 'Section c -> WordXX c
esAddr :: WordXX a
..}) ElfListXX a
xs) Int64
zeros LoadFunc m
f = do
    LoadFunc m
f (WordXX a -> Address
forall a b. (Integral a, Num b) => a -> b
fromIntegral WordXX a
esAddr) (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString -> ByteString
BSL.append ByteString
textData (Int64 -> Word8 -> ByteString
BSL.replicate Int64
zeros Word8
0)
    ElfListXX a -> Int64 -> LoadFunc m -> m ()
forall (m :: * -> *) (a :: ElfClass).
(Monad m, SingElfClassI a) =>
ElfListXX a -> Int64 -> LoadFunc m -> m ()
copyData ElfListXX a
xs Int64
zeros LoadFunc m
f
copyData (ElfListCons ElfXX t a
_ ElfListXX a
xs) Int64
zeros LoadFunc m
f = ElfListXX a -> Int64 -> LoadFunc m -> m ()
forall (m :: * -> *) (a :: ElfClass).
(Monad m, SingElfClassI a) =>
ElfListXX a -> Int64 -> LoadFunc m -> m ()
copyData ElfListXX a
xs Int64
zeros LoadFunc m
f

-- Load an ElfSegment into memory at the given address.
loadSegment :: (Monad m, SingElfClassI a) => LoadFunc m -> ElfXX 'Segment a -> m ()
loadSegment :: forall (m :: * -> *) (a :: ElfClass).
(Monad m, SingElfClassI a) =>
LoadFunc m -> ElfXX 'Segment a -> m ()
loadSegment LoadFunc m
loadFunc ElfSegment{ElfSegmentType
WordXX a
ElfListXX a
epAddMemSize :: forall (c :: ElfClass). ElfXX 'Segment c -> WordXX c
epData :: forall (c :: ElfClass). ElfXX 'Segment c -> ElfListXX c
epType :: forall (c :: ElfClass). ElfXX 'Segment c -> ElfSegmentType
epType :: ElfSegmentType
epAddMemSize :: WordXX a
epData :: ElfListXX a
..} =
    ElfListXX a -> Int64 -> LoadFunc m -> m ()
forall (m :: * -> *) (a :: ElfClass).
(Monad m, SingElfClassI a) =>
ElfListXX a -> Int64 -> LoadFunc m -> m ()
copyData ElfListXX a
epData (WordXX a -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral WordXX a
epAddMemSize) LoadFunc m
loadFunc

------------------------------------------------------------------------

-- | Load a 'BSL.ByteString' into memory at a given address.
type LoadFunc m = Address -> BSL.ByteString -> m ()

-- | Load all loadable segments of an ELF file into memory. An addition to the
-- 'Data.Elf.Elf' file, it requires an implementation of a 'LoadFunc' which is
-- responsible for converting a 'BSL.ByteString' to the internal value
-- representation.
loadElf :: (Monad m) => Elf -> LoadFunc m -> m ()
loadElf :: forall (m :: * -> *). Monad m => Elf -> LoadFunc m -> m ()
loadElf (Elf SingElfClass a
classS ElfListXX a
elfs) LoadFunc m
loadFunc = SingElfClass a -> (SingElfClassI a => m ()) -> m ()
forall (c :: ElfClass) r.
SingElfClass c -> (SingElfClassI c => r) -> r
withSingElfClassI SingElfClass a
classS ((SingElfClassI a => m ()) -> m ())
-> (SingElfClassI a => m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ do
    let loadable :: [ElfXX 'Segment a]
loadable = ElfListXX a -> [ElfXX 'Segment a]
forall (a :: ElfClass). ElfListXX a -> [ElfXX 'Segment a]
loadableSegments ElfListXX a
elfs
    (ElfXX 'Segment a -> m ()) -> [ElfXX 'Segment a] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (LoadFunc m -> ElfXX 'Segment a -> m ()
forall (m :: * -> *) (a :: ElfClass).
(Monad m, SingElfClassI a) =>
LoadFunc m -> ElfXX 'Segment a -> m ()
loadSegment LoadFunc m
loadFunc) [ElfXX 'Segment a]
loadable

-- | Read an ELF from a given 'FilePath'.
readElf :: FilePath -> IO Elf
readElf :: FilePath -> IO Elf
readElf FilePath
path = FilePath -> IO ByteString
readFileLazy FilePath
path IO ByteString -> (ByteString -> IO Elf) -> IO Elf
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ByteString -> IO Elf
forall (m :: * -> *). MonadCatch m => ByteString -> m Elf
parseElf

-- | Return the entry point from the ELF header.
startAddr :: (MonadCatch m) => Elf -> m Word32
startAddr :: forall (m :: * -> *). MonadCatch m => Elf -> m Address
startAddr (Elf SingElfClass a
SELFCLASS32 ElfListXX a
elfs) = ElfXX 'Header a -> Address
ElfXX 'Header a -> WordXX a
forall (c :: ElfClass). ElfXX 'Header c -> WordXX c
ehEntry (ElfXX 'Header a -> Address) -> m (ElfXX 'Header a) -> m Address
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ElfListXX a -> m (ElfXX 'Header a)
forall (a :: ElfClass) (m :: * -> *).
(SingElfClassI a, MonadThrow m) =>
ElfListXX a -> m (ElfXX 'Header a)
elfFindHeader ElfListXX a
elfs
startAddr (Elf SingElfClass a
SELFCLASS64 ElfListXX a
_) = FilePath -> m Address
forall a. HasCallStack => FilePath -> a
error FilePath
"64-bit executables not supported"