-- 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
  , ZydisStatus(..)
  , ZyanCoreStatus(..)
  , ZyanUSize
  , Offset
  , Length
  , initialize
  , decodeBuffer
  , decodeFullBuffer
  )
where

import           Data.ByteString               as BS
import           Data.ByteString.Internal
import           Data.Sequence
import           Data.Word
import           Foreign.ForeignPtr
import           Foreign.Marshal
import           Foreign.Ptr
import           Foreign.Storable

import           Zydis.Types
import           Zydis.Status

-- * FFI types

type MachineModeC = Word32

type AddressWidthC = Word32

type ZyanUSize = Word64

type Offset = ZyanUSize

type Length = ZyanUSize

-- * FFI declarations

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

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

-- * FFI bridges

-- | Initialize a Zydis decoder, required to decode instructions.
initialize :: MachineMode -> AddressWidth -> IO (Either ZyanStatus Decoder)
initialize :: MachineMode -> AddressWidth -> IO (Either ZyanStatus Decoder)
initialize MachineMode
mm AddressWidth
aw = (Ptr Decoder -> IO (Either ZyanStatus Decoder))
-> IO (Either ZyanStatus Decoder)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca Ptr Decoder -> IO (Either ZyanStatus Decoder)
go
 where
  go :: Ptr Decoder -> IO (Either ZyanStatus Decoder)
go Ptr Decoder
decoderPtr = do
    ZyanNativeStatus
r <- Ptr Decoder
-> ZyanNativeStatus -> ZyanNativeStatus -> IO ZyanNativeStatus
c_ZydisDecoderInit Ptr Decoder
decoderPtr
                            (Int -> ZyanNativeStatus
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> ZyanNativeStatus) -> Int -> ZyanNativeStatus
forall a b. (a -> b) -> a -> b
$ MachineMode -> Int
forall a. Enum a => a -> Int
fromEnum MachineMode
mm)
                            (Int -> ZyanNativeStatus
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> ZyanNativeStatus) -> Int -> ZyanNativeStatus
forall a b. (a -> b) -> a -> b
$ AddressWidth -> Int
forall a. Enum a => a -> Int
fromEnum AddressWidth
aw)
    case ZyanNativeStatus -> ZyanStatus
fromZyanNativeStatus ZyanNativeStatus
r of
      Left ZyanCoreStatus
ZyanCoreStatusSuccess -> Decoder -> Either ZyanStatus Decoder
forall a b. b -> Either a b
Right (Decoder -> Either ZyanStatus Decoder)
-> IO Decoder -> IO (Either ZyanStatus 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
decoderPtr
      ZyanStatus
x                          -> Either ZyanStatus Decoder -> IO (Either ZyanStatus Decoder)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either ZyanStatus Decoder -> IO (Either ZyanStatus Decoder))
-> Either ZyanStatus Decoder -> IO (Either ZyanStatus Decoder)
forall a b. (a -> b) -> a -> b
$ ZyanStatus -> Either ZyanStatus Decoder
forall a b. a -> Either a b
Left ZyanStatus
x
{-# INLINE initialize #-}

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

  go :: Ptr Decoder -> IO (Either ZyanStatus 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 ZyanStatus DecodedInstruction))
 -> IO (Either ZyanStatus DecodedInstruction))
-> (Ptr DecodedInstruction
    -> IO (Either ZyanStatus DecodedInstruction))
-> IO (Either ZyanStatus DecodedInstruction)
forall a b. (a -> b) -> a -> b
$ Ptr Decoder
-> Ptr DecodedInstruction
-> IO (Either ZyanStatus DecodedInstruction)
go' Ptr Decoder
decoderPtr

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

  go'' :: Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> IO (Either ZyanStatus 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 ZyanStatus 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 (Seq DecodedInstruction))
decodeFullBuffer :: Decoder
-> ByteString -> IO (Either ZyanStatus (Seq DecodedInstruction))
decodeFullBuffer Decoder
d ByteString
bs = (Ptr Decoder -> IO (Either ZyanStatus (Seq DecodedInstruction)))
-> IO (Either ZyanStatus (Seq DecodedInstruction))
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca @Decoder Ptr Decoder -> IO (Either ZyanStatus (Seq DecodedInstruction))
go
 where
  (ForeignPtr Word8
bufferForeignPtr, Int
_, Int
bufferLength) = ByteString -> (ForeignPtr Word8, Int, Int)
toForeignPtr ByteString
bs

  go :: Ptr Decoder -> IO (Either ZyanStatus (Seq DecodedInstruction))
go = 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 ZyanStatus (Seq DecodedInstruction)))
 -> IO (Either ZyanStatus (Seq DecodedInstruction)))
-> (Ptr Decoder
    -> Ptr DecodedInstruction
    -> IO (Either ZyanStatus (Seq DecodedInstruction)))
-> Ptr Decoder
-> IO (Either ZyanStatus (Seq DecodedInstruction))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr Decoder
-> Ptr DecodedInstruction
-> IO (Either ZyanStatus (Seq DecodedInstruction))
go'

  go' :: Ptr Decoder
