-- Decoder.hs ---

-- Copyright (C) 2020 Nerd Ed

-- Author: Nerd Ed <nerded.nerded@gmail.com>

-- This program is free software; you can redistribute it and/or
-- modify it under the terms of the GNU General Public License
-- as published by the Free Software Foundation; either version 3
-- of the License, or (at your option) any later version.

-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU General Public License for more details.

-- You should have received a copy of the GNU General Public License
-- along with this program. If not, see <http://www.gnu.org/licenses/>.

{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE TypeApplications         #-}

module Zydis.Decoder
  ( ZyanStatus
  , ZyanUSize
  , Offset
  , Length
  , initialize
  , decodeBuffer
  , decodeFullBuffer
  )
where

import           Data.Bits
import           Data.ByteString               as BS
import           Data.ByteString.Internal
import           Data.Vector
import           Data.Word
import           Foreign.ForeignPtr
import           Foreign.Marshal
import           Foreign.Ptr
import           Foreign.Storable
import           Zydis.Types

-- * FFI types

type MachineModeC = Word32

type AddressWidthC = Word32

type ZyanStatus = Word32

type ZyanUSize = Word64

type Offset = ZyanUSize

type Length = ZyanUSize

-- * FFI declarations

foreign import ccall unsafe "ZydisDecoderInit" c_ZydisDecoderInit
  :: Ptr Decoder -> MachineModeC -> AddressWidthC -> IO ZyanStatus

foreign import ccall unsafe "ZydisDecoderDecodeBuffer" c_ZydisDecoderDecodeBuffer
  :: Ptr Decoder -> Ptr Word8 -> ZyanUSize -> Ptr DecodedInstruction -> IO ZyanStatus

-- * FFI bridges

-- | Directly stolen from https://github.com/zyantific/zycore-c/blob/71440fa634d1313db735d3262d453be641bb404f/include/Zycore/Status.h#L81
zyanSuccess :: Word32 -> Bool
zyanSuccess :: Word32 -> Bool
zyanSuccess Word32
x = (Word32
x Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
0x80000000) Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
0
{-# INLINE zyanSuccess #-}

-- | Initialize a Zydis decoder, required to decode instructions.
initialize :: MachineMode -> AddressWidth -> IO (Either ZyanStatus Decoder)
initialize :: MachineMode -> AddressWidth -> IO (Either Word32 Decoder)
initialize MachineMode
mm AddressWidth
aw = (Ptr Decoder -> IO (Either Word32 Decoder))
-> IO (Either Word32 Decoder)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca Ptr Decoder -> IO (Either Word32 Decoder)
go
 where
  go :: Ptr Decoder -> IO (Either Word32 Decoder)
go Ptr Decoder
decoder = do
    Word32
r <- Ptr Decoder -> Word32 -> Word32 -> IO Word32
c_ZydisDecoderInit Ptr Decoder
decoder
                            (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word32) -> Int -> Word32
forall a b. (a -> b) -> a -> b
$ MachineMode -> Int
forall a. Enum a => a -> Int
fromEnum MachineMode
mm)
                            (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word32) -> Int -> Word32
forall a b. (a -> b) -> a -> b
$ AddressWidth -> Int
forall a. Enum a => a -> Int
fromEnum AddressWidth
aw)
    if Word32 -> Bool
