Abstract the string table into a newtype

2024-10-22 01:21:02 +02:00
parent bf5ec1f3e2
commit 57436f664e
6 changed files with 228 additions and 152 deletions
--- a/lib/Language/Elna/Object/ElfCoder.hs
+++ b/lib/Language/Elna/Object/ElfCoder.hs
@@ -0,0 +1,128 @@
+-- | Object file generation.
+module Language.Elna.Object.ElfCoder
+    ( ElfEnvironment(..)
+    , ElfWriter(..)
+    , ElfHeaderResult(..)
+    , elfHeaderSize
+    , addSectionHeader
+    , elfObject
+    , elfSectionsSize
+    , putSectionHeader
+    , module Language.Elna.Object.Elf
+    ) where
+
+import Control.Exception (throwIO)
+import Control.Monad.IO.Class (MonadIO(..))
+import Control.Monad.Trans.State (StateT, runStateT, modify', gets)
+import Data.ByteString (StrictByteString)
+import qualified Data.ByteString as ByteString
+import qualified Data.ByteString.Builder as ByteString.Builder
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import System.IO (Handle, IOMode(..), SeekMode(..), hSeek, withFile)
+import Data.Foldable (traverse_)
+import Language.Elna.Object.Elf
+import Language.Elna.Object.StringTable (StringTable)
+import qualified Language.Elna.Object.StringTable as StringTable
+import GHC.Records (HasField(..))
+
+data ElfEnvironment = ElfEnvironment
+    { objectHeaders :: ElfHeaderResult Elf32_Shdr
+    , objectHandle :: Handle
+    }
+
+newtype ElfWriter a = ElfWriter
+    { runElfWriter :: StateT ElfEnvironment IO a
+    }
+
+data ElfHeaderResult a = ElfHeaderResult
+    { sectionNames :: StringTable
+    , sectionHeaders :: Vector a
+    } deriving Eq
+
+instance Functor ElfWriter
+  where
+    fmap f (ElfWriter x) = ElfWriter $ f <$> x
+
+instance Applicative ElfWriter
+  where
+    pure = ElfWriter . pure
+    (ElfWriter f) <*> (ElfWriter x) = ElfWriter $ f <*> x
+
+instance Monad ElfWriter
+  where
+    (ElfWriter x) >>= f = ElfWriter $ x >>= (runElfWriter . f)
+
+instance MonadIO ElfWriter
+  where
+    liftIO = ElfWriter . liftIO
+
+-- | ELF header size.
+elfHeaderSize :: Elf32_Off
+elfHeaderSize = 52
+
+-- | Calculates the size of all sections based on the 'sh_size' in the given
+-- headers and adds 'elfHeaderSize' to it.
+elfSectionsSize :: Vector Elf32_Shdr -> Elf32_Off
+elfSectionsSize = (elfHeaderSize +)
+    . Vector.foldr ((+) . sh_size) 0
+
+addSectionHeader :: StrictByteString -> Elf32_Shdr -> ElfWriter ()
+addSectionHeader name newHeader = ElfWriter $ modify' modifier
+  where
+    modifier elfEnvironment@ElfEnvironment{ objectHeaders } =
+        let ElfHeaderResult{..} = objectHeaders
+         in elfEnvironment
+            { objectHeaders = ElfHeaderResult
+                { sectionHeaders = Vector.snoc sectionHeaders newHeader
+                , sectionNames = StringTable.append name sectionNames
+                }
+            }
+
+putSectionHeader :: StrictByteString -> Elf32_Shdr -> StrictByteString -> ElfWriter ()
+putSectionHeader name newHeader encoded = do
+    objectHandle' <- ElfWriter $ gets $ getField @"objectHandle"
+    liftIO $ ByteString.hPut objectHandle' encoded
+    addSectionHeader name newHeader
+
+-- Writes an ELF object to the provided file path. The callback writes the
+-- sections, collects headers for those sections and returns the ELF header.
+elfObject :: FilePath -> ElfWriter Elf32_Ehdr -> IO ()
+elfObject outFile putContents = withFile outFile WriteMode withObjectFile
+  where
+    withObjectFile objectHandle
+        = hSeek objectHandle AbsoluteSeek (fromIntegral elfHeaderSize)
+        >> putContents' objectHandle
+        >>= uncurry afterContents
+    putContents' objectHandle
+        = flip runStateT (initialState objectHandle)
+        $ runElfWriter putContents
+    zeroHeader = Elf32_Shdr
+        { sh_type = SHT_NULL
+        , sh_size = 0
+        , sh_offset = 0
+        , sh_name = 0
+        , sh_link = 0
+        , sh_info = 0
+        , sh_flags = 0
+        , sh_entsize = 0
+        , sh_addralign = 0
+        , sh_addr = 0
+        }
+    initialState objectHandle = ElfEnvironment
+        { objectHeaders = ElfHeaderResult
+            { sectionHeaders = Vector.singleton zeroHeader
+            , sectionNames = mempty
+            }
+        , objectHandle = objectHandle
+        }
+    afterContents header ElfEnvironment{ objectHeaders = ElfHeaderResult{..}, ..} =
+        let hPutBuilder = ByteString.Builder.hPutBuilder objectHandle
+            writeSectionHeaders byteOrder' =
+                traverse_ (hPutBuilder . elf32Shdr byteOrder') sectionHeaders
+         in either throwIO pure (byteOrder (e_ident header))
+            >>= writeSectionHeaders
+            >> either throwIO (putHeaders objectHandle) (elf32Ehdr header)
+    putHeaders objectHandle encodedHeader
+        = hSeek objectHandle AbsoluteSeek 0
+        >> ByteString.Builder.hPutBuilder objectHandle encodedHeader