-> Ptr DecodedInstruction
-> IO (Either ZyanStatus (Seq DecodedInstruction))
go' Ptr Decoder
decoderPtr = ForeignPtr Word8
-> (Ptr Word8 -> IO (Either ZyanStatus (Seq DecodedInstruction)))
-> IO (Either ZyanStatus (Seq DecodedInstruction))
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
bufferForeignPtr ((Ptr Word8 -> IO (Either ZyanStatus (Seq DecodedInstruction)))
 -> IO (Either ZyanStatus (Seq DecodedInstruction)))
-> (Ptr DecodedInstruction
    -> Ptr Word8 -> IO (Either ZyanStatus (Seq DecodedInstruction)))
-> Ptr DecodedInstruction
-> IO (Either ZyanStatus (Seq DecodedInstruction))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> IO (Either ZyanStatus (Seq DecodedInstruction))
go'' Ptr Decoder
decoderPtr

  go'' :: Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> IO (Either ZyanStatus (Seq 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
    Seq DecodedInstruction
-> Offset
-> Offset
-> IO (Either ZyanStatus (Seq DecodedInstruction))
loop Seq DecodedInstruction
forall a. Monoid a => a
mempty Offset
0 (Int -> Offset
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bufferLength)
   where
    loop :: Seq DecodedInstruction
-> Offset
-> Offset
-> IO (Either ZyanStatus (Seq DecodedInstruction))
loop !Seq DecodedInstruction
v !Offset
o !Offset
l
      | Offset
l Offset -> Offset -> Bool
forall a. Ord a => a -> a -> Bool
> Offset
0 = do
        Either ZyanStatus DecodedInstruction
x <- Ptr Decoder
-> Ptr DecodedInstruction
-> Ptr Word8
-> Offset
-> Offset
-> IO (Either ZyanStatus DecodedInstruction)
doDecodeInstruction Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr Ptr Word8
bufferPtr Offset
o Offset
l
        case Either ZyanStatus 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
            Seq DecodedInstruction
-> Offset
-> Offset
-> IO (Either ZyanStatus (Seq DecodedInstruction))
loop (Seq DecodedInstruction
v Seq DecodedInstruction
-> DecodedInstruction -> Seq DecodedInstruction
forall a. Seq a -> a -> Seq a
:|> 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 ZyanStatus
s -> Either ZyanStatus (Seq DecodedInstruction)
-> IO (Either ZyanStatus (Seq DecodedInstruction))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either ZyanStatus (Seq DecodedInstruction)
 -> IO (Either ZyanStatus (Seq DecodedInstruction)))
-> Either ZyanStatus (Seq DecodedInstruction)
-> IO (Either ZyanStatus (Seq DecodedInstruction))
forall a b. (a -> b) -> a -> b
$ ZyanStatus -> Either ZyanStatus (Seq DecodedInstruction)
forall a b. a -> Either a b
Left ZyanStatus
s
      | Bool
otherwise = Either ZyanStatus (Seq DecodedInstruction)
-> IO (Either ZyanStatus (Seq DecodedInstruction))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either ZyanStatus (Seq DecodedInstruction)
 -> IO (Either ZyanStatus (Seq DecodedInstruction)))
-> Either ZyanStatus (Seq DecodedInstruction)
-> IO (Either ZyanStatus (Seq DecodedInstruction))
forall a b. (a -> b) -> a -> b
$ Seq DecodedInstruction
-> Either ZyanStatus (Seq DecodedInstruction)
forall a b. b -> Either a b
Right Seq 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 ZyanStatus DecodedInstruction)
doDecodeInstruction Ptr Decoder
decoderPtr Ptr DecodedInstruction
decodedInstructionPtr Ptr Word8
bufferPtr Offset
o Offset
l = do
  ZyanNativeStatus
r <- Ptr Decoder
-> Ptr Word8
-> Offset
-> Ptr DecodedInstruction
-> IO ZyanNativeStatus
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
  case ZyanNativeStatus -> ZyanStatus
fromZyanNativeStatus ZyanNativeStatus
r of
    Left ZyanCoreStatus
ZyanCoreStatusSuccess -> DecodedInstruction -> Either ZyanStatus DecodedInstruction
forall a b. b -> Either a b
Right (DecodedInstruction -> Either ZyanStatus DecodedInstruction)
-> IO DecodedInstruction
-> IO (Either ZyanStatus 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
    ZyanStatus
x                          -> Either ZyanStatus DecodedInstruction
-> IO (Either ZyanStatus DecodedInstruction)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either ZyanStatus DecodedInstruction
 -> IO (Either ZyanStatus DecodedInstruction))
-> Either ZyanStatus DecodedInstruction
-> IO (Either ZyanStatus DecodedInstruction)
forall a b. (a -> b) -> a -> b
$ ZyanStatus -> Either ZyanStatus DecodedInstruction
forall a b. a -> Either a b
Left ZyanStatus
x
{-# INLINE doDecodeInstruction #-}