]> gitweb @ CieloNegro.org - haskell-dns.git/commitdiff
Many changes...
authorPHO <pho@cielonegro.org>
Wed, 20 May 2009 06:43:55 +0000 (15:43 +0900)
committerPHO <pho@cielonegro.org>
Wed, 20 May 2009 06:43:55 +0000 (15:43 +0900)
Network/DNS/Message.hs
dns.cabal

index 2bb5a803b3868e080f51d8ae6cb89587e90c915d..e6aaaa5e2af26fd3671fc812844148f562adb827 100644 (file)
@@ -1,37 +1,59 @@
 module Network.DNS.Message
     ( Message(..)
+    , MessageID
     , MessageType(..)
     , Header(..)
     , Opcode(..)
     , ResponseCode(..)
     , Question(..)
     , ResourceRecord(..)
+    , DomainName
+    , DomainLabel
+    , TTL
+    , SomeRR(..)
     , RecordType(..)
     , RecordClass(..)
+
+    , CNAME(..)
+    , HINFO(..)
     )
     where
 
 import           Control.Monad
 import           Data.Binary
-import           Data.Binary.BitPut
+import           Data.Binary.BitPut as BP
 import           Data.Binary.Get as G
-import           Data.Binary.Put
+import           Data.Binary.Put as P
 import           Data.Binary.Strict.BitGet as BG
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Lazy as LBS
+import           Data.Typeable
+import qualified Data.IntMap as IM
+import           Data.IntMap (IntMap)
 import           Data.Word
 
 
+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
       , msgQuestions   :: ![Question]
-      , msgAnswers     :: ![ResourceRecord]
-      , msgAuthorities :: ![ResourceRecord]
-      , msgAdditionals :: ![ResourceRecord]
+      , msgAnswers     :: ![SomeRR]
+      , msgAuthorities :: ![SomeRR]
+      , msgAdditionals :: ![SomeRR]
       }
 
 data Header
     = Header {
-        hdMessageID             :: !Word16
+        hdMessageID             :: !MessageID
       , hdMessageType           :: !MessageType
       , hdOpcode                :: !Opcode
       , hdIsAuthoritativeAnswer :: !Bool
@@ -47,6 +69,8 @@ data Header
       -- + ARCOUNT
       }
 
+type MessageID = Word16
+
 data MessageType
     = Query
     | Response
@@ -75,7 +99,26 @@ data Question
       }
     deriving (Show, Eq)
 
