import Test.HUnit
-parseMsg :: [Word8] -> Message
-parseMsg = decode . LBS.pack
-
-
-testData :: [Test]
-testData = [ (parseMsg [ 0x22, 0x79, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00
- , 0x00, 0x00, 0x00, 0x00, 0x04, 0x6D, 0x61, 0x69
- , 0x6C, 0x0A, 0x63, 0x69, 0x65, 0x6C, 0x6F, 0x6E
- , 0x65, 0x67, 0x72, 0x6F, 0x03, 0x6F, 0x72, 0x67
- , 0x00, 0x00, 0x05, 0x00, 0x01
- ]
- ~?=
- Message {
- msgHeader = Header {
- hdMessageID = 8825
- , hdMessageType = Query
- , hdOpcode = StandardQuery
- , hdIsAuthoritativeAnswer = False
- , hdIsTruncated = False
- , hdIsRecursionDesired = True
- , hdIsRecursionAvailable = False
- , hdResponseCode = NoError
- }
- , msgQuestions = [ Question {
- qName = mkDomainName "mail.cielonegro.org."
- , qType = wrapQueryType CNAME
- , qClass = IN
- }
- ]
- , msgAnswers = []
- , msgAuthorities = []
- , msgAdditionals = []
- }
+messages :: [([Word8], Message)]
+messages = [ ( [ 0x22, 0x79, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00
+ , 0x00, 0x00, 0x00, 0x00, 0x04, 0x6D, 0x61, 0x69
+ , 0x6C, 0x0A, 0x63, 0x69, 0x65, 0x6C, 0x6F, 0x6E
+ , 0x65, 0x67, 0x72, 0x6F, 0x03, 0x6F, 0x72, 0x67
+ , 0x00, 0x00, 0x05, 0x00, 0x01
+ ]
+ , Message {
+ msgHeader = Header {
+ hdMessageID = 8825
+ , hdMessageType = Query
+ , hdOpcode = StandardQuery
+ , hdIsAuthoritativeAnswer = False
+ , hdIsTruncated = False
+ , hdIsRecursionDesired = True
+ , hdIsRecursionAvailable = False
+ , hdResponseCode = NoError
+ }
+ , msgQuestions = [ Question {
+ qName = mkDomainName "mail.cielonegro.org."
+ , qType = wrapQueryType CNAME
+ , qClass = IN
+ }
+ ]
+ , msgAnswers = []
+ , msgAuthorities = []
+ , msgAdditionals = []
+ }
)
- , (parseMsg [ 0x22, 0x79, 0x85, 0x00, 0x00, 0x01, 0x00, 0x01
- , 0x00, 0x01, 0x00, 0x01, 0x04, 0x6D, 0x61, 0x69
- , 0x6C, 0x0A, 0x63, 0x69, 0x65, 0x6C, 0x6F, 0x6E
- , 0x65, 0x67, 0x72, 0x6F, 0x03, 0x6F, 0x72, 0x67
- , 0x00, 0x00, 0x05, 0x00, 0x01, 0xC0, 0x0C, 0x00
- , 0x05, 0x00, 0x01, 0x00, 0x01, 0x51, 0x80, 0x00
- , 0x06, 0x03, 0x6E, 0x65, 0x6D, 0xC0, 0x11, 0xC0
- , 0x11, 0x00, 0x02, 0x00, 0x01, 0x00, 0x00, 0x0E
- , 0x10, 0x00, 0x02, 0xC0, 0x31, 0xC0, 0x31, 0x00
- , 0x01, 0x00, 0x01, 0x00, 0x00, 0x0E, 0x10, 0x00
- , 0x04, 0xDB, 0x5E, 0x82, 0x8B
- ]
- ~?=
- Message {
- msgHeader = Header {
- hdMessageID = 8825
- , hdMessageType = Response
- , hdOpcode = StandardQuery
- , hdIsAuthoritativeAnswer = True
- , hdIsTruncated = False
- , hdIsRecursionDesired = True
- , hdIsRecursionAvailable = False
- , hdResponseCode = NoError
- }
- , msgQuestions = [ Question {
- qName = mkDomainName "mail.cielonegro.org."
- , qType = wrapQueryType CNAME
- , qClass = IN
- }
- ]
- , msgAnswers = [ wrapRecord $
- ResourceRecord {
- rrName = mkDomainName "mail.cielonegro.org."
- , rrType = CNAME
- , rrClass = IN
- , rrTTL = 86400
- , rrData = mkDomainName "nem.cielonegro.org."
- }
- ]
- , msgAuthorities = [ wrapRecord $
- ResourceRecord {
- rrName = mkDomainName "cielonegro.org."
- , rrType = NS
- , rrClass = IN
- , rrTTL = 3600
- , rrData = mkDomainName "nem.cielonegro.org."
- }
- ]
- , msgAdditionals = [ wrapRecord $
- ResourceRecord {
- rrName = mkDomainName "nem.cielonegro.org."
- , rrType = A
- , rrClass = IN
- , rrTTL = 3600
- , rrData = unsafePerformIO (inet_addr "219.94.130.139")
- }
- ]
- }
+ , ( [ 0x22, 0x79, 0x85, 0x00, 0x00, 0x01, 0x00, 0x01
+ , 0x00, 0x01, 0x00, 0x01, 0x04, 0x6D, 0x61, 0x69
+ , 0x6C, 0x0A, 0x63, 0x69, 0x65, 0x6C, 0x6F, 0x6E
+ , 0x65, 0x67, 0x72, 0x6F, 0x03, 0x6F, 0x72, 0x67
+ , 0x00, 0x00, 0x05, 0x00, 0x01, 0xC0, 0x0C, 0x00
+ , 0x05, 0x00, 0x01, 0x00, 0x01, 0x51, 0x80, 0x00
+ , 0x06, 0x03, 0x6E, 0x65, 0x6D, 0xC0, 0x11, 0xC0
+ , 0x11, 0x00, 0x02, 0x00, 0x01, 0x00, 0x00, 0x0E
+ , 0x10, 0x00, 0x02, 0xC0, 0x31, 0xC0, 0x31, 0x00
+ , 0x01, 0x00, 0x01, 0x00, 0x00, 0x0E, 0x10, 0x00
+ , 0x04, 0xDB, 0x5E, 0x82, 0x8B
+ ]
+ , Message {
+ msgHeader = Header {
+ hdMessageID = 8825
+ , hdMessageType = Response
+ , hdOpcode = StandardQuery
+ , hdIsAuthoritativeAnswer = True
+ , hdIsTruncated = False
+ , hdIsRecursionDesired = True
+ , hdIsRecursionAvailable = False
+ , hdResponseCode = NoError
+ }
+ , msgQuestions = [ Question {
+ qName = mkDomainName "mail.cielonegro.org."
+ , qType = wrapQueryType CNAME
+ , qClass = IN
+ }
+ ]
+ , msgAnswers = [ wrapRecord $
+ ResourceRecord {
+ rrName = mkDomainName "mail.cielonegro.org."
+ , rrType = CNAME
+ , rrClass = IN
+ , rrTTL = 86400
+ , rrData = mkDomainName "nem.cielonegro.org."
+ }
+ ]
+ , msgAuthorities = [ wrapRecord $
+ ResourceRecord {
+ rrName = mkDomainName "cielonegro.org."
+ , rrType = NS
+ , rrClass = IN
+ , rrTTL = 3600
+ , rrData = mkDomainName "nem.cielonegro.org."
+ }
+ ]
+ , msgAdditionals = [ wrapRecord $
+ ResourceRecord {
+ rrName = mkDomainName "nem.cielonegro.org."
+ , rrType = A
+ , rrClass = IN
+ , rrTTL = 3600
+ , rrData = unsafePerformIO (inet_addr "219.94.130.139")
+ }
+ ]
+ }
)
]
+packMsg :: Message -> [Word8]
+packMsg = LBS.unpack . encode
+
+unpackMsg :: [Word8] -> Message
+unpackMsg = decode . LBS.pack
+
+testData :: [Test]
+testData = map mkPackTest messages
+ ++
+ map mkUnpackTest messages
+ where
+ mkPackTest :: ([Word8], Message) -> Test
+ mkPackTest (bin, msg) = packMsg msg ~?= bin
+
+ mkUnpackTest :: ([Word8], Message) -> Test
+ mkUnpackTest (bin, msg) = unpackMsg bin ~?= msg
main :: IO ()
main = runTestTT (test testData) >> return ()
\ No newline at end of file
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
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
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
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
_ -> 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
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
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)
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
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)
{-
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
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
--- /dev/null
+module Network.DNS.Unpacker
+ ( Unpacker
+ , UnpackingState(..)
+
+ , unpack
+ , unpack'
+
+ , getState
+ , setState
+ , modifyState
+
+ , skip
+ , lookAhead
+ , bytesRead
+
+ , getByteString
+ , getLazyByteString
+ , getWord8
+ , getWord16be
+ , getWord32be
+
+ , getBinary
+ , liftToBinary
+ )
+ where
+
+import qualified Data.Binary as Binary
+import qualified Data.Binary.Get as Bin
+import qualified Data.ByteString as Strict
+import qualified Data.ByteString.Lazy as Lazy
+import Data.Bits
+import Data.Int
+import Data.Word
+
+
+data UnpackingState s
+ = UnpackingState {
+ stSource :: !Lazy.ByteString
+ , stBytesRead :: !Int64
+ , stUserState :: s
+ }
+
+newtype Unpacker s a = U { unU :: UnpackingState s -> (a, UnpackingState s) }
+
+instance Monad (Unpacker s) where
+ return a = U (\ s -> (a, s))
+ m >>= k = U (\ s -> let (a, s') = unU m s
+ in
+ unU (k a) s')
+ fail err = do bytes <- get stBytesRead
+ U (error (err
+ ++ ". Failed unpacking at byte position "
+ ++ show bytes))
+
+get :: (UnpackingState s -> a) -> Unpacker s a
+get f = U (\ s -> (f s, s))
+
+set :: (UnpackingState s -> UnpackingState s) -> Unpacker s ()
+set f = U (\ s -> ((), f s))
+
+mkState :: Lazy.ByteString -> Int64 -> s -> UnpackingState s
+mkState xs n s
+ = UnpackingState {
+ stSource = xs
+ , stBytesRead = n
+ , stUserState = s
+ }
+
+unpack' :: Unpacker s a -> s -> Lazy.ByteString -> (a, s)
+unpack' m s xs
+ = let (a, s') = unU m (mkState xs 0 s)
+ in
+ (a, stUserState s')
+
+unpack :: Unpacker s a -> s -> Lazy.ByteString -> a
+unpack = ((fst .) .) . unpack'
+
+getState :: Unpacker s s
+getState = get stUserState
+
+setState :: s -> Unpacker s ()
+setState = modifyState . const
+
+modifyState :: (s -> s) -> Unpacker s ()
+modifyState f
+ = set $ \ st -> st { stUserState = f (stUserState st) }
+
+skip :: Int64 -> Unpacker s ()
+skip n = getLazyByteString n >> return ()
+
+lookAhead :: Unpacker s a -> Unpacker s a
+lookAhead m = U (\ s -> let (a, _) = unU m s
+ in
+ (a, s))
+
+bytesRead :: Integral i => Unpacker s i
+bytesRead = get stBytesRead >>= return . fromIntegral
+
+getByteString :: Int -> Unpacker s Strict.ByteString
+getByteString n = getLazyByteString (fromIntegral n) >>= return . Strict.concat . Lazy.toChunks
+
+getLazyByteString :: Int64 -> Unpacker s Lazy.ByteString
+getLazyByteString n
+ = do src <- get stSource
+ let (xs, ys) = Lazy.splitAt n src
+ if Lazy.length xs /= n then
+ fail "Too few bytes"
+ else
+ do set $ \ st -> st {
+ stSource = ys
+ , stBytesRead = stBytesRead st + n
+ }
+ return xs
+
+getWord8 :: Unpacker s Word8
+getWord8 = getLazyByteString 1 >>= return . (`Lazy.index` 0)
+
+getWord16be :: Unpacker s Word16
+getWord16be = do xs <- getLazyByteString 2
+ return $ (fromIntegral (xs `Lazy.index` 0) `shiftL` 8) .|.
+ (fromIntegral (xs `Lazy.index` 1))
+
+getWord32be :: Unpacker s Word32
+getWord32be = do xs <- getLazyByteString 4
+ return $ (fromIntegral (xs `Lazy.index` 0) `shiftL` 24) .|.
+ (fromIntegral (xs `Lazy.index` 1) `shiftL` 16) .|.
+ (fromIntegral (xs `Lazy.index` 2) `shiftL` 8) .|.
+ (fromIntegral (xs `Lazy.index` 3))
+
+getBinary :: Binary.Binary a => Unpacker s a
+getBinary = do s <- get id
+ let (a, rest, bytes) = Bin.runGetState Binary.get (stSource s) (stBytesRead s)
+ set $ \ st -> st {
+ stSource = rest
+ , stBytesRead = bytes
+ }
+ return a
+
+
+liftToBinary :: s -> Unpacker s a -> Bin.Get a
+liftToBinary s m
+ = do bytes <- Bin.bytesRead
+ src <- Bin.getRemainingLazyByteString
+
+ let (a, s') = unU m (mkState src bytes s)
+
+ -- These bytes was consumed by the unpacker.
+ Bin.skip (fromIntegral (stBytesRead s' - bytes))
+
+ return a
projects / haskell-dns.git / commitdiff