Network/DNS/Unpacker.hs

   1 module Network.DNS.Unpacker
   2     ( Unpacker
   3     , UnpackingState(..)
   4
   5     , unpack
   6     , unpack'
   7
   8     , getState
   9     , setState
  10     , modifyState
  11
  12     , skip
  13     , lookAhead
  14     , bytesRead
  15
  16     , getByteString
  17     , getLazyByteString
  18     , getWord8
  19     , getWord16be
  20     , getWord32be
  21
  22     , getBinary
  23     , liftToBinary
  24     )
  25     where
  26
  27 import qualified Data.Binary as Binary
  28 import qualified Data.Binary.Get as Bin
  29 import qualified Data.ByteString as Strict
  30 import qualified Data.ByteString.Lazy as Lazy
  31 import           Data.Bits
  32 import           Data.Int
  33 import           Data.Word
  34
  35
  36 data UnpackingState s
  37     = UnpackingState {
  38         stSource    :: !Lazy.ByteString
  39       , stBytesRead :: !Int64
  40       , stUserState :: s
  41       }
  42
  43 newtype Unpacker s a = U { unU :: UnpackingState s -> (a, UnpackingState s) }
  44
  45 instance Monad (Unpacker s) where
  46     return a = U (\ s -> (a, s))
  47     m >>= k  = U (\ s -> let (a, s') = unU m s
  48                          in
  49                            unU (k a) s')
  50     fail err = do bytes <- get stBytesRead
  51                   U (error (err
  52                             ++ ". Failed unpacking at byte position "
  53                             ++ show bytes))
  54
  55 get :: (UnpackingState s -> a) -> Unpacker s a
  56 get f = U (\ s -> (f s, s))
  57
  58 set :: (UnpackingState s -> UnpackingState s) -> Unpacker s ()
  59 set f = U (\ s -> ((), f s))
  60
  61 mkState :: Lazy.ByteString -> Int64 -> s -> UnpackingState s
  62 mkState xs n s
  63     = UnpackingState {
  64         stSource    = xs
  65       , stBytesRead = n
  66       , stUserState = s
  67       }
  68
  69 unpack' :: Unpacker s a -> s -> Lazy.ByteString -> (a, s)
  70 unpack' m s xs
  71     = let (a, s') = unU m (mkState xs 0 s)
  72       in
  73         (a, stUserState s')
  74
  75 unpack :: Unpacker s a -> s -> Lazy.ByteString -> a
  76 unpack = ((fst .) .) . unpack'
  77
  78 getState :: Unpacker s s
  79 getState = get stUserState
  80
  81 setState :: s -> Unpacker s ()
  82 setState = modifyState . const
  83
  84 modifyState :: (s -> s) -> Unpacker s ()
  85 modifyState f
  86     = set $ \ st -> st { stUserState = f (stUserState st) }
  87
  88 skip :: Int64 -> Unpacker s ()
  89 skip n = getLazyByteString n >> return ()
  90
  91 lookAhead :: Unpacker s a -> Unpacker s a
  92 lookAhead m = U (\ s -> let (a, _) = unU m s
  93                         in
  94                           (a, s))
  95
  96 bytesRead :: Integral i => Unpacker s i
  97 bytesRead = get stBytesRead >>= return . fromIntegral
  98
  99 getByteString :: Int -> Unpacker s Strict.ByteString
 100 getByteString n = getLazyByteString (fromIntegral n) >>= return . Strict.concat . Lazy.toChunks
 101
 102 getLazyByteString :: Int64 -> Unpacker s Lazy.ByteString
 103 getLazyByteString n
 104     = do src <- get stSource
 105          let (xs, ys) = Lazy.splitAt n src
 106          if Lazy.length xs /= n then
 107              fail "Too few bytes"
 108            else
 109              do set $ \ st -> st {
 110                                 stSource    = ys
 111                               , stBytesRead = stBytesRead st + n
 112                               }
 113                 return xs
 114
 115 getWord8 :: Unpacker s Word8
 116 getWord8 = getLazyByteString 1 >>= return . (`Lazy.index` 0)
 117
 118 getWord16be :: Unpacker s Word16
 119 getWord16be = do xs <- getLazyByteString 2
 120                  return $ (fromIntegral (xs `Lazy.index` 0) `shiftL` 8) .|.
 121                           (fromIntegral (xs `Lazy.index` 1))
 122
 123 getWord32be :: Unpacker s Word32
 124 getWord32be = do xs <- getLazyByteString 4
 125                  return $ (fromIntegral (xs `Lazy.index` 0) `shiftL` 24) .|.
 126                           (fromIntegral (xs `Lazy.index` 1) `shiftL` 16) .|.
 127                           (fromIntegral (xs `Lazy.index` 2) `shiftL`  8) .|.
 128                           (fromIntegral (xs `Lazy.index` 3))
 129
 130 getBinary :: Binary.Binary a => Unpacker s a
 131 getBinary = do s <- get id
 132                let (a, rest, bytes) = Bin.runGetState Binary.get (stSource s) (stBytesRead s)
 133                set $ \ st -> st {
 134                                stSource    = rest
 135                              , stBytesRead = bytes
 136                              }
 137                return a
 138
 139
 140 liftToBinary :: s -> Unpacker s a -> Bin.Get a
 141 liftToBinary s m
 142     = do bytes <- Bin.bytesRead
 143          src   <- Bin.getRemainingLazyByteString
 144
 145          let (a, s') = unU m (mkState src bytes s)
 146
 147          -- These bytes was consumed by the unpacker.
 148          Bin.skip (fromIntegral (stBytesRead s' - bytes))
 149
 150          return a