-type DomainName = [[Word8]]
+putQ :: Question -> Put
+putQ q
+    = do putDomainName $ qName q
+         put $ qType  q
+         put $ qClass q
+
+getQ :: DecompTable -> Get (Question, DecompTable)
+getQ dt
+    = do (nm, dt') <- getDomainName dt
+         ty        <- get
+         cl        <- get
+         let q = Question {
+                   qName  = nm
+                 , qType  = ty
+                 , qClass = cl
+                 }
+         return (q, dt')
+
+type DomainName  = [DomainLabel]
+type DomainLabel = BS.ByteString
 
 data RecordClass
     = IN
@@ -85,15 +128,141 @@ data RecordClass
     | AnyClass -- Only for queries
     deriving (Show, Eq)
 
-data ResourceRecord
-    = ResourceRecord {
-        rrName  :: !DomainName
-      , rrType  :: !RecordType
-      , rrClass :: !RecordClass
-      , rrTTL   :: !Word32
-      , rrData  :: ![Word8]
-      }
-    deriving (Show, Eq)
+class (Typeable rr, Show rr, Eq rr) => ResourceRecord rr where
+    rrName    :: rr -> DomainName
+    rrType    :: rr -> RecordType
+    rrClass   :: rr -> RecordClass
+    rrTTL     :: rr -> TTL
+    rrPutData :: rr -> Put
+    rrGetData :: DecompTable -> DomainName -> RecordClass -> TTL -> Get (rr, DecompTable)
+    toRR      :: rr -> SomeRR
+    fromRR    :: SomeRR -> Maybe rr
+
+    toRR   rr           = SomeRR rr
+    fromRR (SomeRR rr') = cast rr'
+
+putRR :: ResourceRecord rr => rr -> Put
+putRR rr = do putDomainName $ rrName rr
+              put $ rrType  rr
+              put $ rrClass rr
+              putWord32be $ rrTTL rr
+
+              let dat = runPut $ rrPutData rr
+              putWord16be $ fromIntegral $ LBS.length dat
+              putLazyByteString dat
+
+getRR :: DecompTable -> Get (SomeRR, DecompTable)
+getRR dt
+    = do (nm, dt') <- getDomainName dt
+         ty        <- get
+         cl        <- get
+         ttl       <- G.getWord32be
+         case ty of
+           CNAME   -> do (rr, dt'') <- rrGetData dt' nm cl ttl
+                         return (toRR (rr :: CNAME), dt'')
+           HINFO   -> do (rr, dt'') <- rrGetData dt' nm cl ttl
+                         return (toRR (rr :: HINFO), dt'')
+           AXFR    -> onlyForQuestions "AXFR"
+           MAILB   -> onlyForQuestions "MAILB"
+           MAILA   -> onlyForQuestions "MAILA"
+           AnyType -> onlyForQuestions "ANY"
+    where
+      onlyForQuestions name
+          = fail (name ++ " is only for questions, not an actual resource record.")
+
+data SomeRR = forall rr. ResourceRecord rr => SomeRR rr
+              deriving Typeable
+instance ResourceRecord SomeRR where
+    rrName    (SomeRR rr) = rrName  rr
+    rrType    (SomeRR rr) = rrType  rr
+    rrClass   (SomeRR rr) = rrClass rr
+    rrTTL     (SomeRR rr) = rrTTL   rr
+    rrPutData (SomeRR rr) = rrPutData rr
+    rrGetData _ _ _ _     = fail "SomeRR can't directly be constructed."
+    toRR   = id
+    fromRR = Just
+instance Eq SomeRR where
+    (SomeRR a) == (SomeRR b) = Just a == cast b
+instance Show SomeRR where
+    show (SomeRR rr) = show rr
+
+type DecompTable = IntMap BS.ByteString
+type TTL = Word32
+
+data CNAME = CNAME' !DomainName !RecordClass !TTL !DomainName
+             deriving (Eq, Show, Typeable)
+instance ResourceRecord CNAME where
+    rrName    (CNAME' n _ _ _) = n
+    rrType    _                = CNAME
+    rrClass   (CNAME' _ c _ _) = c
+    rrTTL     (CNAME' _ _ t _) = t
+    rrGetData dt n c t         = do (d, dt') <- getDomainName dt
+                                    return (CNAME' n c t d, dt')
+    rrPutData (CNAME' _ _ _ d) = putDomainName d
+
+data HINFO = HINFO' !DomainName !RecordClass !TTL !BS.ByteString !BS.ByteString
+             deriving (Eq, Show, Typeable)
+instance ResourceRecord HINFO where
+    rrName    (HINFO' n _ _ _ _) = n
+    rrType    _                  = HINFO
+    rrClass   (HINFO' _ c _ _ _) = c
+    rrTTL     (HINFO' _ _ t _ _) = t
+    rrGetData dt n c t           = do cpu <- getCharString
+                                      os  <- getCharString
+                                      return (HINFO' n c t cpu os, dt)
+    rrPutData (HINFO' _ _ _ c o) = do putCharString c
+                                      putCharString o
+
+getDomainName :: DecompTable -> Get (DomainName, DecompTable)
+getDomainName = flip worker []
+    where
+      worker :: DecompTable -> [DomainLabel] -> Get ([DomainLabel], DecompTable)
+      worker dt soFar
+          = do (l, dt') <- getDomainLabel dt
+               case BS.null l of
+                 True  -> return (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')
+
+getCharString :: Get BS.ByteString
+getCharString = do len <- G.getWord8
+                   getByteString (fromIntegral len)
+
+putCharString :: BS.ByteString -> Put
+putCharString = putDomainLabel
+
+data LabelHeader
+    = Offset !Int
+    | Length !Int
+
+putDomainName :: DomainName -> Put
+putDomainName = mapM_ putDomainLabel
+
+putDomainLabel :: DomainLabel -> Put
+putDomainLabel l
+    = do putWord8 $ fromIntegral $ BS.length l
+         P.putByteString l
 
 data RecordType
     = A
@@ -120,6 +289,34 @@ data RecordType
     | AnyType
     deriving (Show, Eq)
 
+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
+               mapM_ putRR $ msgAnswers m
+               mapM_ putRR $ msgAuthorities m
+               mapM_ putRR $ 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 getRR dt1
+             (aths, dt3) <- replicateM' nAth getRR dt2
+             (adds, _  ) <- replicateM' nAdd getRR dt3
+             return Message {
+                          msgHeader      = hdr
+                        , msgQuestions   = qs
+                        , msgAnswers     = anss
+                        , msgAuthorities = aths
+                        , msgAdditionals = adds
+                        }
+
 instance Binary Header where
     put h = do putWord16be $ hdMessageID h
                putLazyByteString flags
@@ -249,3 +446,11 @@ instance Enum RecordClass where
     toEnum 4   = HS
     toEnum 255 = AnyClass
     toEnum _   = undefined
+
+instance Binary RecordType where
+    get = liftM (toEnum . fromIntegral) G.getWord16be
+    put = putWord16be . fromIntegral . fromEnum
+
+instance Binary RecordClass where
+    get = liftM (toEnum . fromIntegral) G.getWord16be
+    put = putWord16be . fromIntegral . fromEnum
\ No newline at end of file
index 92cb60c15547c5bd994837f61d1becdda25d4923..a401b3cacb12768808eead75fd7dab593f329722 100644 (file)
--- a/dns.cabal
+++ b/dns.cabal
@@ -12,10 +12,13 @@ Build-Type:          Simple
 
 Library
     Build-Depends:
-        base, binary, binary-strict
+        base, binary, binary-strict, bytestring, containers
 
     Exposed-Modules:
         Network.DNS.Message
 
+    Extensions:
+        DeriveDataTypeable, ExistentialQuantification
+
     GHC-Options:
         -Wall