2 files changed, 70 insertions, 51 deletions

diff --git a/Presence/XMLToByteStrings.hs b/Presence/XMLToByteStrings.hs
new file mode 100644
index 00000000..d8d2e965
--- /dev/null
+++ b/Presence/XMLToByteStrings.hs
@@ -0,0 +1,64 @@
+{-# LANGUAGE CPP #-}
+module XMLToByteStrings (xmlToByteStrings) where
+
+import Control.Monad.IO.Class
+import Control.Monad.Trans.Class
+import Control.Monad.Trans.Maybe
+import Control.Monad.Fix
+import Control.Monad (when)
+import           Data.Conduit
+import qualified Data.Conduit.List as CL
+import Data.XML.Types as XML (Event)
+import qualified Data.ByteString as S (ByteString)
+#ifdef RENDERFLUSH
+import Data.Conduit.Blaze (builderToByteStringFlush)
+import Text.XML.Stream.Render (def,renderBuilderFlush)
+#else
+import Text.XML.Stream.Render (def,renderBytes)
+#endif
+
+fixMaybeT f = (>> return ()) . runMaybeT . fix $ f
+mawait :: Monad m => MaybeT (ConduitM i o m) i
+mawait = MaybeT await
+
+
+xmlToByteStrings ::
+  ( MonadUnsafeIO m
+  , MonadIO m) =>
+                 Source m [XML.Event] -> Sink S.ByteString m b -> m b
+xmlToByteStrings xmlSource snk
+#ifdef RENDERFLUSH
+    = xmlSource 
+        $$  flushList
+        =$= renderBuilderFlush def 
+        =$= builderToByteStringFlush 
+        =$= discardFlush 
+        =$ snk
+#else
+    = xmlSource $$ renderChunks =$ snk
+#endif
+
+#ifdef RENDERFLUSH
+flushList :: Monad m => ConduitM [a] (Flush a) m ()
+flushList = fixMaybeT $ \loop -> do
+    xs <- mawait
+    lift ( CL.sourceList xs $$ CL.mapM_ (yield . Chunk) )
+    lift ( yield Flush )
+    loop
+
+discardFlush :: Monad m => ConduitM (Flush a) a m ()
+discardFlush = fixMaybeT $ \loop -> do
+    x <- mawait
+    let unchunk (Chunk a) = a
+        ischunk (Chunk _) = True
+        ischunk _         = False
+    lift . when (ischunk x) $ yield (unchunk x)
+    loop
+#else
+renderChunks :: (MonadUnsafeIO m, MonadIO m) => ConduitM [Event] S.ByteString m ()
+renderChunks = fixMaybeT $  \loop -> do
+    xs <- mawait
+    lift . when (not . null $ xs) $ ( CL.sourceList xs $= renderBytes def $$ CL.mapM_ yield )
+    loop
+#endif
+
diff --git a/Presence/XMPP.hs b/Presence/XMPP.hs
index e51607b2..ec2feb3a 100644
--- a/Presence/XMPP.hs
+++ b/Presence/XMPP.hs
@@ -1,6 +1,5 @@
 {-# LANGUAGE OverloadedStrings   #-}
 {-# LANGUAGE FlexibleContexts    #-}
-{-# LANGUAGE CPP                 #-}
 {-# LANGUAGE ViewPatterns        #-}
 module XMPP 
     ( module XMPPTypes
@@ -19,6 +18,8 @@ import XMPPTypes
 import SocketLike
 import ByteStringOperators
 import ControlMaybe
+import XMLToByteStrings
+
 
 import Data.Maybe (catMaybes)
 import Data.HList
@@ -70,9 +71,6 @@ import Data.Text.Lazy.Encoding as L (decodeUtf8)
 import Data.Text.Lazy (toStrict)
 import qualified Data.Sequence as Seq
 import Data.Foldable (toList)
-#ifdef RENDERFLUSH
-import Data.Conduit.Blaze
-#endif
 import Data.List (find)
 import qualified Text.Show.ByteString as L
 import NestingXML
@@ -631,16 +629,7 @@ handleClient st src snk = do
     rchan <- subscribeToRoster session
     cmdChan <- atomically newTChan
 
-#ifdef RENDERFLUSH
-    writer <- async (   toClient session pchan cmdChan rchan
-                    $$  flushList
-                    =$= renderBuilderFlush def 
-                    =$= builderToByteStringFlush 
-                    =$= discardFlush 
-                     =$ snk )
-#else
-    writer <- async ( toClient session pchan cmdChan rchan $$ renderChunks =$ snk )
-#endif
+    writer <- async ( toClient session pchan cmdChan rchan `xmlToByteStrings` snk ) 
     finally         ( src $= parseBytes def $$ fromClient session cmdChan )
         $ do
             atomically $ writeTChan cmdChan QuitThread
@@ -654,30 +643,6 @@ listenForXmppClients ::
 listenForXmppClients addr_family session_factory port st = do
     doServer (addr_family .*. port .*. session_factory .*. st) handleClient
 
-#ifdef RENDERFLUSH
-flushList :: Monad m => ConduitM [a] (Flush a) m ()
-flushList = fixMaybeT $ \loop -> do
-    xs <- mawait
-    lift ( CL.sourceList xs $$ CL.mapM_ (yield . Chunk) )
-    lift ( yield Flush )
-    loop
-
-discardFlush :: Monad m => ConduitM (Flush a) a m ()
-discardFlush = fixMaybeT $ \loop -> do
-    x <- mawait
-    let unchunk (Chunk a) = a
-        ischunk (Chunk _) = True
-        ischunk _         = False
-    lift . when (ischunk x) $ yield (unchunk x)
-    loop
-#else
-renderChunks :: (MonadUnsafeIO m, MonadIO m) => ConduitM [Event] ByteString m ()
-renderChunks = fixMaybeT $  \loop -> do
-    xs <- mawait
-    lift . when (not . null $ xs) $ ( CL.sourceList xs $= renderBytes def $$ CL.mapM_ yield )
-    loop
-#endif
-
 
 listenForRemotePeers
   :: (HList l, HHead l (XMPPPeerClass session),
@@ -1149,19 +1114,9 @@ handleOutgoingToPeer sock cache chan snk = do
     let failure cmd = do
             writeIORef failed cmd
             putStrLn $ "Failed: " ++ show cmd
-    finally (
-#ifdef RENDERFLUSH
-        handle (\(IOError _ _ _ _ _ _) -> return ()) $ 
-            toPeer sock cache chan failure
-                $$  flushList
-                =$= renderBuilderFlush def
-                =$= builderToByteStringFlush
-                =$= discardFlush
-                 =$ snk
-#else
-        handle (\(IOError _ _ _ _ _ _) -> return ()) $ toPeer sock cache chan failure $$ renderChunks =$ snk
-#endif
-        ) $ L.putStrLn $ "(>P) disconnected " <++> showPeer (RemotePeer p)
+    finally ( handleIO_  (return ()) 
+                         (toPeer sock cache chan failure `xmlToByteStrings` snk) ) 
+        $ L.putStrLn $ "(>P) disconnected " <++> showPeer (RemotePeer p)
     readIORef failed
 
 connect' :: SockAddr -> Int -> IO (Maybe Socket)