Network/DNS/Message.hs

   1 module Network.DNS.Message
   2     ( Message(..)
   3     , MessageID
   4     , MessageType(..)
   5     , Header(..)
   6     , Opcode(..)
   7     , ResponseCode(..)
   8     , Question(..)
   9     , ResourceRecord(..)
  10     , DomainName
  11     , DomainLabel
  12     , TTL
  13     , SomeRR(..)
  14     , RecordType(..)
  15     , RecordClass(..)
  16
  17     , CNAME(..)
  18     , HINFO(..)
  19     )
  20     where
  21
  22 import           Control.Monad
  23 import           Data.Binary
  24 import           Data.Binary.BitPut as BP
  25 import           Data.Binary.Get as G
  26 import           Data.Binary.Put as P
  27 import           Data.Binary.Strict.BitGet as BG
  28 import qualified Data.ByteString as BS
  29 import qualified Data.ByteString.Lazy as LBS
  30 import           Data.Typeable
  31 import qualified Data.IntMap as IM
  32 import           Data.IntMap (IntMap)
  33 import           Data.Word
  34
  35
  36 replicateM' :: Monad m => Int -> (a -> m (b, a)) -> a -> m ([b], a)
  37 replicateM' = worker []
  38     where
  39       worker :: Monad m => [b] -> Int -> (a -> m (b, a)) -> a -> m ([b], a)
  40       worker soFar 0 _ a = return (reverse soFar, a)
  41       worker soFar n f a = do (b, a') <- f a
  42                               worker (b : soFar) (n - 1) f a'
  43
  44
  45 data Message
  46     = Message {
  47         msgHeader      :: !Header
  48       , msgQuestions   :: ![Question]
  49       , msgAnswers     :: ![SomeRR]
  50       , msgAuthorities :: ![SomeRR]
  51       , msgAdditionals :: ![SomeRR]
  52       }
  53
  54 data Header
  55     = Header {
  56         hdMessageID             :: !MessageID
  57       , hdMessageType           :: !MessageType
  58       , hdOpcode                :: !Opcode
  59       , hdIsAuthoritativeAnswer :: !Bool
  60       , hdIsTruncated           :: !Bool
  61       , hdIsRecursionDesired    :: !Bool
  62       , hdIsRecursionAvailable  :: !Bool
  63       , hdResponseCode          :: !ResponseCode
  64
  65       -- These fields are supressed in this data structure:
  66       -- + QDCOUNT
  67       -- + ANCOUNT
  68       -- + NSCOUNT
  69       -- + ARCOUNT
  70       }
  71
  72 type MessageID = Word16
  73
  74 data MessageType
  75     = Query
  76     | Response
  77     deriving (Show, Eq)
  78
  79 data Opcode
  80     = StandardQuery
  81     | InverseQuery
  82     | ServerStatusRequest
  83     deriving (Show, Eq)
  84
  85 data ResponseCode
  86     = NoError
  87     | FormatError
  88     | ServerFailure
  89     | NameError
  90     | NotImplemented
  91     | Refused
  92     deriving (Show, Eq)
  93
  94 data Question
  95     = Question {
  96         qName  :: !DomainName
  97       , qType  :: !RecordType
  98       , qClass :: !RecordClass
  99       }
 100     deriving (Show, Eq)
 101
 102 putQ :: Question -> Put
 103 putQ q
 104     = do putDomainName $ qName q
 105          put $ qType  q
 106          put $ qClass q
 107
 108 getQ :: DecompTable -> Get (Question, DecompTable)
 109 getQ dt
 110     = do (nm, dt') <- getDomainName dt
 111          ty        <- get
 112          cl        <- get
 113          let q = Question {
 114                    qName  = nm
 115                  , qType  = ty
 116                  , qClass = cl
 117                  }
 118          return (q, dt')
 119
 120 type DomainName  = [DomainLabel]
 121 type DomainLabel = BS.ByteString
 122
 123 data RecordClass
 124     = IN
 125     | CS -- Obsolete
 126     | CH
 127     | HS
 128     | AnyClass -- Only for queries
 129     deriving (Show, Eq)
 130
 131 class (Typeable rr, Show rr, Eq rr) => ResourceRecord rr where
 132     rrName    :: rr -> DomainName
 133     rrType    :: rr -> RecordType
 134     rrClass   :: rr -> RecordClass
 135     rrTTL     :: rr -> TTL
 136     rrPutData :: rr -> Put
 137     rrGetData :: DecompTable -> DomainName -> RecordClass -> TTL -> Get (rr, DecompTable)
 138     toRR      :: rr -> SomeRR
 139     fromRR    :: SomeRR -> Maybe rr
 140
 141     toRR   rr           = SomeRR rr
 142     fromRR (SomeRR rr') = cast rr'
 143
 144 putRR :: ResourceRecord rr => rr -> Put
 145 putRR rr = do putDomainName $ rrName rr
 146               put $ rrType  rr
 147               put $ rrClass rr
 148               putWord32be $ rrTTL rr
 149
 150               let dat = runPut $ rrPutData rr
 151               putWord16be $ fromIntegral $ LBS.length dat
 152               putLazyByteString dat
 153
 154 getRR :: DecompTable -> Get (SomeRR, DecompTable)
 155 getRR dt
 156     = do (nm, dt') <- getDomainName dt
 157          ty        <- get
 158          cl        <- get
 159          ttl       <- G.getWord32be
 160          case ty of
 161            CNAME   -> do (rr, dt'') <- rrGetData dt' nm cl ttl
 162                          return (toRR (rr :: CNAME), dt'')
 163            HINFO   -> do (rr, dt'') <- rrGetData dt' nm cl ttl
 164                          return (toRR (rr :: HINFO), dt'')
 165            AXFR    -> onlyForQuestions "AXFR"
 166            MAILB   -> onlyForQuestions "MAILB"
 167            MAILA   -> onlyForQuestions "MAILA"
 168            AnyType -> onlyForQuestions "ANY"
 169     where
 170       onlyForQuestions name
 171           = fail (name ++ " is only for questions, not an actual resource record.")
 172
 173 data SomeRR = forall rr. ResourceRecord rr => SomeRR rr
 174               deriving Typeable
 175 instance ResourceRecord SomeRR where
 176     rrName    (SomeRR rr) = rrName  rr
 177     rrType    (SomeRR rr) = rrType  rr
 178     rrClass   (SomeRR rr) = rrClass rr
 179     rrTTL     (SomeRR rr) = rrTTL   rr
 180     rrPutData (SomeRR rr) = rrPutData rr
 181     rrGetData _ _ _ _     = fail "SomeRR can't directly be constructed."
 182     toRR   = id
 183     fromRR = Just
 184 instance Eq SomeRR where
 185     (SomeRR a) == (SomeRR b) = Just a == cast b
 186 instance Show SomeRR where
 187     show (SomeRR rr) = show rr
 188
 189 type DecompTable = IntMap BS.ByteString
 190 type TTL = Word32
 191
 192 data CNAME = CNAME' !DomainName !RecordClass !TTL !DomainName
 193              deriving (Eq, Show, Typeable)
 194 instance ResourceRecord CNAME where
 195     rrName    (CNAME' n _ _ _) = n
 196     rrType    _                = CNAME
 197     rrClass   (CNAME' _ c _ _) = c
 198     rrTTL     (CNAME' _ _ t _) = t
 199     rrGetData dt n c t         = do (d, dt') <- getDomainName dt
 200                                     return (CNAME' n c t d, dt')
 201     rrPutData (CNAME' _ _ _ d) = putDomainName d
 202
 203 data HINFO = HINFO' !DomainName !RecordClass !TTL !BS.ByteString !BS.ByteString
 204              deriving (Eq, Show, Typeable)
 205 instance ResourceRecord HINFO where
 206     rrName    (HINFO' n _ _ _ _) = n
 207     rrType    _                  = HINFO
 208     rrClass   (HINFO' _ c _ _ _) = c
 209     rrTTL     (HINFO' _ _ t _ _) = t
 210     rrGetData dt n c t           = do cpu <- getCharString
 211                                       os  <- getCharString
 212                                       return (HINFO' n c t cpu os, dt)
 213     rrPutData (HINFO' _ _ _ c o) = do putCharString c
 214                                       putCharString o
 215
 216 getDomainName :: DecompTable -> Get (DomainName, DecompTable)
 217 getDomainName = flip worker []
 218     where
 219       worker :: DecompTable -> [DomainLabel] -> Get ([DomainLabel], DecompTable)
 220       worker dt soFar
 221           = do (l, dt') <- getDomainLabel dt
 222                case BS.null l of
 223                  True  -> return (reverse (l : soFar), dt')
 224                  False -> worker dt' (l : soFar)
 225
 226 getDomainLabel :: DecompTable -> Get (DomainLabel, DecompTable)
 227 getDomainLabel dt
 228     = do header <- getByteString 1
 229          let Right h
 230                  = runBitGet header $
 231                    do a <- getBit
 232                       b <- getBit
 233                       n <- liftM fromIntegral (getAsWord8 6)
 234                       case (a, b) of
 235                         ( True,  True) -> return $ Offset n
 236                         (False, False) -> return $ Length n
 237                         _              -> fail "Illegal label header"
 238          case h of
 239            Offset n
 240                -> do let Just l = IM.lookup n dt
 241                      return (l, dt)
 242            Length n
 243                -> do offset <- liftM fromIntegral bytesRead
 244                      label  <- getByteString n
 245                      let dt' = IM.insert offset label dt
 246                      return (label, dt')
 247
 248 getCharString :: Get BS.ByteString
 249 getCharString = do len <- G.getWord8
 250                    getByteString (fromIntegral len)
 251
 252 putCharString :: BS.ByteString -> Put
 253 putCharString = putDomainLabel
 254
 255 data LabelHeader
 256     = Offset !Int
 257     | Length !Int
 258
 259 putDomainName :: DomainName -> Put
 260 putDomainName = mapM_ putDomainLabel
 261
 262 putDomainLabel :: DomainLabel -> Put
 263 putDomainLabel l
 264     = do putWord8 $ fromIntegral $ BS.length l
 265          P.putByteString l
 266
 267 data RecordType
 268     = A
 269     | NS
 270     | MD
 271     | MF
 272     | CNAME
 273     | SOA
 274     | MB
 275     | MG
 276     | MR
 277     | NULL
 278     | WKS
 279     | PTR
 280     | HINFO
 281     | MINFO
 282     | MX
 283     | TXT
 284
 285     -- Only for queries:
 286     | AXFR
 287     | MAILB -- Obsolete
 288     | MAILA -- Obsolete
 289     | AnyType
 290     deriving (Show, Eq)
 291
 292 instance Binary Message where
 293     put m = do put $ msgHeader m
 294                putWord16be $ fromIntegral $ length $ msgQuestions m
 295                putWord16be $ fromIntegral $ length $ msgAnswers m
 296                putWord16be $ fromIntegral $ length $ msgAuthorities m
 297                putWord16be $ fromIntegral $ length $ msgAdditionals m
 298                mapM_ putQ  $ msgQuestions m
 299                mapM_ putRR $ msgAnswers m
 300                mapM_ putRR $ msgAuthorities m
 301                mapM_ putRR $ msgAdditionals m
 302
 303     get = do hdr  <- get
 304              nQ   <- liftM fromIntegral G.getWord16be
 305              nAns <- liftM fromIntegral G.getWord16be
 306              nAth <- liftM fromIntegral G.getWord16be
 307              nAdd <- liftM fromIntegral G.getWord16be
 308              (qs  , dt1) <- replicateM' nQ   getQ  IM.empty
 309              (anss, dt2) <- replicateM' nAns getRR dt1
 310              (aths, dt3) <- replicateM' nAth getRR dt2
 311              (adds, _  ) <- replicateM' nAdd getRR dt3
 312              return Message {
 313                           msgHeader      = hdr
 314                         , msgQuestions   = qs
 315                         , msgAnswers     = anss
 316                         , msgAuthorities = aths
 317                         , msgAdditionals = adds
 318                         }
 319
 320 instance Binary Header where
 321     put h = do putWord16be $ hdMessageID h
 322                putLazyByteString flags
 323         where
 324           flags = runBitPut $
 325                   do putNBits 1 $ fromEnum $ hdMessageType h
 326                      putNBits 4 $ fromEnum $ hdOpcode h
 327                      putBit $ hdIsAuthoritativeAnswer h
 328                      putBit $ hdIsTruncated h
 329                      putBit $ hdIsRecursionDesired h
 330                      putBit $ hdIsRecursionAvailable h
 331                      putNBits 3 (0 :: Int)
 332                      putNBits 4 $ fromEnum $ hdResponseCode h
 333
 334     get = do mID   <- G.getWord16be
 335              flags <- getByteString 2
 336              let Right hd
 337                      = runBitGet flags $
 338                        do qr <- liftM (toEnum . fromIntegral) $ getAsWord8 1
 339                           op <- liftM (toEnum . fromIntegral) $ getAsWord8 4
 340                           aa <- getBit
 341                           tc <- getBit
 342                           rd <- getBit
 343                           ra <- getBit
 344                           BG.skip 3
 345                           rc <- liftM (toEnum . fromIntegral) $ getAsWord8 4
 346                           return Header {
 347                                        hdMessageID             = mID
 348                                      , hdMessageType           = qr
 349                                      , hdOpcode                = op
 350                                      , hdIsAuthoritativeAnswer = aa
 351                                      , hdIsTruncated           = tc
 352                                      , hdIsRecursionDesired    = rd
 353                                      , hdIsRecursionAvailable  = ra
 354                                      , hdResponseCode          = rc
 355                                      }
 356              return hd
 357
 358 instance Enum MessageType where
 359     fromEnum Query    = 0
 360     fromEnum Response = 1
 361
 362     toEnum 0 = Query
 363     toEnum 1 = Response
 364     toEnum _ = undefined
 365
 366 instance Enum Opcode where
 367     fromEnum StandardQuery       = 0
 368     fromEnum InverseQuery        = 1
 369     fromEnum ServerStatusRequest = 2
 370
 371     toEnum 0 = StandardQuery
 372     toEnum 1 = InverseQuery
 373     toEnum 2 = ServerStatusRequest
 374     toEnum _ = undefined
 375
 376 instance Enum ResponseCode where
 377     fromEnum NoError        = 0
 378     fromEnum FormatError    = 1
 379     fromEnum ServerFailure  = 2
 380     fromEnum NameError      = 3
 381     fromEnum NotImplemented = 4
 382     fromEnum Refused        = 5
 383
 384     toEnum 0 = NoError
 385     toEnum 1 = FormatError
 386     toEnum 2 = ServerFailure
 387     toEnum 3 = NameError
 388     toEnum 4 = NotImplemented
 389     toEnum 5 = Refused
 390     toEnum _ = undefined
 391
 392 instance Enum RecordType where
 393     fromEnum A       = 1
 394     fromEnum NS      = 2
 395     fromEnum MD      = 3
 396     fromEnum MF      = 4
 397     fromEnum CNAME   = 5
 398     fromEnum SOA     = 6
 399     fromEnum MB      = 7
 400     fromEnum MG      = 8
 401     fromEnum MR      = 9
 402     fromEnum NULL    = 10
 403     fromEnum WKS     = 11
 404     fromEnum PTR     = 12
 405     fromEnum HINFO   = 13
 406     fromEnum MINFO   = 14
 407     fromEnum MX      = 15
 408     fromEnum TXT     = 16
 409     fromEnum AXFR    = 252
 410     fromEnum MAILB   = 253
 411     fromEnum MAILA   = 254
 412     fromEnum AnyType = 255
 413
 414     toEnum 1  = A
 415     toEnum 2  = NS
 416     toEnum 3  = MD
 417     toEnum 4  = MF
 418     toEnum 5  = CNAME
 419     toEnum 6  = SOA
 420     toEnum 7  = MB
 421     toEnum 8  = MG
 422     toEnum 9  = MR
 423     toEnum 10 = NULL
 424     toEnum 11 = WKS
 425     toEnum 12 = PTR
 426     toEnum 13 = HINFO
 427     toEnum 14 = MINFO
 428     toEnum 15 = MX
 429     toEnum 16 = TXT
 430     toEnum 252 = AXFR
 431     toEnum 253 = MAILB
 432     toEnum 254 = MAILA
 433     toEnum 255 = AnyType
 434     toEnum _  = undefined
 435
 436 instance Enum RecordClass where
 437     fromEnum IN       = 1
 438     fromEnum CS       = 2
 439     fromEnum CH       = 3
 440     fromEnum HS       = 4
 441     fromEnum AnyClass = 255
 442
 443     toEnum 1   = IN
 444     toEnum 2   = CS
 445     toEnum 3   = CH
 446     toEnum 4   = HS
 447     toEnum 255 = AnyClass
 448     toEnum _   = undefined
 449
 450 instance Binary RecordType where
 451     get = liftM (toEnum . fromIntegral) G.getWord16be
 452     put = putWord16be . fromIntegral . fromEnum
 453
 454 instance Binary RecordClass where
 455     get = liftM (toEnum . fromIntegral) G.getWord16be
 456     put = putWord16be . fromIntegral . fromEnum