zyanSuccess Word32
r then Decoder -> Either Word32 Decoder
forall a b. b -> Either a b
Right (Decoder -> Either Word32 Decoder)
-> IO Decoder -> IO (Either Word32 Decoder)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr Decoder -> IO Decoder
forall a. Storable a => Ptr a -> IO a
peek Ptr Decoder
decoder else Either Word32 Decoder -> IO (Either Word32 Decoder)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Word32 Decoder -> IO (Either Word32 Decoder))
-> Either Word32 Decoder -> IO (Either Word32 Decoder)
forall a b. (a -> b) -> a -> b
$ Word32 -> Either Word32 Decoder
forall a b. a -> Either a b
Left Word32
r
{-# INLINE initialize #-}

-- | Decode a single intruction.
decodeBuffer
  :: Decoder
  -> ByteString
  -> Offset
  -> Length
  -> IO (Either ZyanStatus DecodedInstruction)
decodeBuffer :: Decoder
-> ByteString
-> Offset
-> Offset
-> IO (Either Word32 DecodedInstruction)
decodeBuffer Decoder
d ByteString
bs Offset
o Offset
l = (Ptr Decoder -> IO (Either Word32 DecodedInstruction))
-> IO (Either Word32 DecodedInstruction)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca @Decoder Ptr Decoder -> IO (Either Word32 DecodedInstruction)
go
 where
  (ForeignPtr Word8
bufferForeignPtr, Int
_, Int
_) = ByteString -> (ForeignPtr Word8, Int, Int)
toForeignPtr ByteString
bs

  go :: Ptr Decoder -> IO (Either Word32 DecodedInstruction)
go Ptr Decoder
decoderPtr = forall b.
Storable DecodedInstruction =>
(Ptr DecodedInstruction -> IO b) -> IO b
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca @DecodedInstruction ((Ptr DecodedInstruction -> IO (Either Word32 DecodedInstruction))
 -> IO (Either Word32 DecodedInstruction))
-> (Ptr DecodedInstruction
    -> IO (Either Word32 DecodedInstruction))
-> IO (Either Word32 DecodedInstruction)
forall a b. (a -> b) -> a -> b
$ Ptr Decoder
-> Ptr DecodedInstruction -> IO (Either Word32 DecodedInstruction)
go' Ptr Decoder
decoderPtr

  go' :: Ptr Decoder
-> Ptr DecodedInstruction -> IO (Either Word32 DecodedInstruction)
go' Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr =
    ForeignPtr Word8
-> (Ptr Word8 -> IO (Either Word32 DecodedInstruction))
-> IO (Either Word32 DecodedInstruction)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
bufferForeignPtr ((Ptr Word8 -> IO (Either Word32 DecodedInstruction))
 -> IO (Either Word32 DecodedInstruction))
-> (Ptr Word8 -> IO (Either Word32 DecodedInstruction))
-> IO (Either Word32 DecodedInstruction)
forall a b. (a -> b) -> a -> b
$ Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> IO (Either Word32 DecodedInstruction)
go'' Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr

  go'' :: Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> IO (Either Word32 DecodedInstruction)
go'' Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr Ptr Word8
bufferPtr = do
    Ptr Decoder -> Decoder -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr Decoder
decoderPtr Decoder
d
    Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> Offset
-> Offset
-> IO (Either Word32 DecodedInstruction)
doDecodeInstruction Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr Ptr Word8
bufferPtr Offset
o Offset
l
{-# INLINE decodeBuffer #-}

-- | Efficiently decode an entire buffer of instructions.
decodeFullBuffer
  :: Decoder -> ByteString -> IO (Either ZyanStatus (Vector DecodedInstruction))
decodeFullBuffer :: Decoder
-> ByteString -> IO (Either Word32 (Vector DecodedInstruction))
decodeFullBuffer Decoder
d ByteString
bs = (Ptr Decoder -> IO (Either Word32 (Vector DecodedInstruction)))
-> IO (Either Word32 (Vector DecodedInstruction))
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca @Decoder Ptr Decoder -> IO (Either Word32 (Vector DecodedInstruction))
go
 where
  (ForeignPtr Word8
bufferForeignPtr, Int
_, Int
bufferLength) = ByteString -> (ForeignPtr Word8, Int, Int)
toForeignPtr ByteString
bs

  go :: Ptr Decoder -> IO (Either Word32 (Vector DecodedInstruction))
go Ptr Decoder
decoderPtr = forall b.
Storable DecodedInstruction =>
(Ptr DecodedInstruction -> IO b) -> IO b
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca @DecodedInstruction ((Ptr DecodedInstruction
  -> IO (Either Word32 (Vector DecodedInstruction)))
 -> IO (Either Word32 (Vector DecodedInstruction)))
-> (Ptr DecodedInstruction
    -> IO (Either Word32 (Vector DecodedInstruction)))
-> IO (Either Word32 (Vector DecodedInstruction))
forall a b. (a -> b) -> a -> b
$ Ptr Decoder
-> Ptr DecodedInstruction
-> IO (Either Word32 (Vector DecodedInstruction))
go' Ptr Decoder
decoderPtr

  go' :: Ptr Decoder
-> Ptr DecodedInstruction
-> IO (Either Word32 (Vector DecodedInstruction))
go' Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr =
    ForeignPtr Word8
-> (Ptr Word8 -> IO (Either Word32 (Vector DecodedInstruction)))
-> IO (Either Word32 (Vector DecodedInstruction))
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
bufferForeignPtr ((Ptr Word8 -> IO (Either Word32 (Vector DecodedInstruction)))
 -> IO (Either Word32 (Vector DecodedInstruction)))
-> (Ptr Word8 -> IO (Either Word32 (Vector DecodedInstruction)))
-> IO (Either Word32 (Vector DecodedInstruction))
forall a b. (a -> b) -> a -> b
$ Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> IO (Either Word32 (Vector DecodedInstruction))
go'' Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr

  go'' :: Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> IO (Either Word32 (Vector DecodedInstruction))
go'' Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr Ptr Word8
bufferPtr = do
    Ptr Decoder -> Decoder -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr Decoder
decoderPtr Decoder
d
    (Vector DecodedInstruction, Offset, Offset)
-> IO (Either Word32 (Vector DecodedInstruction))
forall (f :: * -> *).
(Semigroup (f DecodedInstruction), Applicative f) =>
(f DecodedInstruction, Offset, Offset)
-> IO (Either Word32 (f DecodedInstruction))
loop (Monoid (Vector DecodedInstruction) => Vector DecodedInstruction
forall a. Monoid a => a
mempty @(Vector DecodedInstruction), Offset
0, Int -> Offset
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bufferLength)
   where
    loop :: (f DecodedInstruction, Offset, Offset)
-> IO (Either Word32 (f DecodedInstruction))
loop (!f DecodedInstruction
v, !Offset
o, !Offset
l)
      | Offset
l Offset -> Offset -> Bool
forall a. Ord a => a -> a -> Bool
> Offset
0 = do
        Either Word32 DecodedInstruction
x <- Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> Offset
-> Offset
-> IO (Either Word32 DecodedInstruction)
doDecodeInstruction Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr Ptr Word8
bufferPtr Offset
o Offset
l
        case Either Word32 DecodedInstruction
x of
          Right DecodedInstruction
i -> do
            let il :: Offset
il = Word8 -> Offset
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Offset) -> Word8 -> Offset
forall a b. (a -> b) -> a -> b
$ DecodedInstruction -> Word8
decodedInstructionLength DecodedInstruction
i
            (f DecodedInstruction, Offset, Offset)
-> IO (Either Word32 (f DecodedInstruction))
loop (f DecodedInstruction
v f DecodedInstruction
-> f DecodedInstruction -> f DecodedInstruction
forall a. Semigroup a => a -> a -> a
<> DecodedInstruction -> f DecodedInstruction
forall (f :: * -> *) a. Applicative f => a -> f a
pure DecodedInstruction
i, Offset
o Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
+ Offset
il, Offset
l Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
- Offset
il)
          Left Word32
s -> Either Word32 (f DecodedInstruction)
-> IO (Either Word32 (f DecodedInstruction))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Word32 (f DecodedInstruction)
 -> IO (Either Word32 (f DecodedInstruction)))
