Introduce Packer monad so that we can compress binary packets.

[haskell-dns.git] / Network / DNS / Packer.hs

diff --git a/Network/DNS/Packer.hs b/Network/DNS/Packer.hs
new file mode 100644 (file)
index 0000000..7f8f895
--- /dev/null
+++ b/Network/DNS/Packer.hs
@@ -0,0 +1,154 @@
+module Network.DNS.Packer
+    ( Packer
+
+    , pack
+    , pack'
+
+    , getState
+    , setState
+    , modifyState
+
+    , bytesWrote
+    , withOffset
+
+    , putByteString
+    , putLazyByteString
+
+    , putWord8
+    , putWord16be
+    , putWord32be
+
+    , putBinary
+    , liftToBinary
+    )
+    where
+
+import qualified Data.Binary as Binary
+import qualified Data.Binary.Put as Bin
+import           Data.Bits
+import qualified Data.ByteString as Strict
+import qualified Data.ByteString.Lazy as Lazy
+import           Data.Int
+import           Data.Word
+
+
+data PackingState s
+    = PackingState {
+        stResult     :: !Lazy.ByteString
+      , stBytesWrote :: !Int64
+      , stUserState  :: s
+      }
+
+newtype Packer s a = P { unP :: PackingState s -> (a, PackingState s) }
+
+instance Monad (Packer s) where
+    return a = P (\ s -> (a, s))
+    m >>= k  = P (\ s -> let (a, s') = unP m s
+                         in
+                           unP (k a) s')
+    fail err = do bytes <- get stBytesWrote
+                  P (error (err
+                            ++ ". Failed packing at byte position "
+                            ++ show bytes))
+
+get :: (PackingState s -> a) -> Packer s a
+get f = P (\ s -> (f s, s))
+
+set :: (PackingState s -> PackingState s) -> Packer s ()
+set f = P (\ s -> ((), f s))
+
+mkState :: Lazy.ByteString -> Int64 -> s -> PackingState s
+mkState xs n s
+    = PackingState {
+        stResult     = xs
+      , stBytesWrote = n
+      , stUserState  = s
+      }
+
+pack' :: Packer s a -> s -> (Lazy.ByteString, s, a)
+pack' m s
+    = let (a, s') = unP m (mkState Lazy.empty 0 s)
+      in
+        (stResult s', stUserState s', a)
+
+pack :: Packer s a -> s -> Lazy.ByteString
+pack = (fst' .) . pack'
+    where
+      fst' (xs, _, _) = xs
+
+getState :: Packer s s
+getState = get stUserState
+
+setState :: s -> Packer s ()
+setState = modifyState . const
+
+modifyState :: (s -> s) -> Packer s ()
+modifyState f
+    = set $ \ st -> st { stUserState = f (stUserState st) }
+
+bytesWrote :: Integral i => Packer s i
+bytesWrote = get stBytesWrote >>= return . fromIntegral
+
+withOffset :: Int64 -> Packer s a -> Packer s a
+withOffset n m
+    = P $ \ s -> let (taken, dropped) = Lazy.splitAt n (stResult s)
+                     padded           = Lazy.take n (taken `Lazy.append` Lazy.repeat 0)
+                     tempState        = s {
+                                          stResult     = padded
+                                        , stBytesWrote = stBytesWrote s - Lazy.length dropped
+                                        }
+                     (a, tempState')  = unP m tempState
+                     newState         = tempState {
+                                          stResult     = replaceHead (stResult s) (stResult tempState')
+                                        , stBytesWrote = max (stBytesWrote s) (stBytesWrote tempState')
+                                        }
+                 in
+                   (a, newState)
+      where
+        replaceHead :: Lazy.ByteString -> Lazy.ByteString -> Lazy.ByteString
+        replaceHead world newHead
+            = let rest = Lazy.drop (Lazy.length newHead) world
+              in
+                newHead `Lazy.append` rest
+
+
+putByteString :: Strict.ByteString -> Packer s ()
+putByteString = putLazyByteString . Lazy.fromChunks . (:[])
+
+putLazyByteString :: Lazy.ByteString -> Packer s ()
+putLazyByteString xs
+    = set $ \ st -> st {
+                      stResult     = stResult st `Lazy.append` xs
+                    , stBytesWrote = stBytesWrote st + Lazy.length xs
+                    }
+
+putWord8 :: Word8 -> Packer s ()
+putWord8 w
+    = set $ \ st -> st {
+                      stResult     = stResult st `Lazy.snoc` w
+                    , stBytesWrote = stBytesWrote st + 1
+                    }
+
+putWord16be :: Word16 -> Packer s ()
+putWord16be w
+    = do putWord8 $ fromIntegral $ (w `shiftR`  8) .&. 0xFF
+         putWord8 $ fromIntegral $  w              .&. 0xFF
+
+putWord32be :: Word32 -> Packer s ()
+putWord32be w
+    = do putWord8 $ fromIntegral $ (w `shiftR` 24) .&. 0xFF
+         putWord8 $ fromIntegral $ (w `shiftR` 16) .&. 0xFF
+         putWord8 $ fromIntegral $ (w `shiftR`  8) .&. 0xFF
+         putWord8 $ fromIntegral $  w              .&. 0xFF
+
+
+putBinary :: Binary.Binary a => a -> Packer s ()
+putBinary = putLazyByteString . Binary.encode
+
+
+liftToBinary :: s -> Packer s a -> Bin.PutM a
+liftToBinary s m
+    = do let (a, s') = unP m (mkState Lazy.empty 0 s)
+
+         Bin.putLazyByteString (stResult s')
+         return a