Many changes...

[haskell-dns.git] / Network / DNS / Message.hs

diff --git a/Network/DNS/Message.hs b/Network/DNS/Message.hs
index 6144d13766e037d551838526f18804bf0b896451..570548ced67f03e55c1cbe42466bfb505b641e0c 100644 (file)
--- a/Network/DNS/Message.hs
+++ b/Network/DNS/Message.hs
@@ -13,14 +13,39 @@ module Network.DNS.Message
     , RecordType
     , RecordClass(..)
 
     , RecordType
     , RecordClass(..)
 

-    , SomeRR(..)
-    , SomeRT(..)
-
+    , SOAFields(..)
+    , WKSFields(..)
+
+    , SomeQ
+    , SomeQT
+    , SomeRR
+    , SomeRT
+
+    , A(..)
+    , NS(..)
+    , MD(..)
+    , MF(..)

     , CNAME(..)
     , CNAME(..)

+    , SOA(..)
+    , MB(..)
+    , MG(..)
+    , MR(..)
+    , NULL(..)
+    , WKS(..)
+    , PTR(..)

     , HINFO(..)
     , HINFO(..)

+    , MINFO(..)
+    , MX(..)
+    , TXT(..)
+
+    , AXFR(..)
+    , MAILB(..)
+    , MAILA(..)
+    , ANY(..)

 
 

-    , mkQueryType

     , mkDomainName
     , mkDomainName

+    , wrapQuestion
+    , wrapRecord

     )
     where
 
     )
     where
 


@@ -29,7 +54,7 @@ import           Control.Monad
 import           Data.Binary
 import           Data.Binary.BitPut as BP
 import           Data.Binary.Get as G
 import           Data.Binary
 import           Data.Binary.BitPut as BP
 import           Data.Binary.Get as G

-import           Data.Binary.Put as P
+import           Data.Binary.Put as P'

 import           Data.Binary.Strict.BitGet as BG
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Char8 as C8 hiding (ByteString)
 import           Data.Binary.Strict.BitGet as BG
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Char8 as C8 hiding (ByteString)


@@ -37,22 +62,20 @@ import qualified Data.ByteString.Lazy as LBS
 import           Data.Typeable
 import qualified Data.IntMap as IM
 import           Data.IntMap (IntMap)
 import           Data.Typeable
 import qualified Data.IntMap as IM
 import           Data.IntMap (IntMap)

+import qualified Data.IntSet as IS
+import           Data.IntSet (IntSet)
+import qualified Data.Map as M
+import           Data.Map (Map)

 import           Data.Word
 import           Data.Word