-> Either Word32 (f DecodedInstruction)
-> IO (Either Word32 (f DecodedInstruction))
forall a b. (a -> b) -> a -> b
$ Word32 -> Either Word32 (f DecodedInstruction)
forall a b. a -> Either a b
Left Word32
s
      | Bool
otherwise = Either Word32 (f DecodedInstruction)
-> IO (Either Word32 (f DecodedInstruction))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Word32 (f DecodedInstruction)
 -> IO (Either Word32 (f DecodedInstruction)))
-> Either Word32 (f DecodedInstruction)
-> IO (Either Word32 (f DecodedInstruction))
forall a b. (a -> b) -> a -> b
$ f DecodedInstruction -> Either Word32 (f DecodedInstruction)
forall a b. b -> Either a b
Right f DecodedInstruction
v
{-# INLINE decodeFullBuffer #-}

doDecodeInstruction
  :: Ptr Decoder
  -> Ptr DecodedInstruction
  -> Ptr Word8
  -> Offset
  -> Length
  -> IO (Either ZyanStatus DecodedInstruction)
doDecodeInstruction :: Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> Offset
-> Offset
-> IO (Either Word32 DecodedInstruction)
doDecodeInstruction Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr Ptr Word8
bufferPtr Offset
o Offset
l = do
  Word32
r <- Ptr Decoder
-> Ptr Word8 -> Offset -> Ptr DecodedInstruction -> IO Word32
c_ZydisDecoderDecodeBuffer Ptr Decoder
decoderPtr
                                  (Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
bufferPtr (Offset -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Offset
o))
                                  Offset
l
                                  Ptr DecodedInstruction
decodedInstructionPtr
  if Word32 -> Bool
zyanSuccess Word32
r then DecodedInstruction -> Either Word32 DecodedInstruction
forall a b. b -> Either a b
Right (DecodedInstruction -> Either Word32 DecodedInstruction)
-> IO DecodedInstruction -> IO (Either Word32 DecodedInstruction)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr DecodedInstruction -> IO DecodedInstruction
forall a. Storable a => Ptr a -> IO a
peek Ptr DecodedInstruction
decodedInstructionPtr else Either Word32 DecodedInstruction
-> IO (Either Word32 DecodedInstruction)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Word32 DecodedInstruction
 -> IO (Either Word32 DecodedInstruction))
-> Either Word32 DecodedInstruction
-> IO (Either Word32 DecodedInstruction)
forall a b. (a -> b) -> a -> b
$ Word32 -> Either Word32 DecodedInstruction
forall a b. a -> Either a b
Left Word32
r
{-# INLINE doDecodeInstruction #-}