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