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
-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
type SomeQT = SomeRT
-putQ :: Question -> Put
+putQ :: Question -> Packer CompTable ()
putQ 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)
-type DomainLabel = BS.ByteString
+ putBinary $ qClass q
+
+getQ :: Unpacker DecompTable Question
+getQ = do nm <- getDomainName
+ ty <- getSomeRT
+ cl <- getBinary
+ return Question {
+ qName = nm
+ , qType = ty
+ , qClass = cl
+ }
-nameToLabels :: DomainName -> [DomainLabel]
-nameToLabels (DN ls) = ls
-labelsToName :: [DomainLabel] -> DomainName
-labelsToName = DN
+newtype DomainName = DN [DomainLabel] deriving (Eq, Show, Ord, Typeable)
+type DomainLabel = BS.ByteString
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
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
(SomeRR a) == (SomeRR b) = Just a == cast b
-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')
+putSomeRR :: SomeRR -> Packer CompTable ()
+putSomeRR (SomeRR rr) = putResourceRecord rr
+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 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 <- 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, 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'')
+ U.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
_ -> 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
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
+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
- 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
+
+ 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
+ = 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 Eq SomeRT where
(SomeRT a) == (SomeRT b) = Just a == cast b
-putSomeRT :: SomeRT -> Put
+putSomeRT :: SomeRT -> Packer s ()
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)
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)
+ putRecordData _ = P.putWord32be
+ getRecordData _ = U.getWord32be
data NS = NS deriving (Show, Eq, Typeable)
instance RecordType NS DomainName 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)
{-
-}
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 = 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 getQ
+ anss <- replicateM nAns getSomeRR
+ aths <- replicateM nAth getSomeRR
+ adds <- replicateM nAdd getSomeRR
return Message {
msgHeader = hdr
, msgQuestions = qs
}
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
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
instance Binary RecordClass where
get = liftM (toEnum . fromIntegral) G.getWord16be
- put = putWord16be . fromIntegral . fromEnum
+ put = P'.putWord16be . fromIntegral . fromEnum
defaultRTTable :: IntMap SomeRT
projects / haskell-dns.git / blobdiff