X-Git-Url: http://git.cielonegro.org/gitweb.cgi?p=haskell-dns.git;a=blobdiff_plain;f=Network%2FDNS%2FMessage.hs;fp=Network%2FDNS%2FMessage.hs;h=5c537956bd1657284a8c3b15db3dd9e9941cce62;hp=be0b79a33800a32dcc769c7ff68129f322261bf4;hb=6423ccc375d8b7d61707de4c6e7b2ace5971be0f;hpb=298473c933e7ad1e101f4db7a7ee115745098235 diff --git a/Network/DNS/Message.hs b/Network/DNS/Message.hs index be0b79a..5c53795 100644 --- a/Network/DNS/Message.hs +++ b/Network/DNS/Message.hs @@ -34,15 +34,17 @@ import Control.Monad 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 qualified Data.ByteString.Lazy as LBS import Data.Typeable import qualified Data.IntMap as IM import Data.IntMap (IntMap) +import qualified Data.Map as M +import Data.Map (Map) import Data.Word +import Network.DNS.Packer as P import Network.DNS.Unpacker as U import Network.Socket @@ -108,11 +110,11 @@ data Question type SomeQT = SomeRT -putQ :: Question -> Put +putQ :: Question -> Packer CompTable () putQ q = do putDomainName $ qName q putSomeRT $ qType q - put $ qClass q + putBinary $ qClass q getQ :: Unpacker DecompTable Question getQ = do nm <- getDomainName @@ -125,23 +127,25 @@ getQ = do nm <- getDomainName } -newtype DomainName = DN [DomainLabel] deriving (Eq, Show, Typeable) +newtype DomainName = DN [DomainLabel] deriving (Eq, Show, Ord, Typeable) type DomainLabel = BS.ByteString -nameToLabels :: DomainName -> [DomainLabel] -nameToLabels (DN ls) = ls - -labelsToName :: [DomainLabel] -> DomainName -labelsToName = DN - rootName :: DomainName rootName = DN [BS.empty] +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 = labelsToName . mkLabels [] . notEmpty +mkDomainName = DN . mkLabels [] . notEmpty where notEmpty :: String -> String notEmpty xs = assert (not $ null xs) xs @@ -173,18 +177,6 @@ data RecordType rt dt => ResourceRecord rt dt 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 - - data SomeRR = forall rt dt. RecordType rt dt => SomeRR (ResourceRecord rt dt) instance Show SomeRR where @@ -194,8 +186,8 @@ instance Eq SomeRR where (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 :: Unpacker DecompTable SomeRR getSomeRR = do srt <- U.lookAhead $ @@ -205,8 +197,9 @@ getSomeRR = do srt <- U.lookAhead $ SomeRT rt -> getResourceRecord rt >>= return . SomeRR +type CompTable = Map DomainName Int type DecompTable = IntMap DomainName -type TTL = Word32 +type TTL = Word32 getDomainName :: Unpacker DecompTable DomainName getDomainName = worker @@ -217,7 +210,7 @@ getDomainName = worker hdr <- getLabelHeader case hdr of Offset n - -> do dt <- getState + -> do dt <- U.getState case IM.lookup n dt of Just name -> return name @@ -229,7 +222,7 @@ getDomainName = worker -> do label <- U.getByteString n rest <- worker let name = consLabel label rest - modifyState $ IM.insert offset name + U.modifyState $ IM.insert offset name return name getLabelHeader :: Unpacker s LabelHeader @@ -262,28 +255,64 @@ getCharString :: Unpacker s BS.ByteString getCharString = do len <- U.getWord8 U.getByteString (fromIntegral len) -putCharString :: BS.ByteString -> Put -putCharString = putDomainLabel +putCharString :: BS.ByteString -> Packer s () +putCharString xs = do P.putWord8 $ fromIntegral $ BS.length xs + P.putByteString xs 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 + + if isRootName rest then + P.putWord8 0 + else + putDomainName rest -putDomainLabel :: DomainLabel -> Put -putDomainLabel l - = do putWord8 $ fromIntegral $ BS.length l - P.putByteString l class (Show rt, Show dt, Eq rt, Eq dt, Typeable rt, Typeable dt) => RecordType rt dt | rt -> dt where rtToInt :: rt -> Int - putRecordData :: rt -> dt -> Put + putRecordData :: rt -> dt -> Packer CompTable () getRecordData :: rt -> Unpacker DecompTable dt - putRecordType :: rt -> Put - putRecordType = putWord16be . fromIntegral . rtToInt + putRecordType :: rt -> Packer s () + putRecordType = P.putWord16be . fromIntegral . rtToInt + + putResourceRecord :: ResourceRecord rt dt -> Packer CompTable () + putResourceRecord rr + = do putDomainName $ rrName rr + putRecordType $ rrType rr + putBinary $ rrClass rr + P.putWord32be $ rrTTL rr + + -- First, write a dummy data length. + offset <- bytesWrote + P.putWord16be 0 + + -- Second, write data. + putRecordData (rrType rr) (rrData rr) + + -- Third, rewrite the dummy length to an actual value. + offset' <- bytesWrote + withOffset offset + $ P.putWord16be (fromIntegral (offset' - offset - 2)) getResourceRecord :: rt -> Unpacker DecompTable (ResourceRecord rt dt) getResourceRecord rt @@ -309,7 +338,7 @@ instance Show SomeRT where instance Eq SomeRT where (SomeRT a) == (SomeRT b) = Just a == cast b -putSomeRT :: SomeRT -> Put +putSomeRT :: SomeRT -> Packer s () putSomeRT (SomeRT rt) = putRecordType rt getSomeRT :: Unpacker s SomeRT @@ -323,7 +352,7 @@ getSomeRT = do n <- liftM fromIntegral U.getWord16be data A = A deriving (Show, Eq, Typeable) instance RecordType A HostAddress where rtToInt _ = 1 - putRecordData _ = putWord32be + putRecordData _ = P.putWord32be getRecordData _ = U.getWord32be data NS = NS deriving (Show, Eq, Typeable) @@ -376,17 +405,18 @@ data RecordType -} 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_ putQ $ msgQuestions m mapM_ putSomeRR $ msgAnswers m mapM_ putSomeRR $ msgAuthorities m mapM_ putSomeRR $ msgAdditionals m - get = liftToBinary IM.empty $ + get = U.liftToBinary IM.empty $ do hdr <- getBinary nQ <- liftM fromIntegral U.getWord16be nAns <- liftM fromIntegral U.getWord16be @@ -405,8 +435,8 @@ instance Binary Message 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 @@ -538,7 +568,7 @@ instance Enum RecordClass where instance Binary RecordClass where get = liftM (toEnum . fromIntegral) G.getWord16be - put = putWord16be . fromIntegral . fromEnum + put = P'.putWord16be . fromIntegral . fromEnum defaultRTTable :: IntMap SomeRT