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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
|
{-# LANGUAGE RecordWildCards #-}
module Network.DNS.Query (composeQuery, composeDNSFormat) where
import qualified Data.ByteString.Lazy.Char8 as BL (ByteString)
import qualified Data.ByteString.Char8 as BS (length, null, break, drop)
import Network.DNS.StateBinary
import Network.DNS.Internal
import Data.Monoid
import Control.Monad.State
import Data.Bits
import Data.Word
import Data.IP
(+++) :: Monoid a => a -> a -> a
(+++) = mappend
----------------------------------------------------------------
composeDNSFormat :: DNSFormat -> BL.ByteString
composeDNSFormat fmt = runSPut (encodeDNSFormat fmt)
composeQuery :: Int -> [Question] -> BL.ByteString
composeQuery idt qs = composeDNSFormat qry
where
hdr = header defaultQuery
qry = defaultQuery {
header = hdr {
identifier = idt
, qdCount = length qs
}
, question = qs
}
----------------------------------------------------------------
encodeDNSFormat :: DNSFormat -> SPut
encodeDNSFormat fmt = encodeHeader hdr
+++ mconcat (map encodeQuestion qs)
+++ mconcat (map encodeRR an)
+++ mconcat (map encodeRR au)
+++ mconcat (map encodeRR ad)
where
hdr = header fmt
qs = question fmt
an = answer fmt
au = authority fmt
ad = additional fmt
encodeHeader :: DNSHeader -> SPut
encodeHeader hdr = encodeIdentifier (identifier hdr)
+++ encodeFlags (flags hdr)
+++ decodeQdCount (qdCount hdr)
+++ decodeAnCount (anCount hdr)
+++ decodeNsCount (nsCount hdr)
+++ decodeArCount (arCount hdr)
where
encodeIdentifier = putInt16
decodeQdCount = putInt16
decodeAnCount = putInt16
decodeNsCount = putInt16
decodeArCount = putInt16
encodeFlags :: DNSFlags -> SPut
encodeFlags DNSFlags{..} = put16 word
where
word16 :: Enum a => a -> Word16
word16 = toEnum . fromEnum
set :: Word16 -> State Word16 ()
set byte = modify (.|. byte)
st :: State Word16 ()
st = sequence_
[ set (word16 rcode)
, when recAvailable $ set (bit 7)
, when recDesired $ set (bit 8)
, when trunCation $ set (bit 9)
, when authAnswer $ set (bit 10)
, set (word16 opcode `shiftL` 11)
, when (qOrR==QR_Response) $ set (bit 15)
]
word = execState st 0
encodeQuestion :: Question -> SPut
encodeQuestion Question{..} =
encodeDomain qname
+++ putInt16 (typeToInt qtype)
+++ put16 1
encodeRR :: ResourceRecord -> SPut
encodeRR ResourceRecord{..} =
mconcat
[ encodeDomain rrname
, putInt16 (typeToInt rrtype)
, put16 1
, putInt32 rrttl
, putInt16 rdlen
, encodeRDATA rdata
]
encodeRDATA :: RDATA -> SPut
encodeRDATA rd = case rd of
(RD_A ip) -> mconcat $ map putInt8 (fromIPv4 ip)
(RD_AAAA ip) -> mconcat $ map putInt16 (fromIPv6 ip)
(RD_NS dom) -> encodeDomain dom
(RD_CNAME dom) -> encodeDomain dom
(RD_PTR dom) -> encodeDomain dom
(RD_MX prf dom) -> mconcat [putInt16 prf, encodeDomain dom]
(RD_TXT txt) -> putByteString txt
(RD_OTH bytes) -> mconcat $ map putInt8 bytes
(RD_SOA d1 d2 serial refresh retry expire min') -> mconcat $
[ encodeDomain d1
, encodeDomain d2
, putInt32 serial
, putInt32 refresh
, putInt32 retry
, putInt32 expire
, putInt32 min'
]
(RD_SRV prio weight port dom) -> mconcat $
[ putInt16 prio
, putInt16 weight
, putInt16 port
, encodeDomain dom
]
----------------------------------------------------------------
encodeDomain :: Domain -> SPut
encodeDomain dom | BS.null dom = put8 0
encodeDomain dom = do
mpos <- wsPop dom
cur <- gets wsPosition
case mpos of
Just pos -> encodePointer pos
Nothing -> wsPush dom cur >>
mconcat [ encodePartialDomain hd
, encodeDomain tl
]
where
(hd, tl') = BS.break (=='.') dom
tl = if BS.null tl' then tl' else BS.drop 1 tl'
encodePointer :: Int -> SPut
encodePointer pos = let w = (pos .|. 0xc000) in putInt16 w
encodePartialDomain :: Domain -> SPut
encodePartialDomain sub = putInt8 (BS.length sub)
+++ putByteString sub
|