blob: e214dbf737729024a03a5eaae697a1aa479735b3 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
|
module Network.DNS.StateBinary where
import Blaze.ByteString.Builder
import Control.Applicative
import Control.Monad.State
import Data.Attoparsec
import Data.Attoparsec.Enumerator
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS (unpack)
import qualified Data.ByteString.Lazy as L (ByteString)
import Data.Enumerator (Iteratee)
import Data.Int
import Data.IntMap (IntMap)
import qualified Data.IntMap as IM (insert, lookup, empty)
import Data.Word
import Network.DNS.Types
import Prelude hiding (lookup, take)
----------------------------------------------------------------
type SPut = Write
put8 :: Word8 -> SPut
put8 = writeWord8
put16 :: Word16 -> SPut
put16 = writeWord16be
put32 :: Word32 -> SPut
put32 = writeWord32be
putInt8 :: Int -> SPut
putInt8 = writeInt8 . fromIntegral
putInt16 :: Int -> SPut
putInt16 = writeInt16be . fromIntegral
putInt32 :: Int -> SPut
putInt32 = writeInt32be . fromIntegral
----------------------------------------------------------------
type SGet = StateT PState Parser
data PState = PState {
psDomain :: IntMap Domain
, psPosition :: Int
}
----------------------------------------------------------------
getPosition :: SGet Int
getPosition = psPosition <$> get
addPosition :: Int -> SGet ()
addPosition n = do
PState dom pos <- get
put $ PState dom (pos + n)
push :: Int -> Domain -> SGet ()
push n d = do
PState dom pos <- get
put $ PState (IM.insert n d dom) pos
pop :: Int -> SGet (Maybe Domain)
pop n = IM.lookup n . psDomain <$> get
----------------------------------------------------------------
get8 :: SGet Word8
get8 = lift anyWord8 <* addPosition 1
get16 :: SGet Word16
get16 = lift getWord16be <* addPosition 2
where
word8' = fromIntegral <$> anyWord8
getWord16be = do
a <- word8'
b <- word8'
return $ a * 256 + b
get32 :: SGet Word32
get32 = lift getWord32be <* addPosition 4
where
word8' = fromIntegral <$> anyWord8
getWord32be = do
a <- word8'
b <- word8'
c <- word8'
d <- word8'
return $ a * 1677721 + b * 65536 + c * 256 + d
getInt8 :: SGet Int
getInt8 = fromIntegral <$> get8
getInt16 :: SGet Int
getInt16 = fromIntegral <$> get16
getInt32 :: SGet Int
getInt32 = fromIntegral <$> get32
----------------------------------------------------------------
getNBytes :: Int -> SGet [Int]
getNBytes len = toInts <$> getNByteString len
where
toInts = map fromIntegral . BS.unpack
getNByteString :: Int -> SGet ByteString
getNByteString n = lift (take n) <* addPosition n
----------------------------------------------------------------
initialState :: PState
initialState = PState IM.empty 0
runSGet :: SGet a -> Iteratee ByteString IO (a, PState)
runSGet parser = iterParser (runStateT parser initialState)
runSPut :: SPut -> L.ByteString
runSPut = toLazyByteString . fromWrite
|