1 module Network.DNS.Packer
26 import qualified Data.Binary as Binary
27 import qualified Data.Binary.Put as Bin
29 import qualified Data.ByteString as Strict
30 import qualified Data.ByteString.Lazy as Lazy
37 stResult :: !Lazy.ByteString
38 , stBytesWrote :: !Int64
42 newtype Packer s a = P { unP :: PackingState s -> (a, PackingState s) }
44 instance Monad (Packer s) where
45 return a = P (\ s -> (a, s))
46 m >>= k = P (\ s -> let (a, s') = unP m s
49 fail err = do bytes <- get stBytesWrote
51 ++ ". Failed packing at byte position "
54 get :: (PackingState s -> a) -> Packer s a
55 get f = P (\ s -> (f s, s))
57 set :: (PackingState s -> PackingState s) -> Packer s ()
58 set f = P (\ s -> ((), f s))
60 mkState :: Lazy.ByteString -> Int64 -> s -> PackingState s
68 pack' :: Packer s a -> s -> (Lazy.ByteString, s, a)
70 = let (a, s') = unP m (mkState Lazy.empty 0 s)
72 (stResult s', stUserState s', a)
74 pack :: Packer s a -> s -> Lazy.ByteString
75 pack = (fst' .) . pack'
79 getState :: Packer s s
80 getState = get stUserState
82 setState :: s -> Packer s ()
83 setState = modifyState . const
85 modifyState :: (s -> s) -> Packer s ()
87 = set $ \ st -> st { stUserState = f (stUserState st) }
89 bytesWrote :: Integral i => Packer s i
90 bytesWrote = get stBytesWrote >>= return . fromIntegral
92 withOffset :: Int64 -> Packer s a -> Packer s a
94 = P $ \ s -> let (taken, dropped) = Lazy.splitAt n (stResult s)
95 padded = Lazy.take n (taken `Lazy.append` Lazy.repeat 0)
98 , stBytesWrote = stBytesWrote s - Lazy.length dropped
100 (a, tempState') = unP m tempState
101 newState = tempState {
102 stResult = replaceHead (stResult s) (stResult tempState')
103 , stBytesWrote = max (stBytesWrote s) (stBytesWrote tempState')
108 replaceHead :: Lazy.ByteString -> Lazy.ByteString -> Lazy.ByteString
109 replaceHead world newHead
110 = let rest = Lazy.drop (Lazy.length newHead) world
112 newHead `Lazy.append` rest
115 putByteString :: Strict.ByteString -> Packer s ()
116 putByteString = putLazyByteString . Lazy.fromChunks . (:[])
118 putLazyByteString :: Lazy.ByteString -> Packer s ()
121 stResult = stResult st `Lazy.append` xs
122 , stBytesWrote = stBytesWrote st + Lazy.length xs
125 putWord8 :: Word8 -> Packer s ()
128 stResult = stResult st `Lazy.snoc` w
129 , stBytesWrote = stBytesWrote st + 1
132 putWord16be :: Word16 -> Packer s ()
134 = do putWord8 $ fromIntegral $ (w `shiftR` 8) .&. 0xFF
135 putWord8 $ fromIntegral $ w .&. 0xFF
137 putWord32be :: Word32 -> Packer s ()
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
145 putBinary :: Binary.Binary a => a -> Packer s ()
146 putBinary = putLazyByteString . Binary.encode
149 liftToBinary :: s -> Packer s a -> Bin.PutM a
151 = do let (a, s') = unP m (mkState Lazy.empty 0 s)
153 Bin.putLazyByteString (stResult s')