{-# LANGUAGE
BangPatterns
, RecordWildCards
, TransformListComp
#-}
module Network.DNS.Encode.Builders (
putDNSMessage
, putDNSFlags
, putHeader
, putDomain
, putMailbox
, putResourceRecord
) where
import Control.Monad.State (State, modify, execState, gets)
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Char8 as BS
import qualified Data.ByteString.Lazy.Char8 as LBS
import qualified Data.IP
import Data.List (foldl')
import Data.IP (IP(..), fromIPv4, fromIPv6b, makeAddrRange)
import GHC.Exts (the, groupWith)
import Network.DNS.Imports
import Network.DNS.StateBinary
import Network.DNS.Types
putDNSMessage :: DNSMessage -> SPut
putDNSMessage msg = putHeader hd
<> putNums
<> mconcat (map putQuestion qs)
<> mconcat (map putResourceRecord an)
<> mconcat (map putResourceRecord au)
<> mconcat (map putResourceRecord ad)
where
putNums = mconcat $ fmap putInt16 [ length qs
, length an
, length au
, length ad
]
hm = header msg
fl = flags hm
eh = ednsHeader msg
qs = question msg
an = answer msg
au = authority msg
hd = ifEDNS eh hm $ hm { flags = fl { rcode = rc } }
rc = ifEDNS eh <$> id <*> nonEDNSrcode $ rcode fl
where
nonEDNSrcode code | fromRCODE code < 16 = code
| otherwise = FormatErr
ad = prependOpt $ additional msg
where
prependOpt ads = mapEDNS eh (fromEDNS ads $ fromRCODE rc) ads
where
fromEDNS :: AdditionalRecords -> Word16 -> EDNS -> AdditionalRecords
fromEDNS rrs rc' edns = ResourceRecord name' type' class' ttl' rdata' : rrs
where
name' = BS.singleton '.'
type' = OPT
class' = maxUdpSize `min` (minUdpSize `max` ednsUdpSize edns)
ttl0' = fromIntegral (rc' .&. 0xff0) `shiftL` 20
vers' = fromIntegral (ednsVersion edns) `shiftL` 16
ttl'
| ednsDnssecOk edns = ttl0' `setBit` 15 .|. vers'
| otherwise = ttl0' .|. vers'
rdata' = RD_OPT $ ednsOptions edns
putHeader :: DNSHeader -> SPut
putHeader hdr = putIdentifier (identifier hdr)
<> putDNSFlags (flags hdr)
where
putIdentifier = put16
putDNSFlags :: DNSFlags -> SPut
putDNSFlags DNSFlags{..} = put16 word
where
set :: Word16 -> State Word16 ()
set byte = modify (.|. byte)
st :: State Word16 ()
st = sequence_
[ set (fromRCODE rcode .&. 0x0f)
, when chkDisable $ set (bit 4)
, when authenData $ set (bit 5)
, when recAvailable $ set (bit 7)
, when recDesired $ set (bit 8)
, when trunCation $ set (bit 9)
, when authAnswer $ set (bit 10)
, set (fromOPCODE opcode `shiftL` 11)
, when (qOrR==QR_Response) $ set (bit 15)
]
word = execState st 0
putQuestion :: Question -> SPut
putQuestion Question{..} = putDomain qname
<> put16 (fromTYPE qtype)
<> put16 classIN
putResourceRecord :: ResourceRecord -> SPut
putResourceRecord ResourceRecord{..} = mconcat [
putDomain rrname
, put16 (fromTYPE rrtype)
, put16 rrclass
, put32 rrttl
, putResourceRData rdata
]
where
putResourceRData :: RData -> SPut
putResourceRData rd = do
addPositionW 2
rDataBuilder <- putRData rd
let rdataLength = fromIntegral . LBS.length . BB.toLazyByteString $ rDataBuilder
let rlenBuilder = BB.int16BE rdataLength
return $ rlenBuilder <> rDataBuilder
putRData :: RData -> SPut
putRData rd = case rd of
RD_A address -> mconcat $ map putInt8 (fromIPv4 address)
RD_NS nsdname -> putDomain nsdname
RD_CNAME cname -> putDomain cname
RD_SOA a b c d e f g -> putSOA a b c d e f g
RD_NULL bytes -> putByteString bytes
RD_PTR ptrdname -> putDomain ptrdname
RD_MX pref exch -> mconcat [put16 pref, putDomain exch]
RD_TXT textstring -> putTXT textstring
RD_AAAA address -> mconcat $ map putInt8 (fromIPv6b address)
RD_SRV pri wei prt tgt -> putSRV pri wei prt tgt
RD_DNAME dname -> putDomain dname
RD_OPT options -> mconcat $ fmap putOData options
RD_DS kt ka dt d -> putDS kt ka dt d
RD_RRSIG rrsig -> putRRSIG rrsig
RD_NSEC next types -> putDomain next <> putNsecTypes types
RD_DNSKEY f p alg key -> putDNSKEY f p alg key
RD_NSEC3 a f i s h types -> putNSEC3 a f i s h types
RD_NSEC3PARAM a f iter salt -> putNSEC3PARAM a f iter salt
RD_TLSA u s m dgst -> putTLSA u s m dgst
UnknownRData bytes -> putByteString bytes
where
putSOA mn mr serial refresh retry expire minttl = mconcat
[ putDomain mn
, putMailbox mr
, put32 serial
, put32 refresh
, put32 retry
, put32 expire
, put32 minttl
]
putTXT textstring =
let (!h, !t) = BS.splitAt 255 textstring
in putByteStringWithLength h <> if BS.null t
then mempty
else putTXT t
putSRV priority weight port target = mconcat
[ put16 priority
, put16 weight
, put16 port
, putDomain target
]
putDS keytag keyalg digestType digest = mconcat
[ put16 keytag
, put8 keyalg
, put8 digestType
, putByteString digest
]
putRRSIG RDREP_RRSIG{..} = mconcat
[ put16 $ fromTYPE rrsigType
, put8 rrsigKeyAlg
, put8 rrsigNumLabels
, put32 rrsigTTL
, put32 $ fromIntegral rrsigExpiration
, put32 $ fromIntegral rrsigInception
, put16 rrsigKeyTag
, putDomain rrsigZone
, putByteString rrsigValue
]
putDNSKEY flags protocol alg key = mconcat
[ put16 flags
, put8 protocol
, put8 alg
, putByteString key
]
putNSEC3 alg flags iterations salt hash types = mconcat
[ put8 alg
, put8 flags
, put16 iterations
, putByteStringWithLength salt
, putByteStringWithLength hash
, putNsecTypes types
]
putNSEC3PARAM alg flags iterations salt = mconcat
[ put8 alg
, put8 flags
, put16 iterations
, putByteStringWithLength salt
]
putTLSA usage selector mtype assocData = mconcat
[ put8 usage
, put8 selector
, put8 mtype
, putByteString assocData
]
putNsecTypes :: [TYPE] -> SPut
putNsecTypes types = putTypeList $ map fromTYPE types
where
putTypeList :: [Word16] -> SPut
putTypeList ts =
mconcat [ putWindow (the top8) bot8 |
t <- ts,
let top8 = fromIntegral t `shiftR` 8,
let bot8 = fromIntegral t .&. 0xff,
then group by top8
using groupWith ]
putWindow :: Int -> [Int] -> SPut
putWindow top8 bot8s =
let blks = maximum bot8s `shiftR` 3
in putInt8 top8
<> put8 (1 + fromIntegral blks)
<> putBits 0 [ (the block, foldl' mergeBits 0 bot8) |
bot8 <- bot8s,
let block = bot8 `shiftR` 3,
then group by block
using groupWith ]
where
mergeBits acc b = setBit acc (7 - b.&.0x07)
putBits :: Int -> [(Int, Word8)] -> SPut
putBits _ [] = pure mempty
putBits n ((block, octet) : rest) =
putReplicate (block-n) 0
<> put8 octet
<> putBits (block + 1) rest
putODWords :: Word16 -> [Word8] -> SPut
putODWords code ws =
mconcat [ put16 code
, putInt16 $ length ws
, mconcat $ map put8 ws
]
putODBytes :: Word16 -> ByteString -> SPut
putODBytes code bs =
mconcat [ put16 code
, putInt16 $ BS.length bs
, putByteString bs
]
putOData :: OData -> SPut
putOData (OD_NSID nsid) = putODBytes (fromOptCode NSID) nsid
putOData (OD_DAU as) = putODWords (fromOptCode DAU) as
putOData (OD_DHU hs) = putODWords (fromOptCode DHU) hs
putOData (OD_N3U hs) = putODWords (fromOptCode N3U) hs
putOData (OD_ClientSubnet srcBits scpBits ip) =
let octets = fromIntegral $ (srcBits + 7) `div` 8
prefix addr = Data.IP.addr $ makeAddrRange addr $ fromIntegral srcBits
(family, raw) = case ip of
IPv4 ip4 -> (1, take octets $ fromIPv4 $ prefix ip4)
IPv6 ip6 -> (2, take octets $ fromIPv6b $ prefix ip6)
dataLen = 2 + 2 + octets
in mconcat [ put16 $ fromOptCode ClientSubnet
, putInt16 dataLen
, put16 family
, put8 srcBits
, put8 scpBits
, mconcat $ fmap putInt8 raw
]
putOData (OD_ECSgeneric family srcBits scpBits addr) =
mconcat [ put16 $ fromOptCode ClientSubnet
, putInt16 $ 4 + BS.length addr
, put16 family
, put8 srcBits
, put8 scpBits
, putByteString addr
]
putOData (UnknownOData code bs) = putODBytes code bs
putByteStringWithLength :: BS.ByteString -> SPut
putByteStringWithLength bs = putInt8 (fromIntegral $ BS.length bs)
<> putByteString bs
rootDomain :: Domain
rootDomain = BS.pack "."
putDomain :: Domain -> SPut
putDomain = putDomain' '.'
putMailbox :: Mailbox -> SPut
putMailbox = putDomain' '@'
putDomain' :: Char -> ByteString -> SPut
putDomain' sep dom
| BS.null dom || dom == rootDomain = put8 0
| otherwise = do
mpos <- wsPop dom
cur <- gets wsPosition
case mpos of
Just pos -> putPointer pos
Nothing -> wsPush dom cur >>
mconcat [ putPartialDomain hd
, putDomain' '.' tl
]
where
(hd, tl) =
let p = parseLabel (c2w sep) dom
in if sep /= '.' && BS.null (snd p)
then parseLabel (c2w '.') dom
else p
c2w = fromIntegral . fromEnum
putPointer :: Int -> SPut
putPointer pos = putInt16 (pos .|. 0xc000)
putPartialDomain :: Domain -> SPut
putPartialDomain = putByteStringWithLength