-
-
-replicateM' :: Monad m => Int -> (a -> m (b, a)) -> a -> m ([b], a)
-replicateM' = worker []
-    where
-      worker :: Monad m => [b] -> Int -> (a -> m (b, a)) -> a -> m ([b], a)
-      worker soFar 0 _ a = return (reverse soFar, a)
-      worker soFar n f a = do (b, a') <- f a
-                              worker (b : soFar) (n - 1) f a'
+import           Network.DNS.Packer as P
+import           Network.DNS.Unpacker as U
+import           Network.Socket

 
 
 data Message
     = Message {
         msgHeader      :: !Header
 
 
 data Message
     = Message {
         msgHeader      :: !Header

-      , msgQuestions   :: ![Question]
+      , msgQuestions   :: ![SomeQ]

       , msgAnswers     :: ![SomeRR]
       , msgAuthorities :: ![SomeRR]
       , msgAdditionals :: ![SomeRR]
       , msgAnswers     :: ![SomeRR]
       , msgAuthorities :: ![SomeRR]
       , msgAdditionals :: ![SomeRR]


@@ -100,48 +123,87 @@ data ResponseCode
     | Refused
     deriving (Show, Eq)
 
     | Refused
     deriving (Show, Eq)
 

-data Question
+data QueryType qt => Question qt

     = Question {
         qName  :: !DomainName
     = Question {
         qName  :: !DomainName

-      , qType  :: !SomeQT
+      , qType  :: !qt

       , qClass :: !RecordClass
       }
       , qClass :: !RecordClass
       }

-    deriving (Show, Eq)
+    deriving (Typeable)
+
+instance QueryType qt => Show (Question qt) where
+    show q = "Question { qName = " ++ show (qName q) ++
+             ", qType = " ++ show (qType q) ++
+             ", qClass = " ++ show (qClass q) ++ " }"
+
+instance QueryType qt => Eq (Question qt) where
+    a == b = qName  a == qName  b &&
+             qType  a == qType  b &&
+             qClass a == qClass b
+
+data SomeQ = forall qt. QueryType qt => SomeQ (Question qt)
+
+instance Show SomeQ where
+    show (SomeQ q) = show q
+
+instance Eq SomeQ where
+    (SomeQ a) == (SomeQ b) = Just a == cast b

 
 

-type SomeQT = SomeRT
+data SomeQT = forall qt. QueryType qt => SomeQT qt

 
 

-mkQueryType :: RecordType rt dt => rt -> SomeQT
-mkQueryType = SomeRT
+instance Show SomeQT where
+    show (SomeQT qt) = show qt

 
 

-putQ :: Question -> Put
-putQ q
+instance Eq SomeQT where
+    (SomeQT a) == (SomeQT b) = Just a == cast b
+
+putSomeQ :: SomeQ -> Packer CompTable ()
+putSomeQ (SomeQ q)

     = do putDomainName $ qName q
     = do putDomainName $ qName q

-         putSomeRT $ qType q
-         put $ qClass q
-
-getQ :: DecompTable -> Get (Question, DecompTable)
-getQ dt
-    = do (nm, dt') <- getDomainName dt
-         ty        <- getSomeRT
-         cl        <- get
-         let q = Question {
-                   qName  = nm
-                 , qType  = ty
-                 , qClass = cl
-                 }
-         return (q, dt')
-
-newtype DomainName  = DN [DomainLabel] deriving (Eq, Show, Typeable)
+         putQueryType $ qType q
+         putBinary $ qClass q
+
+getSomeQ :: Unpacker DecompTable SomeQ
+getSomeQ
+    = do nm <- getDomainName
+         ty <- getSomeQT
+         cl <- getBinary
+         case ty of
+           SomeQT qt -> return $ SomeQ $
+                        Question {
+                          qName  = nm
+                        , qType  = qt
+                        , qClass = cl
+                        }
+
+getSomeQT :: Unpacker s SomeQT
+getSomeQT = do n <- liftM fromIntegral U.getWord16be
+               case IM.lookup n defaultQTTable of
+                 Just sqt
+                     -> return sqt
+                 Nothing
+                     -> fail ("Unknown query type: " ++ show n)
+
+
+newtype DomainName  = DN [DomainLabel] deriving (Eq, Show, Ord, Typeable)

 type DomainLabel    = BS.ByteString
 
 type DomainLabel    = BS.ByteString
 

-nameToLabels :: DomainName -> [DomainLabel]
-nameToLabels (DN ls) = ls
+rootName :: DomainName
+rootName = DN [BS.empty]

 
 

-labelsToName :: [DomainLabel] -> DomainName
-labelsToName = DN
+isRootName :: DomainName -> Bool
+isRootName (DN [_]) = True
+isRootName _        = False
+
+consLabel :: DomainLabel -> DomainName -> DomainName
+consLabel x (DN ys) = DN (x:ys)
+
+unconsLabel :: DomainName -> (DomainLabel, DomainName)
+unconsLabel (DN (x:xs)) = (x, DN xs)
+unconsLabel x           = error ("Illegal use of unconsLabel: " ++ show x)

 
 mkDomainName :: String -> DomainName
 
 mkDomainName :: String -> DomainName

-mkDomainName = labelsToName . mkLabels [] . notEmpty
+mkDomainName = DN . mkLabels [] . notEmpty

     where
       notEmpty :: String -> String
       notEmpty xs = assert (not $ null xs) xs
     where
       notEmpty :: String -> String
       notEmpty xs = assert (not $ null xs) xs


@@ -173,37 +235,6 @@ data RecordType rt dt => ResourceRecord rt dt
     deriving (Show, Eq, Typeable)
 
 
     deriving (Show, Eq, Typeable)
 
 

-putRR :: forall rt dt. RecordType rt dt => ResourceRecord rt dt -> Put
-putRR rr = do putDomainName $ rrName rr
-              putRecordType $ rrType  rr
-              put $ rrClass rr
-              putWord32be $ rrTTL rr
-
-              let dat = runPut $
-                        putRecordData (undefined :: rt) (rrData rr)
-              putWord16be $ fromIntegral $ LBS.length dat
-              putLazyByteString dat
-
-
-getRR :: forall rt dt. RecordType rt dt => DecompTable -> rt -> Get (ResourceRecord rt dt, DecompTable)
-getRR dt rt
-    = do (nm, dt1)  <- getDomainName dt
-         G.skip 2   -- record type
-         cl         <- get
-         ttl        <- G.getWord32be
-         G.skip 2   -- data length
-         (dat, dt2) <- getRecordData (undefined :: rt) dt1
-
-         let rr = ResourceRecord {
-                    rrName  = nm
-                  , rrType  = rt
-                  , rrClass = cl
-                  , rrTTL   = ttl
-                  , rrData  = dat
-                  }
-         return (rr, dt2)
-
-

 data SomeRR = forall rt dt. RecordType rt dt => SomeRR (ResourceRecord rt dt)
 
 instance Show SomeRR where
 data SomeRR = forall rt dt. RecordType rt dt => SomeRR (ResourceRecord rt dt)
 
 instance Show SomeRR where


@@ -213,79 +244,178 @@ instance Eq SomeRR where
     (SomeRR a) == (SomeRR b) = Just a == cast b
 
 
     (SomeRR a) == (SomeRR b) = Just a == cast b
 
 

-putSomeRR :: SomeRR -> Put
-putSomeRR (SomeRR rr) = putRR rr
+putSomeRR :: SomeRR -> Packer CompTable ()
+putSomeRR (SomeRR rr) = putResourceRecord rr

 
 

-getSomeRR :: DecompTable -> Get (SomeRR, DecompTable)
-getSomeRR dt
-    = do srt <- lookAhead $
-                do getDomainName dt -- skip
-                   getSomeRT
-         case srt of
-           SomeRT rt -> getRR dt rt >>= \ (rr, dt') -> return (SomeRR rr, dt')
+getSomeRR :: Unpacker DecompTable SomeRR
+getSomeRR = do srt <- U.lookAhead $
+                      do getDomainName -- skip
+                         getSomeRT
+               case srt of
+                 SomeRT rt
+                     -> getResourceRecord rt >>= return . SomeRR

 
 

+type CompTable   = Map DomainName Int
+type DecompTable = IntMap DomainName
+type TTL         = Word32

 
 

-type DecompTable = IntMap BS.ByteString
-type TTL = Word32
-
-getDomainName :: DecompTable -> Get (DomainName, DecompTable)
-getDomainName = flip worker []
+getDomainName :: Unpacker DecompTable DomainName
+getDomainName = worker

     where
     where

-      worker :: DecompTable -> [DomainLabel] -> Get (DomainName, DecompTable)
-      worker dt soFar
-          = do (l, dt') <- getDomainLabel dt
-               case BS.null l of
-                 True  -> return (labelsToName (reverse (l : soFar)), dt')
-                 False -> worker dt' (l : soFar)
-
-getDomainLabel :: DecompTable -> Get (DomainLabel, DecompTable)
-getDomainLabel dt
-    = do header <- getByteString 1
-         let Right h
-                 = runBitGet header $
-                   do a <- getBit
-                      b <- getBit
-                      n <- liftM fromIntegral (getAsWord8 6)
-                      case (a, b) of
-                        ( True,  True) -> return $ Offset n
-                        (False, False) -> return $ Length n
-                        _              -> fail "Illegal label header"
-         case h of
-           Offset n
-               -> do let Just l = IM.lookup n dt
-                     return (l, dt)
-           Length n
-               -> do offset <- liftM fromIntegral bytesRead
-                     label  <- getByteString n
-                     let dt' = IM.insert offset label dt
-                     return (label, dt')
-
-getCharString :: Get BS.ByteString
-getCharString = do len <- G.getWord8
-                   getByteString (fromIntegral len)
-
-putCharString :: BS.ByteString -> Put
-putCharString = putDomainLabel
+      worker :: Unpacker DecompTable DomainName
+      worker
+          = do offset <- U.bytesRead
+               hdr    <- getLabelHeader
+               case hdr of
+                 Offset n
+                     -> do dt <- U.getState
+                           case IM.lookup n dt of
+                             Just name
+                                 -> return name
+                             Nothing
+                                 -> fail ("Illegal offset of label pointer: " ++ show (n, dt))
+                 Length 0
+                     -> return rootName
+                 Length n
+                     -> do label <- U.getByteString n
+                           rest  <- worker
+                           let name = consLabel label rest
+                           U.modifyState $ IM.insert offset name
+                           return name
+
+      getLabelHeader :: Unpacker s LabelHeader
+      getLabelHeader
+          = do header <- U.lookAhead $ U.getByteString 1
+               let Right h
+                       = runBitGet header $
+                         do a <- getBit
+                            b <- getBit
+                            n <- liftM fromIntegral (getAsWord8 6)
+                            case (a, b) of
+                              ( True,  True) -> return $ Offset n
+                              (False, False) -> return $ Length n
+                              _              -> fail "Illegal label header"
+               case h of
+                 Offset _
+                     -> do header' <- U.getByteString 2 -- Pointers have 2 octets.
+                           let Right h'
+                                   = runBitGet header' $
+                                     do BG.skip 2
+                                        n <- liftM fromIntegral (getAsWord16 14)
+                                        return $ Offset n
+                           return h'
+                 len@(Length _)
+                     -> do U.skip 1
+                           return len
+
+
+getCharString :: Unpacker s BS.ByteString
+getCharString = do len <- U.getWord8
+                   U.getByteString (fromIntegral len)
+
+putCharString :: BS.ByteString -> Packer s ()
+putCharString xs = do P.putWord8 $ fromIntegral $ BS.length xs
+                      P.putByteString xs

 
 data LabelHeader
     = Offset !Int
     | Length !Int
 
 
 data LabelHeader
     = Offset !Int
     | Length !Int
 

-putDomainName :: DomainName -> Put
-putDomainName = mapM_ putDomainLabel . nameToLabels
+putDomainName :: DomainName -> Packer CompTable ()
+putDomainName name
+    = do ct <- P.getState
+         case M.lookup name ct of
+           Just n
+               -> do let ptr = runBitPut $
+                               do putBit True
+                                  putBit True
+                                  putNBits 14 n
+                     P.putLazyByteString ptr
+           Nothing
+               -> do offset <- bytesWrote
+                     P.modifyState $ M.insert name offset
+
+                     let (label, rest) = unconsLabel name
+
+                     putCharString label

 
 

-putDomainLabel :: DomainLabel -> Put
-putDomainLabel l
-    = do putWord8 $ fromIntegral $ BS.length l
-         P.putByteString l
+                     if isRootName rest then
+                         P.putWord8 0
+                       else
+                         putDomainName rest
+
+class (Show qt, Eq qt, Typeable qt) => QueryType qt where
+    qtToInt :: qt -> Int
+
+    putQueryType :: qt -> Packer s ()
+    putQueryType = P.putWord16be . fromIntegral . qtToInt
+
+instance RecordType rt dt => QueryType rt where
+    qtToInt = rtToInt

 
 class (Show rt, Show dt, Eq rt, Eq dt, Typeable rt, Typeable dt) => RecordType rt dt | rt -> dt where
     rtToInt       :: rt -> Int
 
 class (Show rt, Show dt, Eq rt, Eq dt, Typeable rt, Typeable dt) => RecordType rt dt | rt -> dt where
     rtToInt       :: rt -> Int

-    putRecordType :: rt -> Put
-    putRecordData :: rt -> dt -> Put
-    getRecordData :: rt -> DecompTable -> Get (dt, DecompTable)
-
-    putRecordType = putWord16be . fromIntegral . rtToInt
+    putRecordData :: rt -> dt -> Packer CompTable ()
+    getRecordData :: rt -> Unpacker DecompTable dt
+
+    putRecordType :: rt -> Packer s ()
+    putRecordType = P.putWord16be . fromIntegral . rtToInt
+
+    putRecordDataWithLength :: rt -> dt -> Packer CompTable ()
+    putRecordDataWithLength rt dt
+        = do -- First, write a dummy data length.
+             offset <- bytesWrote
+             P.putWord16be 0
+
+             -- Second, write data.
+             putRecordData rt dt
+
+             -- Third, rewrite the dummy length to an actual value.
+             offset' <- bytesWrote
+             let len = offset' - offset - 2
+             if len <= 0xFFFF then
+                 withOffset offset
+                    $ P.putWord16be $ fromIntegral len
+               else
+                 fail ("putRecordData " ++ show rt ++ " wrote " ++ show len
+                       ++ " bytes, which is way too long")
+
+    putResourceRecord :: ResourceRecord rt dt -> Packer CompTable ()
+    putResourceRecord rr
+        = do putDomainName $ rrName  rr
+             putRecordType $ rrType  rr
+             putBinary     $ rrClass rr
+             P.putWord32be $ rrTTL   rr
+             putRecordDataWithLength (rrType rr) (rrData rr)
+
+    getRecordDataWithLength :: rt -> Unpacker DecompTable dt
+    getRecordDataWithLength rt
+        = do len     <- U.getWord16be
+             offset  <- U.bytesRead
+             dat     <- getRecordData rt
+             offset' <- U.bytesRead
+
+             let consumed = offset' - offset
+             when (consumed /= len)
+                      $ fail ("getRecordData " ++ show rt ++ " consumed " ++ show consumed ++
+                              " bytes but it had to consume " ++ show len ++ " bytes")
+
+             return dat
+
+    getResourceRecord :: rt -> Unpacker DecompTable (ResourceRecord rt dt)
+    getResourceRecord rt
+        = do name     <- getDomainName
+             U.skip 2 -- record type
+             cl       <- getBinary
+             ttl      <- U.getWord32be
+             dat      <- getRecordDataWithLength rt
+             return $ ResourceRecord {
+                          rrName  = name
+                        , rrType  = rt
+                        , rrClass = cl
+                        , rrTTL   = ttl
+                        , rrData  = dat
+                        }

 
 data SomeRT = forall rt dt. RecordType rt dt => SomeRT rt
 
 
 data SomeRT = forall rt dt. RecordType rt dt => SomeRT rt
 


@@ -295,79 +425,258 @@ instance Show SomeRT where
 instance Eq SomeRT where
     (SomeRT a) == (SomeRT b) = Just a == cast b
 
 instance Eq SomeRT where
     (SomeRT a) == (SomeRT b) = Just a == cast b
 

-putSomeRT :: SomeRT -> Put
-putSomeRT (SomeRT rt) = putRecordType rt
-
-getSomeRT :: Get SomeRT
-getSomeRT = do n <- liftM fromIntegral G.getWord16be
+getSomeRT :: Unpacker s SomeRT
+getSomeRT = do n <- liftM fromIntegral U.getWord16be

                case IM.lookup n defaultRTTable of
                  Nothing
                      -> fail ("Unknown resource record type: " ++ show n)
                  Just srt
                      -> return srt
 
                case IM.lookup n defaultRTTable of
                  Nothing
                      -> fail ("Unknown resource record type: " ++ show n)
                  Just srt
                      -> return srt
 

+
+data SOAFields
+    = SOAFields {
+        soaMasterNameServer   :: !DomainName
+      , soaResponsibleMailbox :: !DomainName
+      , soaSerialNumber       :: !Word32
+      , soaRefreshInterval    :: !Word32
+      , soaRetryInterval      :: !Word32
+      , soaExpirationLimit    :: !Word32
+      , soaMinimumTTL         :: !Word32
+      }
+    deriving (Show, Eq, Typeable)
+
+data WKSFields
+    = WKSFields {
+        wksAddress  :: !HostAddress
+      , wksProtocol :: !ProtocolNumber
+      , wksServices :: !IntSet
+      }
+    deriving (Show, Eq, Typeable)
+
+
+data A = A deriving (Show, Eq, Typeable)
+instance RecordType A HostAddress where
+    rtToInt       _ = 1
+    putRecordData _ = P.putWord32be
+    getRecordData _ = U.getWord32be
+
+data NS = NS deriving (Show, Eq, Typeable)
+instance RecordType NS DomainName where
+    rtToInt       _ = 2
+    putRecordData _ = putDomainName
+    getRecordData _ = getDomainName
+
+data MD = MD deriving (Show, Eq, Typeable)
+instance RecordType MD DomainName where
+    rtToInt       _ = 3
+    putRecordData _ = putDomainName
+    getRecordData _ = getDomainName
+
+data MF = MF deriving (Show, Eq, Typeable)
+instance RecordType MF DomainName where
+    rtToInt       _ = 4
+    putRecordData _ = putDomainName
+    getRecordData _ = getDomainName
+

 data CNAME = CNAME deriving (Show, Eq, Typeable)
 instance RecordType CNAME DomainName where
     rtToInt       _ = 5
     putRecordData _ = putDomainName
     getRecordData _ = getDomainName
 
 data CNAME = CNAME deriving (Show, Eq, Typeable)
 instance RecordType CNAME DomainName where
     rtToInt       _ = 5
     putRecordData _ = putDomainName
     getRecordData _ = getDomainName
 

+data SOA = SOA deriving (Show, Eq, Typeable)
+instance RecordType SOA SOAFields where
+    rtToInt       _ = 6
+    putRecordData _ = \ soa ->
+                      do putDomainName $ soaMasterNameServer soa
+                         putDomainName $ soaResponsibleMailbox soa
+                         P.putWord32be $ soaSerialNumber soa
+                         P.putWord32be $ soaRefreshInterval soa
+                         P.putWord32be $ soaRetryInterval soa
+                         P.putWord32be $ soaExpirationLimit soa
+                         P.putWord32be $ soaMinimumTTL soa
+    getRecordData _ = do master  <- getDomainName
+                         mail    <- getDomainName
+                         serial  <- U.getWord32be
+                         refresh <- U.getWord32be
+                         retry   <- U.getWord32be
+                         expire  <- U.getWord32be
+                         ttl     <- U.getWord32be
+                         return SOAFields {
+                                      soaMasterNameServer   = master
+                                    , soaResponsibleMailbox = mail
+                                    , soaSerialNumber       = serial
+                                    , soaRefreshInterval    = refresh
+                                    , soaRetryInterval      = retry
+                                    , soaExpirationLimit    = expire
+                                    , soaMinimumTTL         = ttl
+                                    }
+
+data MB = MB deriving (Show, Eq, Typeable)
+instance RecordType MB DomainName where
+    rtToInt       _ = 7
+    putRecordData _ = putDomainName
+    getRecordData _ = getDomainName
+
+data MG = MG deriving (Show, Eq, Typeable)
+instance RecordType MG DomainName where
+    rtToInt       _ = 8
+    putRecordData _ = putDomainName
+    getRecordData _ = getDomainName
+
+data MR = MR deriving (Show, Eq, Typeable)
+instance RecordType MR DomainName where
+    rtToInt       _ = 9
+    putRecordData _ = putDomainName
+    getRecordData _ = getDomainName
+
+data NULL = NULL deriving (Show, Eq, Typeable)
+instance RecordType NULL BS.ByteString where
+    rtToInt                 _ = 10
+    putRecordData         _ _ = fail "putRecordData NULL can't be defined"
+    getRecordData           _ = fail "getRecordData NULL can't be defined"
+    putRecordDataWithLength _ = \ dat ->
+                                do P.putWord16be $ fromIntegral $ BS.length dat
+                                   P.putByteString dat
+    getRecordDataWithLength _ = do len <- U.getWord16be
+                                   U.getByteString $ fromIntegral len
+
+data WKS = WKS deriving (Show, Eq, Typeable)
+instance RecordType WKS WKSFields where
+    rtToInt       _ = 11
+    putRecordData _ = \ wks ->
+                      do P.putWord32be $ wksAddress wks
+                         P.putWord8 $ fromIntegral $ wksProtocol wks
+                         P.putLazyByteString $ toBitmap $ wksServices wks
+        where
+          toBitmap :: IntSet -> LBS.ByteString
+          toBitmap is
+              = let maxPort   = IS.findMax is
+                    range     = [0 .. maxPort]
+                    isAvail p = p `IS.member` is
+                in
+                  runBitPut $ mapM_ putBit $ map isAvail range
+    getRecordData _ = fail "getRecordData WKS can't be defined"
+
+    getRecordDataWithLength _
+        = do len   <- U.getWord16be
+             addr  <- U.getWord32be
+             proto <- liftM fromIntegral U.getWord8
+             bits  <- U.getByteString $ fromIntegral $ len - 4 - 1
+             return WKSFields {
+                          wksAddress  = addr
+                        , wksProtocol = proto
+                        , wksServices = fromBitmap bits
+                        }
+        where
+          fromBitmap :: BS.ByteString -> IntSet
+          fromBitmap bs
+              = let Right is = runBitGet bs $ worker 0 IS.empty
+                in
+                  is
+
+          worker :: Int -> IntSet -> BitGet IntSet
+          worker pos is
+              = do remain <- BG.remaining
+                   if remain == 0 then
+                       return is
+                     else
+                       do bit <- getBit
+                          if bit then
+                              worker (pos + 1) (IS.insert pos is)
+                            else
+                              worker (pos + 1) is
+
+
+data PTR = PTR deriving (Show, Eq, Typeable)
+instance RecordType PTR DomainName where
+    rtToInt       _ = 12
+    putRecordData _ = putDomainName
+    getRecordData _ = getDomainName
+

 data HINFO = HINFO deriving (Show, Eq, Typeable)
 instance RecordType HINFO (BS.ByteString, BS.ByteString) where
 data HINFO = HINFO deriving (Show, Eq, Typeable)
 instance RecordType HINFO (BS.ByteString, BS.ByteString) where

-    rtToInt       _           = 13
-    putRecordData _ (cpu, os) = do putCharString cpu
-                                   putCharString os
-    getRecordData _ dt        = do cpu <- getCharString
-                                   os  <- getCharString
-                                   return ((cpu, os), dt)
-
-{-
-data RecordType
-    = A
-    | NS
-    | MD
-    | MF
-    | CNAME
-    | SOA
-    | MB
-    | MG
-    | MR
-    | NULL
-    | WKS
-    | PTR
-    | HINFO
-    | MINFO
-    | MX
-    | TXT
-
-    -- Only for queries:
-    | AXFR
-    | MAILB -- Obsolete
-    | MAILA -- Obsolete
-    | AnyType
-    deriving (Show, Eq)
--}
+    rtToInt       _ = 13
+    putRecordData _ = \ (cpu, os) ->
+                      do putCharString cpu
+                         putCharString os
+    getRecordData _ = do cpu <- getCharString
+                         os  <- getCharString
+                         return (cpu, os)
+
+data MINFO = MINFO deriving (Show, Eq, Typeable)
+instance RecordType MINFO (DomainName, DomainName) where
+    rtToInt       _ = 14
+    putRecordData _ = \ (r, e) ->
+                      do putDomainName r
+                         putDomainName e
+    getRecordData _ = do r <- getDomainName
+                         e <- getDomainName
+                         return (r, e)
+
+data MX = MX deriving (Show, Eq, Typeable)
+instance RecordType MX (Word16, DomainName) where
+    rtToInt       _ = 15
+    putRecordData _ = \ (pref, exch) ->
+                      do P.putWord16be pref
+                         putDomainName exch
+    getRecordData _ = do pref <- U.getWord16be
+                         exch <- getDomainName
+                         return (pref, exch)
+
+data TXT = TXT deriving (Show, Eq, Typeable)
+instance RecordType TXT [BS.ByteString] where
+    rtToInt       _ = 16
+    putRecordData _ = mapM_ putCharString
+    getRecordData _ = fail "getRecordData TXT can't be defined"
+
+    getRecordDataWithLength _ = U.getWord16be >>= worker [] . fromIntegral
+        where
+          worker :: [BS.ByteString] -> Int -> Unpacker s [BS.ByteString]
+          worker soFar 0 = return (reverse soFar)
+          worker soFar n = do str <- getCharString
+                              worker (str : soFar) (0 `max` n - 1 - BS.length str)
+
+data AXFR = AXFR deriving (Show, Eq, Typeable)
+instance QueryType AXFR where
+    qtToInt _ = 252
+
+data MAILB = MAILB deriving (Show, Eq, Typeable)
+instance QueryType MAILB where
+    qtToInt _ = 253
+
+data MAILA = MAILA deriving (Show, Eq, Typeable)
+instance QueryType MAILA where
+    qtToInt _ = 254
+
+data ANY = ANY deriving (Show, Eq, Typeable)
+instance QueryType ANY where
+    qtToInt _ = 255
+

 
 instance Binary Message where
 
 instance Binary Message where

-    put m = do put $ msgHeader m
-               putWord16be $ fromIntegral $ length $ msgQuestions m
-               putWord16be $ fromIntegral $ length $ msgAnswers m
-               putWord16be $ fromIntegral $ length $ msgAuthorities m
-               putWord16be $ fromIntegral $ length $ msgAdditionals m
-               mapM_ putQ  $ msgQuestions m
+    put m = P.liftToBinary M.empty $
+            do putBinary $ msgHeader m
+               P.putWord16be $ fromIntegral $ length $ msgQuestions m
+               P.putWord16be $ fromIntegral $ length $ msgAnswers m
+               P.putWord16be $ fromIntegral $ length $ msgAuthorities m
+               P.putWord16be $ fromIntegral $ length $ msgAdditionals m
+               mapM_ putSomeQ  $ msgQuestions m

                mapM_ putSomeRR $ msgAnswers m
                mapM_ putSomeRR $ msgAuthorities m
                mapM_ putSomeRR $ msgAdditionals m
 
                mapM_ putSomeRR $ msgAnswers m
                mapM_ putSomeRR $ msgAuthorities m
                mapM_ putSomeRR $ msgAdditionals m
 

-    get = do hdr  <- get
-             nQ   <- liftM fromIntegral G.getWord16be
-             nAns <- liftM fromIntegral G.getWord16be
-             nAth <- liftM fromIntegral G.getWord16be
-             nAdd <- liftM fromIntegral G.getWord16be
-             (qs  , dt1) <- replicateM' nQ   getQ IM.empty
-             (anss, dt2) <- replicateM' nAns getSomeRR dt1
-             (aths, dt3) <- replicateM' nAth getSomeRR dt2
-             (adds, _  ) <- replicateM' nAdd getSomeRR dt3
+    get = U.liftToBinary IM.empty $
+          do hdr  <- getBinary
+             nQ   <- liftM fromIntegral U.getWord16be
+             nAns <- liftM fromIntegral U.getWord16be
+             nAth <- liftM fromIntegral U.getWord16be
+             nAdd <- liftM fromIntegral U.getWord16be
+             qs   <- replicateM nQ   getSomeQ
+             anss <- replicateM nAns getSomeRR
+             aths <- replicateM nAth getSomeRR
+             adds <- replicateM nAdd getSomeRR

              return Message {
                           msgHeader      = hdr
                         , msgQuestions   = qs
              return Message {
                           msgHeader      = hdr
                         , msgQuestions   = qs


@@ -377,8 +686,8 @@ instance Binary Message where
                         }
 
 instance Binary Header where
                         }
 
 instance Binary Header where

-    put h = do putWord16be $ hdMessageID h
-               putLazyByteString flags
+    put h = do P'.putWord16be $ hdMessageID h
+               P'.putLazyByteString flags

         where
           flags = runBitPut $
                   do putNBits 1 $ fromEnum $ hdMessageType h
         where
           flags = runBitPut $
                   do putNBits 1 $ fromEnum $ hdMessageType h


@@ -391,7 +700,7 @@ instance Binary Header where
                      putNBits 4 $ fromEnum $ hdResponseCode h
 
     get = do mID   <- G.getWord16be
                      putNBits 4 $ fromEnum $ hdResponseCode h
 
     get = do mID   <- G.getWord16be

-             flags <- getByteString 2
+             flags <- G.getByteString 2

              let Right hd
                      = runBitGet flags $
                        do qr <- liftM (toEnum . fromIntegral) $ getAsWord8 1
              let Right hd
                      = runBitGet flags $
                        do qr <- liftM (toEnum . fromIntegral) $ getAsWord8 1


@@ -448,52 +757,6 @@ instance Enum ResponseCode where
     toEnum 5 = Refused
     toEnum _ = undefined
 
     toEnum 5 = Refused
     toEnum _ = undefined
 

-{-
-instance Enum RecordType where
-    fromEnum A       = 1
-    fromEnum NS      = 2
-    fromEnum MD      = 3
-    fromEnum MF      = 4
-    fromEnum CNAME   = 5
-    fromEnum SOA     = 6
-    fromEnum MB      = 7
-    fromEnum MG      = 8
-    fromEnum MR      = 9
-    fromEnum NULL    = 10
-    fromEnum WKS     = 11
-    fromEnum PTR     = 12
-    fromEnum HINFO   = 13
-    fromEnum MINFO   = 14
-    fromEnum MX      = 15
-    fromEnum TXT     = 16
-    fromEnum AXFR    = 252
-    fromEnum MAILB   = 253
-    fromEnum MAILA   = 254
-    fromEnum AnyType = 255
-
-    toEnum 1  = A
-    toEnum 2  = NS
-    toEnum 3  = MD
-    toEnum 4  = MF
-    toEnum 5  = CNAME
-    toEnum 6  = SOA
-    toEnum 7  = MB
-    toEnum 8  = MG
-    toEnum 9  = MR
-    toEnum 10 = NULL
-    toEnum 11 = WKS
-    toEnum 12 = PTR
-    toEnum 13 = HINFO
-    toEnum 14 = MINFO
-    toEnum 15 = MX
-    toEnum 16 = TXT
-    toEnum 252 = AXFR
-    toEnum 253 = MAILB
-    toEnum 254 = MAILA
-    toEnum 255 = AnyType
-    toEnum _  = undefined
--}
-

 instance Enum RecordClass where
     fromEnum IN       = 1
     fromEnum CS       = 2
 instance Enum RecordClass where
     fromEnum IN       = 1
     fromEnum CS       = 2


@@ -510,13 +773,63 @@ instance Enum RecordClass where
 
 instance Binary RecordClass where
     get = liftM (toEnum . fromIntegral) G.getWord16be
 
 instance Binary RecordClass where
     get = liftM (toEnum . fromIntegral) G.getWord16be

-    put = putWord16be . fromIntegral . fromEnum
+    put = P'.putWord16be . fromIntegral . fromEnum

 
 
 defaultRTTable :: IntMap SomeRT
 defaultRTTable = IM.fromList $ map toPair $
 
 
 defaultRTTable :: IntMap SomeRT
 defaultRTTable = IM.fromList $ map toPair $

-                 [ SomeRT CNAME
+                 [ wrapRecordType A
+                 , wrapRecordType NS
+                 , wrapRecordType MD
+                 , wrapRecordType MF
+                 , wrapRecordType CNAME
+                 , wrapRecordType SOA
+                 , wrapRecordType MB
+                 , wrapRecordType MG
+                 , wrapRecordType MR
+                 , wrapRecordType NULL
+                 , wrapRecordType WKS
+                 , wrapRecordType PTR
+                 , wrapRecordType HINFO
+                 , wrapRecordType MINFO
+                 , wrapRecordType MX
+                 , wrapRecordType TXT

                  ]
     where
       toPair :: SomeRT -> (Int, SomeRT)
       toPair srt@(SomeRT rt) = (rtToInt rt, srt)
                  ]
     where
       toPair :: SomeRT -> (Int, SomeRT)
       toPair srt@(SomeRT rt) = (rtToInt rt, srt)

+
+
+defaultQTTable :: IntMap SomeQT
+defaultQTTable = mergeWithRTTable defaultRTTable $ IM.fromList $ map toPair $
+                 [ wrapQueryType AXFR
+                 , wrapQueryType MAILB
+                 , wrapQueryType MAILA
+                 , wrapQueryType ANY
+                 ]
+    where
+      toPair :: SomeQT -> (Int, SomeQT)
+      toPair sqt@(SomeQT qt) = (qtToInt qt, sqt)
+
+      mergeWithRTTable :: IntMap SomeRT -> IntMap SomeQT -> IntMap SomeQT
+      mergeWithRTTable rts qts
+          = IM.union (toQTTable rts) qts
+
+      toQTTable :: IntMap SomeRT -> IntMap SomeQT
+      toQTTable = IM.map toSomeQT
+
+      toSomeQT :: SomeRT -> SomeQT
+      toSomeQT (SomeRT rt) = SomeQT rt
+
+
+wrapQueryType :: QueryType qt => qt -> SomeQT
+wrapQueryType = SomeQT
+
+wrapRecordType :: RecordType rt dt => rt -> SomeRT
+wrapRecordType = SomeRT
+
+wrapQuestion :: QueryType qt => Question qt -> SomeQ
+wrapQuestion = SomeQ
+
+wrapRecord :: RecordType rt dt => ResourceRecord rt dt -> SomeRR
+wrapRecord = SomeRR
\ No newline at end of file