Network/DNS/Packer.hs

   1 module Network.DNS.Packer
   2     ( Packer
   3
   4     , pack
   5     , pack'
   6
   7     , getState
   8     , setState
   9     , modifyState
  10
  11     , bytesWrote
  12     , withOffset
  13
  14     , putByteString
  15     , putLazyByteString
  16
  17     , putWord8
  18     , putWord16be
  19     , putWord32be
  20
  21     , putBinary
  22     , liftToBinary
  23     )
  24     where
  25
  26 import qualified Data.Binary as Binary
  27 import qualified Data.Binary.Put as Bin
  28 import           Data.Bits
  29 import qualified Data.ByteString as Strict
  30 import qualified Data.ByteString.Lazy as Lazy
  31 import           Data.Int
  32 import           Data.Word
  33
  34
  35 data PackingState s
  36     = PackingState {
  37         stResult     :: !Lazy.ByteString
  38       , stBytesWrote :: !Int64
  39       , stUserState  :: s
  40       }
  41
  42 newtype Packer s a = P { unP :: PackingState s -> (a, PackingState s) }
  43
  44 instance Monad (Packer s) where
  45     return a = P (\ s -> (a, s))
  46     m >>= k  = P (\ s -> let (a, s') = unP m s
  47                          in
  48                            unP (k a) s')
  49     fail err = do bytes <- get stBytesWrote
  50                   P (error (err
  51                             ++ ". Failed packing at byte position "
  52                             ++ show bytes))
  53
  54 get :: (PackingState s -> a) -> Packer s a
  55 get f = P (\ s -> (f s, s))
  56
  57 set :: (PackingState s -> PackingState s) -> Packer s ()
  58 set f = P (\ s -> ((), f s))
  59
  60 mkState :: Lazy.ByteString -> Int64 -> s -> PackingState s
  61 mkState xs n s
  62     = PackingState {
  63         stResult     = xs
  64       , stBytesWrote = n
  65       , stUserState  = s
  66       }
  67
  68 pack' :: Packer s a -> s -> (Lazy.ByteString, s, a)
  69 pack' m s
  70     = let (a, s') = unP m (mkState Lazy.empty 0 s)
  71       in
  72         (stResult s', stUserState s', a)
  73
  74 pack :: Packer s a -> s -> Lazy.ByteString
  75 pack = (fst' .) . pack'
  76     where
  77       fst' (xs, _, _) = xs
  78
  79 getState :: Packer s s
  80 getState = get stUserState
  81
  82 setState :: s -> Packer s ()
  83 setState = modifyState . const
  84
  85 modifyState :: (s -> s) -> Packer s ()
  86 modifyState f
  87     = set $ \ st -> st { stUserState = f (stUserState st) }
  88
  89 bytesWrote :: Integral i => Packer s i
  90 bytesWrote = get stBytesWrote >>= return . fromIntegral
  91
  92 withOffset :: Int64 -> Packer s a -> Packer s a
  93 withOffset n m
  94     = P $ \ s -> let (taken, dropped) = Lazy.splitAt n (stResult s)
  95                      padded           = Lazy.take n (taken `Lazy.append` Lazy.repeat 0)
  96                      tempState        = s {
  97                                           stResult     = padded
  98                                         , stBytesWrote = stBytesWrote s - Lazy.length dropped
  99                                         }
 100                      (a, tempState')  = unP m tempState
 101                      newState         = tempState {
 102                                           stResult     = replaceHead (stResult s) (stResult tempState')
 103                                         , stBytesWrote = max (stBytesWrote s) (stBytesWrote tempState')
 104                                         }
 105                  in
 106                    (a, newState)
 107       where
 108         replaceHead :: Lazy.ByteString -> Lazy.ByteString -> Lazy.ByteString
 109         replaceHead world newHead
 110             = let rest = Lazy.drop (Lazy.length newHead) world
 111               in
 112                 newHead `Lazy.append` rest
 113
 114
 115 putByteString :: Strict.ByteString -> Packer s ()
 116 putByteString = putLazyByteString . Lazy.fromChunks . (:[])
 117
 118 putLazyByteString :: Lazy.ByteString -> Packer s ()
 119 putLazyByteString xs
 120     = set $ \ st -> st {
 121                       stResult     = stResult st `Lazy.append` xs
 122                     , stBytesWrote = stBytesWrote st + Lazy.length xs
 123                     }
 124
 125 putWord8 :: Word8 -> Packer s ()
 126 putWord8 w
 127     = set $ \ st -> st {
 128                       stResult     = stResult st `Lazy.snoc` w
 129                     , stBytesWrote = stBytesWrote st + 1
 130                     }
 131
 132 putWord16be :: Word16 -> Packer s ()
 133 putWord16be w
 134     = do putWord8 $ fromIntegral $ (w `shiftR`  8) .&. 0xFF
 135          putWord8 $ fromIntegral $  w              .&. 0xFF
 136
 137 putWord32be :: Word32 -> Packer s ()
 138 putWord32be w
 139     = do putWord8 $ fromIntegral $ (w `shiftR` 24) .&. 0xFF
 140          putWord8 $ fromIntegral $ (w `shiftR` 16) .&. 0xFF
 141          putWord8 $ fromIntegral $ (w `shiftR`  8) .&. 0xFF
 142          putWord8 $ fromIntegral $  w              .&. 0xFF
 143
 144
 145 putBinary :: Binary.Binary a => a -> Packer s ()
 146 putBinary = putLazyByteString . Binary.encode
 147
 148
 149 liftToBinary :: s -> Packer s a -> Bin.PutM a
 150 liftToBinary s m
 151     = do let (a, s') = unP m (mkState Lazy.empty 0 s)
 152
 153          Bin.putLazyByteString (stResult s')
 154          return a