Response parsing

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

diff --git a/Network/DNS/Message.hs b/Network/DNS/Message.hs
index 6144d13766e037d551838526f18804bf0b896451..7bedacf5a0922b1816280a8ae4162f9aaf3ff698 100644 (file)
--- a/Network/DNS/Message.hs
+++ b/Network/DNS/Message.hs
@@ -13,14 +13,19 @@ module Network.DNS.Message
     , RecordType
     , RecordClass(..)
 
-    , SomeRR(..)
-    , SomeRT(..)
+    , SomeQT
+    , SomeRR
+    , SomeRT
 
+    , A(..)
+    , NS(..)
     , CNAME(..)
     , HINFO(..)
 
-    , mkQueryType
     , mkDomainName
+    , wrapQueryType
+    , wrapRecordType
+    , wrapRecord
     )
     where
 
@@ -38,6 +43,7 @@ import           Data.Typeable
 import qualified Data.IntMap as IM
 import           Data.IntMap (IntMap)
 import           Data.Word
+import           Network.Socket
 
 
 replicateM' :: Monad m => Int -> (a -> m (b, a)) -> a -> m ([b], a)
@@ -110,9 +116,6 @@ data Question
 
 type SomeQT = SomeRT
 
-mkQueryType :: RecordType rt dt => rt -> SomeQT
-mkQueryType = SomeRT
-
 putQ :: Question -> Put
 putQ q
     = do putDomainName $ qName q
@@ -140,6 +143,12 @@ nameToLabels (DN ls) = ls
 labelsToName :: [DomainLabel] -> DomainName
 labelsToName = DN
 
+rootName :: DomainName
+rootName = DN [BS.empty]
+
+consLabel :: DomainLabel -> DomainName -> DomainName
+consLabel x (DN ys) = DN (x:ys)
+
 mkDomainName :: String -> DomainName
 mkDomainName = labelsToName . mkLabels [] . notEmpty
     where
@@ -225,40 +234,57 @@ getSomeRR dt
            SomeRT rt -> getRR dt rt >>= \ (rr, dt') -> return (SomeRR rr, dt')
 
 
-type DecompTable = IntMap BS.ByteString
+type DecompTable = IntMap DomainName
 type TTL = Word32
 
 getDomainName :: DecompTable -> Get (DomainName, DecompTable)
-getDomainName = flip worker []
+getDomainName = worker
     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')
+      worker :: DecompTable -> Get (DomainName, DecompTable)
+      worker dt
+          = do offset <- liftM fromIntegral bytesRead
+               hdr    <- getLabelHeader
+               case hdr of
+                 Offset n
+                     -> case IM.lookup n dt of
+                          Just name
+                              -> return (name, dt)
+                          Nothing
+                              -> fail ("Illegal offset of label pointer: " ++ show (n, dt))
+                 Length 0
+                     -> return (rootName, dt)
+                 Length n
+                     -> do label       <- getByteString n
+                           (rest, dt') <- worker dt
+                           let name = consLabel label rest
+                               dt'' = IM.insert offset name dt'
+                           return (name, dt'')
+
+      getLabelHeader :: Get LabelHeader
+      getLabelHeader
+          = do header <- lookAhead $ 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' <- 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 G.skip 1
+                           return len
+
 
 getCharString :: Get BS.ByteString
 getCharString = do len <- G.getWord8
@@ -306,6 +332,20 @@ getSomeRT = do n <- liftM fromIntegral G.getWord16be
                  Just srt
                      -> return srt
 
+data A = A deriving (Show, Eq, Typeable)
+instance RecordType A HostAddress where
+    rtToInt       _ = 1
+    putRecordData _ = putWord32be
+    getRecordData _ = \ dt ->
+                      do addr <- G.getWord32be
+                         return (addr, dt)
+
+data NS = NS deriving (Show, Eq, Typeable)
+instance RecordType NS DomainName where
+    rtToInt       _ = 2
+    putRecordData _ = putDomainName
+    getRecordData _ = getDomainName
+
 data CNAME = CNAME deriving (Show, Eq, Typeable)
 instance RecordType CNAME DomainName where
     rtToInt       _ = 5
@@ -321,6 +361,7 @@ instance RecordType HINFO (BS.ByteString, BS.ByteString) where
                                    os  <- getCharString
                                    return ((cpu, os), dt)
 
+
 {-
 data RecordType
     = A
@@ -515,8 +556,21 @@ instance Binary RecordClass where
 
 defaultRTTable :: IntMap SomeRT
 defaultRTTable = IM.fromList $ map toPair $
-                 [ SomeRT CNAME
+                 [ wrapRecordType A
+                 , wrapRecordType NS
+                 , wrapRecordType CNAME
+                 , wrapRecordType HINFO
                  ]
     where
       toPair :: SomeRT -> (Int, SomeRT)
       toPair srt@(SomeRT rt) = (rtToInt rt, srt)
+
+
+wrapQueryType :: RecordType rt dt => rt -> SomeQT
+wrapQueryType = SomeRT
+
+wrapRecordType :: RecordType rt dt => rt -> SomeRT
+wrapRecordType = SomeRT
+
+wrapRecord :: RecordType rt dt => ResourceRecord rt dt -> SomeRR
+wrapRecord = SomeRR
\ No newline at end of file