blob: 05f8468aac817b58ec1afbf6ffb6d60e6d45c374 (
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 BL (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 -> BL.ByteString
runSPut = toLazyByteString . fromWrite
|
