]> gitweb @ CieloNegro.org - haskell-dns.git/blob - Network/DNS/Unpacker.hs
Many changes...
[haskell-dns.git] / 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           Data.Bits
30 import qualified Data.ByteString as Strict
31 import qualified Data.ByteString.Lazy as Lazy
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