X-Git-Url: https://git.cielonegro.org/gitweb.cgi?a=blobdiff_plain;f=Network%2FDNS%2FMessage.hs;h=be0b79a33800a32dcc769c7ff68129f322261bf4;hb=298473c933e7ad1e101f4db7a7ee115745098235;hp=7bedacf5a0922b1816280a8ae4162f9aaf3ff698;hpb=86893ea772a5628f813bc83ff4f36327a8d13842;p=haskell-dns.git diff --git a/Network/DNS/Message.hs b/Network/DNS/Message.hs index 7bedacf..be0b79a 100644 --- a/Network/DNS/Message.hs +++ b/Network/DNS/Message.hs @@ -43,18 +43,10 @@ import Data.Typeable import qualified Data.IntMap as IM import Data.IntMap (IntMap) import Data.Word +import Network.DNS.Unpacker as U import Network.Socket -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' - - data Message = Message { msgHeader :: !Header @@ -122,17 +114,16 @@ putQ 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') +getQ :: Unpacker DecompTable Question +getQ = do nm <- getDomainName + ty <- getSomeRT + cl <- getBinary + return Question { + qName = nm + , qType = ty + , qClass = cl + } + newtype DomainName = DN [DomainLabel] deriving (Eq, Show, Typeable) type DomainLabel = BS.ByteString @@ -194,25 +185,6 @@ putRR rr = do putDomainName $ rrName rr 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 @@ -225,44 +197,44 @@ instance Eq SomeRR where putSomeRR :: SomeRR -> Put putSomeRR (SomeRR rr) = putRR 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 DecompTable = IntMap DomainName type TTL = Word32 -getDomainName :: DecompTable -> Get (DomainName, DecompTable) +getDomainName :: Unpacker DecompTable DomainName getDomainName = worker where - worker :: DecompTable -> Get (DomainName, DecompTable) - worker dt - = do offset <- liftM fromIntegral bytesRead + worker :: Unpacker DecompTable DomainName + worker + = do offset <- U.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)) + -> do dt <- 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, dt) + -> return rootName Length n - -> do label <- getByteString n - (rest, dt') <- worker dt + -> do label <- U.getByteString n + rest <- worker let name = consLabel label rest - dt'' = IM.insert offset name dt' - return (name, dt'') + modifyState $ IM.insert offset name + return name - getLabelHeader :: Get LabelHeader + getLabelHeader :: Unpacker s LabelHeader getLabelHeader - = do header <- lookAhead $ getByteString 1 + = do header <- U.lookAhead $ U.getByteString 1 let Right h = runBitGet header $ do a <- getBit @@ -274,7 +246,7 @@ getDomainName = worker _ -> fail "Illegal label header" case h of Offset _ - -> do header' <- getByteString 2 -- Pointers have 2 octets. + -> do header' <- U.getByteString 2 -- Pointers have 2 octets. let Right h' = runBitGet header' $ do BG.skip 2 @@ -282,13 +254,13 @@ getDomainName = worker return $ Offset n return h' len@(Length _) - -> do G.skip 1 + -> do U.skip 1 return len -getCharString :: Get BS.ByteString -getCharString = do len <- G.getWord8 - getByteString (fromIntegral len) +getCharString :: Unpacker s BS.ByteString +getCharString = do len <- U.getWord8 + U.getByteString (fromIntegral len) putCharString :: BS.ByteString -> Put putCharString = putDomainLabel @@ -307,12 +279,28 @@ putDomainLabel l 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) + getRecordData :: rt -> Unpacker DecompTable dt + putRecordType :: rt -> Put putRecordType = putWord16be . fromIntegral . rtToInt + getResourceRecord :: rt -> Unpacker DecompTable (ResourceRecord rt dt) + getResourceRecord rt + = do name <- getDomainName + U.skip 2 -- record type + cl <- getBinary + ttl <- U.getWord32be + U.skip 2 -- data length + dat <- getRecordData rt + return $ ResourceRecord { + rrName = name + , rrType = rt + , rrClass = cl + , rrTTL = ttl + , rrData = dat + } + data SomeRT = forall rt dt. RecordType rt dt => SomeRT rt instance Show SomeRT where @@ -324,8 +312,8 @@ instance Eq SomeRT where 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) @@ -336,9 +324,7 @@ 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) + getRecordData _ = U.getWord32be data NS = NS deriving (Show, Eq, Typeable) instance RecordType NS DomainName where @@ -357,9 +343,9 @@ instance RecordType HINFO (BS.ByteString, BS.ByteString) where rtToInt _ = 13 putRecordData _ (cpu, os) = do putCharString cpu putCharString os - getRecordData _ dt = do cpu <- getCharString + getRecordData _ = do cpu <- getCharString os <- getCharString - return ((cpu, os), dt) + return (cpu, os) {- @@ -400,15 +386,16 @@ instance Binary Message where 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 = 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 getQ + anss <- replicateM nAns getSomeRR + aths <- replicateM nAth getSomeRR + adds <- replicateM nAdd getSomeRR return Message { msgHeader = hdr , msgQuestions = qs @@ -432,7 +419,7 @@ instance Binary Header where 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