diff --git a/benchmarks/Channels.hs b/benchmarks/Channels.hs
--- a/benchmarks/Channels.hs
+++ b/benchmarks/Channels.hs
@@ -1,13 +1,21 @@
+{-# LANGUAGE LambdaCase, OverloadedStrings #-}
+module Channels where
+
 -- | Like Latency, but creating lots of channels
+
 import System.Environment
-import Control.Monad
+import Control.Monad (void, forM_, forever, replicateM_)
+import Control.Concurrent.MVar
+import Control.Concurrent (forkOS, threadDelay)
 import Control.Applicative
 import Control.Distributed.Process
 import Control.Distributed.Process.Node
-import Network.Transport.ZMQ (createTransport, defaultZMQParameters)
+import Criterion.Measurement
 import Data.Binary (encode, decode)
 import Data.ByteString.Char8 (pack)
+import Network.Transport.ZMQ (createTransport, defaultZMQParameters)
 import qualified Data.ByteString.Lazy as BSL
+import Text.Printf
 
 pingServer :: Process ()
 pingServer = forever $ do
@@ -22,21 +30,43 @@
     (sc, rc) <- newChan :: Process (SendPort (), ReceivePort ())
     send them sc
     receiveChan rc
-  liftIO . putStrLn $ "Did " ++ show n ++ " pings"
 
-initialProcess :: String -> Process ()
-initialProcess "SERVER" = do
+initialServer :: Process ()
+initialServer = do
   us <- getSelfPid
   liftIO $ BSL.writeFile "pingServer.pid" (encode us)
   pingServer
-initialProcess "CLIENT" = do
-  n <- liftIO $ getLine
+
+initialClient :: Int -> Process ()
+initialClient n = do
   them <- liftIO $ decode <$> BSL.readFile "pingServer.pid"
-  pingClient (read n) them
+  pingClient n them
 
 main :: IO ()
-main = do
-  [role, host] <- getArgs
-  Right transport <- createTransport defaultZMQParameters (pack host)
-  node <- newLocalNode transport initRemoteTable
-  runProcess node $ initialProcess role
+main = getArgs >>= \case
+    [] -> defaultBench 
+    [role, host] -> do
+       Right transport <- createTransport defaultZMQParameters (pack host)
+       node <- newLocalNode transport initRemoteTable
+       case role of
+         "SERVER" -> runProcess node initialServer
+         "CLIENT" -> fmap read getLine >>= runProcess node .  initialClient
+         _       -> error "Role should be either SERVER or CLIENT"
+    _ -> error "either call benchmark with [SERVER|CLIENT] host or without arguments"
+  where
+    defaultBench = do
+      void . forkOS $ do
+        Right transport <- createTransport defaultZMQParameters "127.0.0.1"
+        node <- newLocalNode transport initRemoteTable
+        runProcess node $ initialServer
+      threadDelay 1000000
+      e <- newEmptyMVar
+      void . forkOS $ do
+        putStrLn "pings        time\n---          ---\n"
+        forM_ [100,200,600,800,1000,2000,5000,8000,10000] $ \i -> do
+            Right transport <- createTransport defaultZMQParameters "127.0.0.1"
+            node <- newLocalNode transport initRemoteTable
+            d <- time_ (runProcess node $ initialClient i)
+            printf "%-8i %10.4f\n" i d
+        putMVar e ()
+      takeMVar e
diff --git a/benchmarks/Latency.hs b/benchmarks/Latency.hs
--- a/benchmarks/Latency.hs
+++ b/benchmarks/Latency.hs
@@ -1,12 +1,19 @@
-import System.Environment
-import Control.Monad
+{-# LANGUAGE OverloadedStrings, LambdaCase #-}
+module Latency where
+
 import Control.Applicative
+import Control.Monad (void, forM_, forever, replicateM_)
+import Control.Concurrent (forkOS, threadDelay)
+import Control.Concurrent.MVar
 import Control.Distributed.Process
 import Control.Distributed.Process.Node
-import Network.Transport.ZMQ (createTransport, defaultZMQParameters)
+import Criterion.Measurement
 import Data.Binary (encode, decode)
 import Data.ByteString.Char8 (pack)
 import qualified Data.ByteString.Lazy as BSL
+import Network.Transport.ZMQ (createTransport, defaultZMQParameters)
+import System.Environment
+import Text.Printf
 
 pingServer :: Process ()
 pingServer = forever $ do
@@ -17,21 +24,46 @@
 pingClient n them = do
   us <- getSelfPid
   replicateM_ n $ send them us >> (expect :: Process ())
-  liftIO . putStrLn $ "Did " ++ show n ++ " pings"
 
-initialProcess :: String -> Process ()
-initialProcess "SERVER" = do
+initialServer :: Process ()
+initialServer = do
   us <- getSelfPid
   liftIO $ BSL.writeFile "pingServer.pid" (encode us)
   pingServer
-initialProcess "CLIENT" = do
-  n <- liftIO $ getLine
+
+initialClient :: Int -> Process ()
+initialClient n = do
   them <- liftIO $ decode <$> BSL.readFile "pingServer.pid"
-  pingClient (read n) them
+  pingClient n them
 
 main :: IO ()
-main = do
-  [role, host] <- getArgs
-  Right transport <- createTransport defaultZMQParameters (pack host)
-  node <- newLocalNode transport initRemoteTable
-  runProcess node $ initialProcess role
+main = getArgs >>= \case
+  [] -> defaultBenchmark
+  [role, host] -> do
+    Right transport <- createTransport defaultZMQParameters (pack host)
+    node <- newLocalNode transport initRemoteTable
+    case role of
+        "SERVER" -> runProcess node initialServer
+        "CLIENT" -> fmap read getLine >>= runProcess node . initialClient
+        _        -> error "wrong role"
+  _ -> error "either call benchmark with [SERVER|CLIENT] host or without arguments"
+
+defaultBenchmark :: IO ()
+defaultBenchmark = do
+  -- server
+  void . forkOS $ do
+    Right transport <- createTransport defaultZMQParameters "127.0.0.1"
+    node <- newLocalNode transport initRemoteTable
+    runProcess node $ initialServer
+  
+  threadDelay 1000000
+  e <- newEmptyMVar
+  void . forkOS $ do
+    putStrLn "pings        time\n---          ---\n"
+    forM_ [100,200,600,800,1000,2000,5000,8000,10000] $ \i -> do
+        Right transport <- createTransport defaultZMQParameters "127.0.0.1"
+        node <- newLocalNode transport initRemoteTable
+        d <- time_ (runProcess node $ initialClient i)
+        printf "%-8i %10.4f\n" i d
+    putMVar e ()
+  takeMVar e
diff --git a/benchmarks/Throughput.hs b/benchmarks/Throughput.hs
--- a/benchmarks/Throughput.hs
+++ b/benchmarks/Throughput.hs
@@ -1,16 +1,31 @@
+{-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE DeriveDataTypeable #-}
-import System.Environment
-import Control.Monad
-import Control.Applicative
-import Control.Distributed.Process
-import Control.Distributed.Process.Node
-import Network.Transport.ZMQ (createTransport, defaultZMQParameters)
-import Data.Binary
+{-# LANGUAGE LambdaCase #-}
+import           Control.Applicative
+import           Control.Concurrent
+      ( forkOS
+      , threadDelay
+      )
+import           Control.Concurrent.MVar
+import           Control.Distributed.Process
+import           Control.Distributed.Process.Node
+import           Control.Monad
+      ( void
+      , forM_
+      , replicateM_
+      )
+import           Criterion.Measurement
+import           Data.Binary
 import           Data.ByteString.Char8 ( pack )
 import qualified Data.ByteString.Lazy as BSL
-import Data.Typeable
+import           Data.Typeable
+import           Text.Printf
 
-data SizedList a = SizedList { size :: Int , elems :: [a] }
+import           System.Environment
+
+import           Network.Transport.ZMQ (createTransport, defaultZMQParameters)
+
+data SizedList a = SizedList { size :: Int , _elems :: [a] }
   deriving (Typeable)
 
 instance Binary a => Binary (SizedList a) where
@@ -50,22 +65,50 @@
   us <- getSelfPid
   replicateM_ packets $ send them (nats sz)
   send them us
-  n' <- expect
-  liftIO $ print (packets * sz, n' == packets * sz)
+  _ <- expect :: Process Int
+  return ()
 
-initialProcess :: String -> Process ()
-initialProcess "SERVER" = do
+initialServer :: Process ()
+initialServer = do
   us <- getSelfPid
   liftIO $ BSL.writeFile "counter.pid" (encode us)
   counter
-initialProcess "CLIENT" = do
-  n <- liftIO getLine
+
+initialClient :: (Int, Int) -> Process ()
+initialClient n = do
   them <- liftIO $ decode <$> BSL.readFile "counter.pid"
-  count (read n) them
+  count n them
 
 main :: IO ()
-main = do
-  [role, host] <- getArgs
-  Right transport <- createTransport defaultZMQParameters (pack host)
-  node <- newLocalNode transport initRemoteTable
-  runProcess node $ initialProcess role
+main = getArgs >>= \case 
+  [] -> defaultBenchmark
+  [role, host] -> do 
+      Right transport <- createTransport defaultZMQParameters (pack host)
+      node <- newLocalNode transport initRemoteTable
+      case role of
+        "SERVER" -> runProcess node initialServer
+        "CLIENT" -> fmap read getLine >>= runProcess node . initialClient
+        _        -> error "wrong role"
+  _ -> error "either call benchmark with [SERVER|CLIENT] host or without arguments"
+
+
+
+defaultBenchmark :: IO ()
+defaultBenchmark = do
+  -- server
+  void . forkOS $ do
+    Right transport <- createTransport defaultZMQParameters "127.0.0.1"
+    node <- newLocalNode transport initRemoteTable
+    runProcess node $ initialServer
+  
+  threadDelay 1000000
+  e <- newEmptyMVar
+  void . forkOS $ do
+    putStrLn "packet size  time\n---          ---\n"
+    forM_ [1,10,100,200,600,800,1000,2000,4000] $ \i -> do
+        Right transport <- createTransport defaultZMQParameters "127.0.0.1"
+        node <- newLocalNode transport initRemoteTable
+        d <- time_ $ runProcess node $ initialClient (1000,i)
+        printf "%-8i %10.4f\n" i d
+    putMVar e ()
+  takeMVar e
diff --git a/network-transport-zeromq.cabal b/network-transport-zeromq.cabal
--- a/network-transport-zeromq.cabal
+++ b/network-transport-zeromq.cabal
@@ -1,11 +1,16 @@
 name:                network-transport-zeromq
-version:             0.1.0.0
+version:             0.2
 synopsis:            ZeroMQ backend for network-transport
+description:
+  Implementation of the
+  <http://hackage.haskell.org/package/network-transport network-transport>
+  API over ZeroMQ. This provides access to the wealth of transports implemented in ZeroMQ, such as in-process, inter-process, TCP, TPIC and multicast. Furthermore, this makes it possible to encrypt and authenticate clients using ZeroMQ's security mechanisms, and increase throughput using ZeroMQ's intelligent message batching.
+  .
 license:             BSD3
 license-file:        LICENSE
-author:              Alexander Vershilov
-maintainer:          Alexander Vershilov <alexander.vershilov@tweag.io>
 copyright:           (c) 2014, EURL Tweag
+author:              Tweag I/O
+maintainer:          Alexander Vershilov <alexander.vershilov@tweag.io>
 category:            Network
 build-type:          Simple
 cabal-version:       >=1.10
@@ -14,14 +19,18 @@
   Type: git
   Location: https://github.com/tweag/network-transport-zeromq
 
-flag benchmarks
-  description: build benchmarks
+flag install-benchmarks
+  description: Install benchmarks executables (default benchmarks can be triggered during build using \-\-enable-benchmarks option).
   default:     False
 
-flag unsafe-constructs
-  description: use faster primitives, but it will introduce undefined behaviour on some usecases
+flag distributed-process-tests
+  description: build test suites that require distributed-process to be installed
   default:     False
 
+flag unsafe
+  description: Use faster but unsafe primitives.
+  default:     False
+
 library
   build-depends:      base >=4.6 && < 4.8,
                       binary >= 0.6,
@@ -38,11 +47,11 @@
                       void >= 0.6,
                       random >= 1.0
   exposed-modules:    Network.Transport.ZMQ
-                      Network.Transport.ZMQ.Types
-  other-modules:      System.ZMQ4.Utils
+                      Network.Transport.ZMQ.Internal
+                      Network.Transport.ZMQ.Internal.Types
   hs-source-dirs:     src
   ghc-options:        -Wall
-  if flag(unsafe-constructs)
+  if flag(unsafe)
     cpp-options:      -DUNSAFE_SEND
   default-extensions: DeriveGeneric
                       DeriveDataTypeable
@@ -50,9 +59,9 @@
                       LambdaCase
                       ScopedTypeVariables
                       StandaloneDeriving
+  other-extensions:   CPP
   default-language:   Haskell2010
 
-
 Test-Suite test-zeromq
   type:               exitcode-stdio-1.0
   main-is:            TestZMQ.hs
@@ -71,18 +80,23 @@
   build-depends:      base >= 4.4 && < 5,
                       network-transport >= 0.2 && < 0.4,
                       network-transport-zeromq,
-                      zeromq4-haskell >= 0.2
+                      zeromq4-haskell >= 0.2,
+                      tasty >= 0.7,
+                      tasty-hunit >= 0.7
   ghc-options:        -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind
   hs-source-dirs:     tests
   default-language:   Haskell2010
 
-
-
 Test-Suite test-ch-core
   Type:              exitcode-stdio-1.0
   Main-Is:           test-ch.hs
   CPP-Options:       -DTEST_SUITE_MODULE=Control.Distributed.Process.Tests.CH
-  Build-Depends:     base >= 4.4 && < 5,
+  default-extensions:        CPP
+  ghc-options:       -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind
+  HS-Source-Dirs:    tests
+  default-language:  Haskell2010
+  if flag(distributed-process-tests)
+    Build-Depends:   base >= 4.4 && < 5,
                      network-transport-zeromq,
                      distributed-process-tests >= 0.4,
                      network >= 2.3 && < 2.5,
@@ -92,17 +106,20 @@
                      stm,
                      stm-chans,
                      bytestring
-  default-extensions:        CPP
-  ghc-options:       -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind
-  HS-Source-Dirs:    tests
-  default-language:  Haskell2010
-
+    Buildable:       True
+  else
+    Buildable:       False
 
 Test-Suite test-ch-closure
   Type:              exitcode-stdio-1.0
   Main-Is:           test-ch.hs
   CPP-Options:       -DTEST_SUITE_MODULE=Control.Distributed.Process.Tests.Closure
-  Build-Depends:     base >= 4.4 && < 5,
+  default-extensions:        CPP
+  ghc-options:       -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind -caf-all -auto-all
+  HS-Source-Dirs:    tests
+  default-language:  Haskell2010
+  if flag(distributed-process-tests)
+    Build-Depends:   base >= 4.4 && < 5,
                      network-transport-zeromq,
                      distributed-process-tests >= 0.4,
                      network >= 2.3 && < 2.5,
@@ -112,16 +129,20 @@
                      stm,
                      stm-chans,
                      bytestring
-  default-extensions:        CPP
-  ghc-options:       -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind -caf-all -auto-all
-  HS-Source-Dirs:    tests
-  default-language:  Haskell2010
+    Buildable:       True
+  else
+    Buildable:       False
 
 Test-Suite test-ch-stat
   Type:              exitcode-stdio-1.0
   Main-Is:           test-ch.hs
   CPP-Options:       -DTEST_SUITE_MODULE=Control.Distributed.Process.Tests.Stats
-  Build-Depends:     base >= 4.4 && < 5,
+  default-extensions:        CPP
+  ghc-options:       -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind
+  HS-Source-Dirs:    tests
+  default-language:  Haskell2010
+  if flag(distributed-process-tests)
+    Build-Depends:   base >= 4.4 && < 5,
                      network-transport-zeromq,
                      distributed-process-tests >= 0.4,
                      network >= 2.3 && < 2.5,
@@ -131,30 +152,71 @@
                      stm,
                      stm-chans,
                      bytestring
-  default-extensions:        CPP
-  ghc-options:       -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind
-  HS-Source-Dirs:    tests
+    Buildable:       True
+  else
+    Buildable:       False
+
+benchmark bench-channels-local
+  type:              exitcode-stdio-1.0
+  main-is:           Channels.hs
+  build-depends:     base >= 4.4 && < 5,
+                     network-transport-zeromq,
+                     bytestring,
+                     binary,
+                     distributed-process,
+                     criterion
+  hs-source-dirs:    benchmarks
+  ghc-options:       -O2 -Wall -threaded
   default-language:  Haskell2010
 
+benchmark bench-latency-local
+  type:              exitcode-stdio-1.0
+  main-is:           Latency.hs
+  build-depends:     base >= 4.4 && < 5,
+                     network-transport-zeromq,
+                     bytestring,
+                     binary,
+                     distributed-process,
+                     criterion
+  hs-source-dirs:    benchmarks
+  ghc-options:       -O2 -Wall -threaded
+  default-language:  Haskell2010
 
+benchmark bench-throughput-local
+  type:              exitcode-stdio-1.0
+  main-is:           Throughput.hs
+  build-depends:     base >= 4.4 && < 5,
+                     network-transport-zeromq,
+                     bytestring,
+                     binary,
+                     distributed-process,
+                     criterion
+  hs-source-dirs:    benchmarks
+  ghc-options:       -O2 -Wall -threaded
+  default-language:  Haskell2010
+  default-extensions: BangPatterns
 
---- benchmark files
+-- Installable benchmark executables, so as to allow passing arguments
+-- to the benchmarks on the command line and hence run more complex
+-- benchmarks over several hosts.
+
 executable bench-dp-latency
+  ghc-options:       -O2 -Wall
   main-is: Latency.hs
-  build-depends:     base >= 4.4 && < 5,
+  hs-source-dirs:    benchmarks
+  if flag(install-benchmarks)
+    build-depends:   base >= 4.4 && < 5,
                      network-transport-zeromq,
                      bytestring,
                      binary,
                      distributed-process
-  hs-source-dirs:    benchmarks
-  if flag(benchmarks)
     buildable:       True
   else
     buildable:       False
   default-language:  Haskell2010
 
-Executable bench-dp-throughput
-  if flag(benchmarks)
+executable bench-dp-throughput
+  if flag(install-benchmarks)
     Build-Depends:   base >= 4.4 && < 5,
                      distributed-process,
                      network-transport-zeromq,
@@ -167,9 +229,10 @@
   hs-source-dirs:    benchmarks
   ghc-options:       -Wall
   default-extensions:        BangPatterns
+  default-language:  Haskell2010
 
-Executable bench-dp-channels
-    if flag(benchmarks)
+executable bench-dp-channels
+    if flag(install-benchmarks)
       Build-Depends:   base >= 4.4 && < 5,
                        distributed-process,
                        network-transport-zeromq,
@@ -180,5 +243,6 @@
       buildable: False
     main-is:           Channels.hs
     hs-source-dirs:    benchmarks
-    ghc-options:       -Wall
-    default-extensions:        BangPatterns
+    ghc-options:       -Wall -O2 -Wall -threaded
+    default-extensions: BangPatterns
+    default-language:   Haskell2010
diff --git a/src/Network/Transport/ZMQ.hs b/src/Network/Transport/ZMQ.hs
--- a/src/Network/Transport/ZMQ.hs
+++ b/src/Network/Transport/ZMQ.hs
@@ -1,27 +1,43 @@
 -- |
--- Module:    Network.Transport.ZMQ
--- Copyright: (C) 2014, EURL Tweag
--- License:   BSD-3
-{-# LANGUAGE DeriveGeneric, StandaloneDeriving, OverloadedStrings, DeriveDataTypeable, CPP #-}
+-- Copyright: (C) 2014 EURL Tweag
+--
+-- 0MQ implemmentation of the network-transport API.
+--
+-- The 0MQ implementation guarantees that only a single connection (socket) will
+-- be used between endpoints, provided that the addresses specified are
+-- canonical. If /A/ connects to /B/ and reports its address as
+-- @192.168.0.1:8080@ and /B/ subsequently connects tries to connect to /A/ as
+-- @client1.local:http-alt@ then the transport layer will not realize that the
+-- TCP connection can be reused.
+--
+-- This module is intended to be imported qualified.
+
+{-# LANGUAGE CPP #-}
+
 module Network.Transport.ZMQ
   ( -- * Main API
-    createTransport       -- :: ZMQParameters -> ByteString -> IO (Either ZMQError Transport)
-  , ZMQParameters(..)     
-  , ZMQAuthType(..)
-  , defaultZMQParameters  -- :: ZMQParameters
+    createTransport
+  , ZMQParameters(..)
+  , SecurityMechanism(..)
+  , defaultZMQParameters
   -- * Internals
   -- $internals
-  , createTransportEx
-  , breakConnectionEndPoint -- :: ZMQTransport -> EndPointAddress -> EndPointAddress -> IO ()
-  , breakConnection         -- :: ZMQTransport -> EndPointAddress -> EndPointAddress -> IO ()
+  , createTransportExposeInternals
+  , breakConnectionEndPoint
+  , breakConnection
   , unsafeConfigurePush
+  -- $cleanup
+  , registerCleanupAction
+  , registerValidCleanupAction
+  , applyCleanupAction
   -- * Design
   -- $design
   -- ** Multicast
   -- $multicast
   ) where
 
-import Network.Transport.ZMQ.Types
+import Network.Transport.ZMQ.Internal.Types
+import qualified Network.Transport.ZMQ.Internal as ZMQ
 
 import           Control.Applicative
 import           Control.Concurrent
@@ -37,12 +53,11 @@
       , forever
       , unless
       , join
-      , forM_
       , foldM
       , when
       , (<=<)
       )
-import           Control.Exception 
+import           Control.Exception
       ( AsyncException )
 import           Control.Monad.Catch
       ( finally
@@ -68,28 +83,26 @@
       , modifyIORef
       , readIORef
       , writeIORef
+      , atomicModifyIORef'
       )
+import qualified Data.IntMap.Strict as IntMap
 import           Data.List.NonEmpty
 import qualified Data.Map.Strict as Map
--- import           Data.Set (Set)
 import qualified Data.Set as Set
 import qualified Data.Traversable as Traversable
 import           Data.Typeable
+import           Data.Unique
 import           Data.Void
 import           GHC.Generics
       ( Generic )
 
 import Network.Transport
-import Network.Transport.ZMQ.Types
 import           System.IO
       ( fixIO )
 import           System.ZMQ4
       ( Context )
 import qualified System.ZMQ4 as ZMQ
-import qualified System.ZMQ4.Utils   as ZMQ
 
-import Text.Printf
-
 --------------------------------------------------------------------------------
 --- Internal datatypes                                                        --
 --------------------------------------------------------------------------------
@@ -103,25 +116,25 @@
 --  |       +---+---------------- can be configured by user, one host,
 --  |                             port pair per distributed process
 --  |                             instance
---  +---------------------------- In feature it will be possible to add another schemas. 
+--  +---------------------------- In future it will be possible to add other
+--                                schemas.
 -- @
 --
 -- Transport specifies host that will be used, and port will be
--- automatically generated by the transport for each endpoint.                               
+-- automatically generated by the transport for each endpoint.
 --
 -- * Connections Reliability
 --
--- Currently only reliable connections are supportred. In case if
--- Unreliable type was passed reliable connection will be created.
--- This may be changed in future versions.
+-- Currently only reliable connections are supported. In case 'Unreliable' was
+-- passed the connection will nonetheless be reliable, since it is not incorrect
+-- to do so. This may be changed in future versions.
 --
 -- Network-transport-zeromq maintains one thread for each endpoint that is
--- used to read incomming requests, all sends and connection requests are
--- handled from the user thread and mostly asynchronous.
+-- used to read incoming requests. All sends and connection requests are
+-- handled from the user thread and are mostly asynchronous.
 --
--- Each endpoint obtains one pull socket and push socket for each remote
--- end point, all lightweight threads abrigged into one heavyweight
--- connection.
+-- Each endpoint obtains one pull socket and one push socket for each remote end
+-- point. All lightweight threads abrigged into one heavyweight connection.
 --
 -- Virually connections looks like the following plot:
 --
@@ -133,34 +146,32 @@
 -- |               | pull |<-----------------------| push |           |
 -- +               +------+                        +------+           |
 -- |                  |     +--------------------+     |              |
--- +               +------+/~~~~ connection 1 ~~~~\\+------+           |
+-- +               +------+/~~~~ connection 1 ~~~~\\+------+          |
 -- |               | push |~~~~~ connection 2 ~~~~~| pull |           |
--- |               +------+\\                      /+------+           |
+-- |               +------+\\                      /+------+          |
 -- |                  |     +--------------------+     |              |
 -- +------------------+                                +--------------+
 -- @
 --
--- Physically 0mq may choose better representations and mapping on
--- a real connections.
---
--- ZeroMQ connection takes care of reliability thus for correct lifeness,
--- it stores messages so in case of connecdtion break in may be restarted
--- with no message loss. So heartbeating procedure should be introduced on 
--- the top on network-transport, and 'breakConnection' should be used to notify
--- connection death. However if High Water Mark will be reached connection
--- to endpoint considered failed and connection break procedure started
--- automatically.
+-- Physically ZeroMQ may choose better representations and mapping on a real
+-- connections.
 --
+-- The ZeroMQ library takes care of reliability. It keeps message in a queue
+-- such that in case of connection break, it may be restarted with no message
+-- loss. So heartbeating procedure should be introduced on the top on
+-- network-transport, and 'breakConnection' should be used to notify connection
+-- death. However if High Water Mark will be reached connection to endpoint
+-- considered failed and connection break procedure started automatically.
 
 -- Naming conventions.
 -- api* -- functions that can be called by user
--- localEndPoint -- internal functions that are used in endpoints
+-- localEndPoint  -- internal functions that are used in endpoints
 -- remoteEndPoint -- internal functions that are used in endpoints
 --
--- Transport 
+-- Transport
 --   |--- Valid                transport itself
 --   |      |--- newEndpoint
---   |      |--- close         
+--   |      |--- close
 --   |--- Closed
 --
 -- LocaEndPoint
@@ -169,17 +180,17 @@
 
 -- Connection closing procedure.
 -- Once endpoint is gracefully closed by user or by transport close
--- RemoteEndPoint goes into 'RemoteEndPointClosing' state and sends
--- a message to remote endpoint if link is known as alive, in 
--- RemoteEndPointClosing all messages are ignored and EndPointClosed
--- is returned. Uppon a EndPointClose delivery endpoint replies with
--- EndPointClose message, and starts cleanup procedure (marking link
--- as not alive). When endpoint receives EndPointCloseOk message it
--- moved RemoteEndPoint to close state and removes all data structures.
+-- 'RemoteEndPoint' goes into 'RemoteEndPointClosing' state and sends
+-- a message to remote endpoint if link is known as alive, in
+-- 'RemoteEndPointClosing' all messages are ignored and 'EndPointClosed'
+-- is returned. Uppon a 'EndPointClose' delivery endpoint replies with
+-- 'EndPointClose' message, and starts cleanup procedure (marking link
+-- as not alive). When endpoint receives 'EndPointCloseOk' message it
+-- moved 'RemoteEndPoint' to close state and removes all data structures.
 -- Together with marking a endpoint as closed asynchronous timeout is
--- set, and if EndPointCloseOk is not delivered withing that timeperiod
--- EndPoint is closed.
--- 
+-- set, and if 'EndPointCloseOk' is not delivered withing that timeperiod
+-- 'EndPoint' is closed.
+--
 -- endpoint-1                                 endpoint-2
 --  1. mark as closing
 --  [RemoteEndPoint:Closing] ----endPointClose->    [remoteEndPoint:Valid]
@@ -188,9 +199,9 @@
 --  2. mark as closed
 --  [RemoteEndPoint:Closed]                   4. cleanup remote end point
 --
--- EndPoint can be closed in a two ways: normally and abnormally (in case
--- of exception or invariant vionation on a remove side). If EndPoint is
--- closed normally all opened connections will recevice ConnectionClosed 
+-- 'EndPoint' can be closed in a two ways: normally and abnormally (in case
+-- of exception or invariant vionation on a remove side). If 'EndPoint' is
+-- closed normally all opened connections will recevice ConnectionClosed
 -- events, otherwise.
 --
 -- Multicast group schema:
@@ -223,76 +234,55 @@
 -- message. This is done in order to keep ability to notify receiver about
 -- multcast group close.
 --
--- Current solution incorrectly keeps track of incomming connections, this
--- means that we can't really guarantee that close message is delivered to
--- every recipient and it disconnected correctly.
---
-
+-- Current solution incorrectly keeps track of incoming connections. This means
+-- that we can't really guarantee that a "close"" message is delivered to every
+-- recipient and it disconnected correctly.
 
--- $internals 
--- Internal function provides additional 0mq specific set of configuration
--- options, this functionality should be used with caution as it may break
--- required socket properties.
+-- $internals Internal functions provide additional ZeroMQ specific set of
+-- configuration options. This functionality should be used with caution as it
+-- may break required socket properties.
 
--- | Messages
+-- | Messages.
 data ZMQMessage
-      = MessageConnect !EndPointAddress -- ^ Connection greeting
-      | MessageInitConnection !EndPointAddress !ConnectionId !Reliability
-      | MessageInitConnectionOk !EndPointAddress !ConnectionId !ConnectionId
-      | MessageCloseConnection !ConnectionId
-      | MessageData !ConnectionId
-      | MessageEndPointClose !EndPointAddress !Bool
-      | MessageEndPointCloseOk !EndPointAddress
-      deriving (Generic)
-
-{- 
-instance Binary ZMQMessage where
-  put (MessageConnect ep) = putWord64le 0 >> put ep
-  put (MessageInitConnection ep cid rel)   = putWord64le 2 >> put ep >> put cid >> put rel
-  put (MessageInitConnectionOk ep cid rid) = putWord64le 3 >> put ep >> put cid >> put rid
-  put (MessageCloseConnection cid)         = putWord64le 4 >> put cid
-  put (MessageEndPointClose ep)            = putWord64le 5 >> put ep
-  put (MessageData cid)                    = putWord64le cid
-  get = do x <- getWord64be
-           case x of
-             0 -> MessageConnect <$> get
-             2 -> MessageInitConnection   <$> get <*> get <*> get
-             3 -> MessageInitConnectionOk <$> get <*> get <*> get
-             4 -> MessageCloseConnection  <$> get
-             5 -> MessageEndPointClose    <$> get
-             6 -> MessageData <$> pure x
-
-reservedConnectionId = 7
--}
+  = MessageConnect !EndPointAddress -- ^ Connection greeting
+  | MessageInitConnection !EndPointAddress !ConnectionId !Reliability
+  | MessageInitConnectionOk !EndPointAddress !ConnectionId !ConnectionId
+  | MessageCloseConnection !ConnectionId
+  | MessageData !ConnectionId
+  | MessageEndPointClose !EndPointAddress !Bool
+  | MessageEndPointCloseOk !EndPointAddress
+  deriving (Generic)
 
 instance Binary ZMQMessage
 
-data ZMQError = InvariantViolation String
-              | IncorrectState String
-              | ConnectionFailed
-              deriving (Typeable, Show)
+data ZMQError
+  = InvariantViolation String
+  | IncorrectState String
+  | ConnectionFailed
+  deriving (Typeable, Show)
 
 instance Exception ZMQError
 
 -- | Create 0MQ based transport.
 createTransport :: ZMQParameters
-                -> ByteString     -- ^ Transport address (IP or hostname)
+                -> ByteString            -- ^ Transport address (IP or hostname)
                 -> IO (Either (TransportError Void) Transport)
-createTransport z b = fmap (fmap snd) (createTransportEx z b)
+createTransport z b = fmap (fmap snd) (createTransportExposeInternals z b)
 
--- | Create 'Transport' and export internal transport state that may be
--- used in API.
-createTransportEx :: ZMQParameters    -- ^ Transport features.
-                  -> ByteString       -- ^ Host name or IP address
-                  -> IO (Either (TransportError Void) (ZMQTransport, Transport))
-createTransportEx params host = do
+-- | You should probably not use this function (used for unit testing only)
+createTransportExposeInternals
+  :: ZMQParameters                       -- ^ Configuration parameters for ZeroMQ
+  -> ByteString                          -- ^ Host name or IP address
+  -> IO (Either (TransportError Void) (TransportInternals, Transport))
+createTransportExposeInternals params host = do
     ctx       <- ZMQ.context
-    mtid <- case authorizationType params of
-              ZMQAuthPlain user pass -> Just <$> ZMQ.authManager ctx user pass
-              _ -> return Nothing
-    transport <- ZMQTransport 
-    	<$> pure addr 
-        <*> newMVar (TransportValid $ ValidTransportState ctx Map.empty mtid)
+    mtid <- Traversable.sequenceA $
+            fmap (\(SecurityPlain user pass) -> ZMQ.authManager ctx user pass) $
+                 zmqSecurityMechanism params
+    mcl  <- newIORef IntMap.empty
+    transport <- TransportInternals
+    	<$> pure addr
+        <*> newMVar (TransportValid $ ValidTransportState ctx Map.empty mtid mcl)
     return $ Right (transport, Transport
       { newEndPoint    = apiNewEndPoint params transport
       , closeTransport = apiTransportClose transport
@@ -301,27 +291,27 @@
     addr = B.concat ["tcp://", host]
 
 -- Synchronous
-apiTransportClose :: ZMQTransport -> IO ()
+apiTransportClose :: TransportInternals -> IO ()
 apiTransportClose transport = mask_ $ do
     old <- swapMVar (_transportState transport) TransportClosed
     case old of
       TransportClosed -> return ()
-      TransportValid (ValidTransportState ctx m mtid) -> do
+      TransportValid (ValidTransportState ctx m mtid mcl) -> do
         case mtid of
           Nothing -> return ()
           Just tid -> Async.cancel tid
-        forM_ (Map.elems m) $ apiCloseEndPoint transport
-	ZMQ.term ctx
+        Foldable.sequence_ $ Map.map (apiCloseEndPoint transport) m
+        Foldable.sequence_ =<< atomicModifyIORef' mcl (\x -> (IntMap.empty, x))
+        ZMQ.term ctx
 
-apiNewEndPoint :: ZMQParameters -> ZMQTransport -> IO (Either (TransportError NewEndPointErrorCode) EndPoint)
+apiNewEndPoint :: ZMQParameters -> TransportInternals -> IO (Either (TransportError NewEndPointErrorCode) EndPoint)
 apiNewEndPoint params transport = do
---    printf "[transport] endpoint create\n"
     elep <- modifyMVar (_transportState transport) $ \case
        TransportClosed -> return (TransportClosed, Left $ TransportError NewEndPointFailed "Transport is closed.")
-       v@(TransportValid i@(ValidTransportState ctx _ _)) -> do
+       v@(TransportValid i@(ValidTransportState ctx _ _ _)) -> do
          eEndPoint <- endPointCreate params ctx (B8.unpack addr)
          case eEndPoint of
-           Right (_port, ep, chan) -> return 
+           Right (_port, ep, chan) -> return
 	   	  ( TransportValid i{_transportEndPoints = Map.insert (localEndPointAddress ep) ep (_transportEndPoints i)}
                   , Right (ep, ctx, chan))
            Left _ -> return (v, Left $ TransportError NewEndPointFailed "Failed to create new endpoint.")
@@ -344,12 +334,12 @@
     addr = transportAddress transport
 
 -- | Asynchronous operation, shutdown of the remote end point may take a while
-apiCloseEndPoint :: ZMQTransport
+apiCloseEndPoint :: TransportInternals
                  -> LocalEndPoint
                  -> IO ()
 apiCloseEndPoint transport lep = mask_ $ do
---    printf "[%s][go] close endpoint\n"
---           (B8.unpack $ endPointAddressToByteString $ _localEndPointAddress lep)
+    -- we don't close endpoint here because other threads,
+    -- should be able to access local endpoint state
     old <- readMVar (localEndPointState lep)
     case old of
       LocalEndPointValid (ValidLocalEndPoint x _ _ threadId _ _) -> do
@@ -358,10 +348,11 @@
           writeTMChan x EndPointClosed
           closeTMChan x
         Async.cancel threadId
-        void $ Async.waitCatch threadId         
+        void $ Async.waitCatch threadId
       LocalEndPointClosed -> return ()
+    void $ swapMVar (localEndPointState lep) LocalEndPointClosed
     modifyMVar_ (_transportState transport) $ \case
-      TransportClosed  -> return TransportClosed 
+      TransportClosed  -> return TransportClosed
       TransportValid v -> return $ TransportValid
         v{_transportEndPoints = Map.delete (localEndPointAddress lep) (_transportEndPoints v)}
 
@@ -370,14 +361,14 @@
                -> String
                -> IO (Either (TransportError NewEndPointErrorCode) (Int,LocalEndPoint, TMChan Event))
 endPointCreate params ctx addr = do
-    em <- try $ do 
+    em <- try $ do
       pull <- ZMQ.socket ctx ZMQ.Pull
-      case authorizationType params of
-          ZMQNoAuth -> return ()
-          ZMQAuthPlain{} -> do
+      case zmqSecurityMechanism params of
+          Nothing -> return ()
+          Just SecurityPlain{} -> do
               ZMQ.setPlainServer True pull
-      ZMQ.setSendHighWM (ZMQ.restrict (highWaterMark params)) pull
-      port <- ZMQ.bindFromRangeRandom pull addr (minPort params) (maxPort params) (maxTries params)
+      ZMQ.setSendHighWM (ZMQ.restrict (zmqHighWaterMark params)) pull
+      port <- ZMQ.bindRandomPort pull addr
       return (port, pull)
     case em of
       Right (port, pull) -> (do
@@ -387,19 +378,19 @@
                                  <*> pure port
           opened <- newIORef True
           mask $ \restore -> do
-              thread <- Async.async $ (restore (receiver pull lep chOut)) 
+              thread <- Async.async $ (restore (receiver pull lep chOut))
                                `finally` finalizeEndPoint lep port pull
-              putMVar (localEndPointState lep) $ LocalEndPointValid 
+              putMVar (localEndPointState lep) $ LocalEndPointValid
                 (ValidLocalEndPoint chOut (Counter 0 Map.empty) Map.empty thread opened Map.empty)
               return $ Right (port, lep, chOut))
-          `onException` (ZMQ.close pull)
+          `onException` (ZMQ.closeZeroLinger pull)
       Left (_e::SomeException)  -> do
           return $ Left $ TransportError NewEndPointInsufficientResources "no free sockets"
   where
 
     finalizer pull ourEp = forever $ do
       (cmd:_) <- ZMQ.receiveMulti pull
-      case decode' cmd of
+      mask_ $ case decode' cmd of
         MessageEndPointCloseOk theirAddress ->  getRemoteEndPoint ourEp theirAddress >>= \case
           Nothing -> return ()
           Just rep -> do
@@ -416,14 +407,8 @@
       case decode' cmd of
         MessageData idx -> atomically $ writeTMChan chan (Received idx msgs)
         MessageConnect theirAddress -> do
---          printf "[%s] message connect from %s\n"
---                 (B8.unpack $ endPointAddressToByteString ourAddr)
---                 (B8.unpack $ endPointAddressToByteString theirAddress)
           void $ createOrGetRemoteEndPoint params ctx ourEp theirAddress
         MessageInitConnection theirAddress theirId rel -> do
---        printf "[%s] message init connection from %s\n"
---                (B8.unpack $ endPointAddressToByteString ourAddr)
---                (B8.unpack $ endPointAddressToByteString theirAddress)
           join $ do
             modifyMVar (localEndPointState ourEp) $ \case
                 LocalEndPointValid v ->
@@ -468,7 +453,7 @@
             register i RemoteEndPointClosed = do
               atomically $ do
                 writeTMChan chan (ConnectionOpened i rel theirAddress)
-                writeTMChan chan (ConnectionClosed i)                        
+                writeTMChan chan (ConnectionClosed i)
               return RemoteEndPointClosed
             register i (RemoteEndPointClosing x) = do
               atomically $ do
@@ -482,9 +467,6 @@
               return $ RemoteEndPointValid
                 v{_remoteEndPointIncommingConnections = Set.insert i s}
         MessageCloseConnection idx -> join $ do
---          printf "[%s] message close connection: %i\n"
---                 (B8.unpack $ endPointAddressToByteString ourAddr)
---                 idx
           modifyMVar (localEndPointState ourEp) $ \case
             LocalEndPointValid v ->
                 case idx `Map.lookup` m of
@@ -503,12 +485,8 @@
                 (Counter i m) = _localEndPointConnections v
 	    LocalEndPointClosed -> return (LocalEndPointClosed, return ())
         MessageInitConnectionOk theirAddress ourId theirId -> do
---          printf "[%s] message init connection ok: %i -> %i\n"
---                 (B8.unpack $ endPointAddressToByteString ourAddr)
---                 ourId
---                 theirId
           join $ withMVar (localEndPointState ourEp) $ \case
-            LocalEndPointValid v -> 
+            LocalEndPointValid v ->
                 case theirAddress `Map.lookup` r of
                   Nothing  -> return (return ()) -- XXX: send message to the host
                   Just rep -> modifyMVar (remoteEndPointState rep) $ \case
@@ -520,7 +498,7 @@
                           Nothing -> return (t, return ())     -- XXX: send message to the hostv
                           Just c  -> mask_ $ do
                             return (RemoteEndPointValid (ValidRemoteEndPoint sock (Counter x (ourId `Map.delete` m)) s (z+1))
-                                   , do 
+                                   , do
                                         modifyMVar_ (connectionState c) $ \case
                                           ZMQConnectionFailed -> return ZMQConnectionFailed
                                           ZMQConnectionInit -> return $ ZMQConnectionValid (ValidZMQConnection (Just sock) theirId)
@@ -530,12 +508,10 @@
                                           ZMQConnectionValid _ -> throwM $ InvariantViolation "RemoteEndPoint should be closed"
                                         void $ tryPutMVar (connectionReady c) ()
                                    )
-                    RemoteEndPointPending p -> return (RemoteEndPointPending p, undefined)
-              where 
+                    RemoteEndPointPending p -> return (RemoteEndPointPending p, throwM $ InvariantViolation "RemoteEndPoint should be closed")
+              where
                 r = _localEndPointRemotes v
             LocalEndPointClosed -> return $ return ()
---          printf "[%s] message init connection ok                      [ok]\n"
---                          (B8.unpack $ endPointAddressToByteString ourAddr)
         MessageEndPointClose theirAddress True -> getRemoteEndPoint ourEp theirAddress >>= \case
           Nothing  -> return ()
           Just rep -> do
@@ -559,9 +535,7 @@
             closeRemoteEndPoint ourEp rep state
       where
         ourAddr = localEndPointAddress ourEp
-    finalizeEndPoint ourEp port pull = do
---      printf "[%s] finalize-end-point\n"
---             (B8.unpack $ endPointAddressToByteString $ _localEndPointAddress ourEp)
+    finalizeEndPoint ourEp _port pull = do
       join $ withMVar (localEndPointState ourEp) $ \case
         LocalEndPointClosed  -> afterP ()
         LocalEndPointValid v -> do
@@ -571,14 +545,14 @@
             void $ Async.mapConcurrently (remoteEndPointClose False ourEp)
                  $ _localEndPointRemotes v
             Async.cancel tid
-            ZMQ.unbind pull (addr ++ ":" ++ show port)
-            ZMQ.close pull
+            void $ Async.waitCatch tid
+            ZMQ.closeZeroLinger pull
       void $ swapMVar (localEndPointState ourEp) LocalEndPointClosed
 
 apiSend :: ZMQConnection -> [ByteString] -> IO (Either (TransportError SendErrorCode) ())
 #ifdef UNSAFE_SEND
 apiSend c@(ZMQConnection l e _ s _) b = mask_ $ join $ withMVar s $ \case
-    ZMQConnectionInit   -> return $ yield >> apiSend c b 
+    ZMQConnectionInit   -> return $ yield >> apiSend c b
     ZMQConnectionClosed -> afterP $ Left $ TransportError SendClosed "Connection is closed"
     ZMQConnectionFailed -> afterP $ Left $ TransportError SendFailed "Connection is failed"
     ZMQConnectionValid (ValidZMQConnection (Just ch) idx) -> do
@@ -601,7 +575,7 @@
       else afterP $ Left $ TransportError SendFailed "Connection broken."
     _ -> afterP $ Left $ TransportError SendFailed "Incorrect channel."
 #else
-apiSend (ZMQConnection l e _ s _) b = mask_ $ do 
+apiSend (ZMQConnection l e _ s _) b = mask_ $ do
    result <- trySome inner
    case result of
      Left ex -> do
@@ -633,7 +607,7 @@
                 closeRemoteEndPoint l e z
             return $ Left $ TransportError SendFailed "Connection broken."
    cleanup = do
-     void $ cleanupRemoteEndPoint l e 
+     void $ cleanupRemoteEndPoint l e
        (Just $ \v -> ZMQ.send (_remoteEndPointChan v) [] $ encode' (MessageEndPointClose (localEndPointAddress l) False))
      onValidEndPoint l $ \v -> atomically $ do
        writeTMChan (_localEndPointChan v) $ ErrorEvent $ TransportError
@@ -678,19 +652,16 @@
            -> ConnectHints
            -> IO (Either (TransportError ConnectErrorCode) Connection)
 apiConnect params ctx ourEp theirAddr reliability _hints = do
---  printf "[%s] apiConnect to %s\n"
---       (B8.unpack $ endPointAddressToByteString $ _localEndPointAddress ourEp)
---       (B8.unpack $ endPointAddressToByteString theirAddr)
     eRep <- createOrGetRemoteEndPoint params ctx ourEp theirAddr
     case eRep of
       Left{} -> return $ Left $ TransportError ConnectFailed "LocalEndPoint is closed."
-      Right rep -> do 
+      Right rep -> do
         conn <- ZMQConnection <$> pure ourEp
                               <*> pure rep
                               <*> pure reliability
                               <*> newMVar ZMQConnectionInit
                               <*> newEmptyMVar
-        let apiConn = Connection 
+        let apiConn = Connection
               { send = apiSend conn
               , close = apiClose conn
               }
@@ -712,7 +683,7 @@
                    , return $ Left $ TransportError ConnectFailed "RemoteEndPoint failed.")
   where
     waitReady conn apiConn = join $ withMVar (connectionState conn) $ \case
-      ZMQConnectionInit{}   -> return $ yield {-readMVar (connectionReady conn)-} >> waitReady conn apiConn
+      ZMQConnectionInit{}   -> return $ yield >> waitReady conn apiConn
       ZMQConnectionValid{}  -> afterP $ Right apiConn
       ZMQConnectionFailed{} -> afterP $ Left $ TransportError ConnectFailed "Connection failed."
       ZMQConnectionClosed{} -> throwM $ InvariantViolation "Connection closed."
@@ -740,9 +711,6 @@
                           -> EndPointAddress
                           -> IO (Either ZMQError RemoteEndPoint)
 createOrGetRemoteEndPoint params ctx ourEp theirAddr = join $ do
---    printf "[%s] apiConnect to %s\n"
---           saddr
---           (B8.unpack $ endPointAddressToByteString theirAddr)
     modifyMVar (localEndPointState ourEp) $ \case
       LocalEndPointValid v@(ValidLocalEndPoint _ _ m _ o _) -> do
         opened <- readIORef o
@@ -757,17 +725,14 @@
         else return (LocalEndPointValid v, return $ Left $ IncorrectState "EndPointClosing")
       LocalEndPointClosed ->
         return  ( LocalEndPointClosed
-                , return $ Left $ IncorrectState "EndPoint is closed" 
+                , return $ Left $ IncorrectState "EndPoint is closed"
                 )
   where
     create v m = do
---    printf "[%s] apiConnect: remoteEndPoint not found, creating %s\n"
---           (B8.unpack $ endPointAddressToByteString $ _localEndPointAddress ourEp)
---           (B8.unpack $ endPointAddressToByteString theirAddr)
       push <- ZMQ.socket ctx ZMQ.Push
-      case authorizationType params of
-          ZMQNoAuth -> return ()
-          ZMQAuthPlain p u -> do
+      case zmqSecurityMechanism params of
+          Nothing -> return ()
+          Just (SecurityPlain p u) -> do
               ZMQ.setPlainPassword (ZMQ.restrict p) push
               ZMQ.setPlainUserName (ZMQ.restrict u) push
       state <- newMVar . RemoteEndPointPending =<< newIORef []
@@ -804,7 +769,7 @@
             RemoteEndPointFailed    -> return RemoteEndPointFailed
 
 -- | Close all endpoint connections, return previous state in case
--- if it was alive, for further cleanup actions. 
+-- if it was alive, for further cleanup actions.
 cleanupRemoteEndPoint :: LocalEndPoint
                       -> RemoteEndPoint
                       -> (Maybe (ValidRemoteEndPoint -> IO ()))
@@ -817,13 +782,13 @@
         RemoteEndPointValid w -> do
           let (Counter _ cn) = _remoteEndPointPendingConnections w
           traverse_ (\c -> void $ swapMVar (connectionState c) ZMQConnectionFailed) cn
-          cn' <- foldM 
+          cn' <- foldM
                (\(Counter i' cn') idx -> case idx `Map.lookup` cn of
-                   Nothing -> return (Counter i' cn')  
+                   Nothing -> return (Counter i' cn')
                    Just c  -> do
                      void $ swapMVar (connectionState c) ZMQConnectionFailed
                      return $ Counter i' (Map.delete idx cn')
-               ) 
+               )
                (_localEndPointConnections v)
                (Set.toList (_remoteEndPointIncommingConnections w))
           case actions of
@@ -843,7 +808,7 @@
    step1 = modifyMVar_ (localEndPointState lep) $ \case
      LocalEndPointValid v -> return $
        LocalEndPointValid v{_localEndPointRemotes=Map.delete
-                               (remoteEndPointAddress rep) 
+                               (remoteEndPointAddress rep)
                                (_localEndPointRemotes v)}
      c -> return c
    step2 (RemoteEndPointValid v) = do
@@ -856,19 +821,16 @@
 
 remoteEndPointClose :: Bool -> LocalEndPoint -> RemoteEndPoint -> IO ()
 remoteEndPointClose silent lep rep = do
---   printf "[???] remoteEndPointClose %s\n"          
---          (B8.unpack $ endPointAddressToByteString $ lAddr)
---          (B8.unpack $ endPointAddressToByteString $ remoteEndPointAddress rep)
    modifyIORef (remoteEndPointOpened rep) (const False)
    join $ modifyMVar (remoteEndPointState rep) $ \o -> case o of
      RemoteEndPointFailed        -> return (o, return ())
      RemoteEndPointClosed        -> return (o, return ())
      RemoteEndPointClosing (ClosingRemoteEndPoint _ l) -> return (o, void $ readMVar l)
-     RemoteEndPointPending _ -> closing undefined o -- XXX: store socket, or delay
+     RemoteEndPointPending _ -> closing (error "Pending actions should not be executed") o -- XXX: store socket, or delay
      RemoteEndPointValid v   -> closing (_remoteEndPointChan v) o
  where
    closing sock old = do
-     lock <- newEmptyMVar  
+     lock <- newEmptyMVar
      return (RemoteEndPointClosing (ClosingRemoteEndPoint sock lock), go lock old)
    go lock old@(RemoteEndPointValid (ValidRemoteEndPoint c _ s i)) = do
      -- close all connections
@@ -889,15 +851,15 @@
        void $ Async.race (readMVar lock) (threadDelay 1000000)
        void $ tryPutMVar lock ()
        closeRemoteEndPoint lep rep old
-       return ()       
-   go _ _ = return ()    
+       return ()
+   go _ _ = return ()
 
 connectionCleanup :: RemoteEndPoint -> ConnectionId -> IO ()
 connectionCleanup rep cid = modifyMVar_ (remoteEndPointState rep) $ \case
    RemoteEndPointValid v -> return $
       RemoteEndPointValid v{_remoteEndPointIncommingConnections = Set.delete cid (_remoteEndPointIncommingConnections v)}
    c -> return c
-         
+
 encode' :: Binary a => a -> ByteString
 encode' = B.concat . BL.toChunks . encode
 
@@ -925,11 +887,11 @@
   Right x -> return $ Right x
 
 -- | Break endpoint connection, all endpoints that will be affected.
-breakConnection :: ZMQTransport
+breakConnection :: TransportInternals
                 -> EndPointAddress
                 -> EndPointAddress
                 -> IO ()
-breakConnection zmqt from to = Foldable.sequence_ <=<  withMVar (_transportState zmqt) $ \case
+breakConnection zmqt _from to = Foldable.sequence_ <=<  withMVar (_transportState zmqt) $ \case
     TransportValid v -> Traversable.forM (_transportEndPoints v) $ \x ->
       withMVar (localEndPointState x) $ \case
         LocalEndPointValid w -> return $ Foldable.sequence_ $ flip Map.mapWithKey (_localEndPointRemotes w) $ \key rep ->
@@ -949,7 +911,7 @@
 
 -- | Break connection between two endpoints, other endpoints on the same
 -- remove will not be affected.
-breakConnectionEndPoint :: ZMQTransport
+breakConnectionEndPoint :: TransportInternals
                         -> EndPointAddress
                         -> EndPointAddress
                         -> IO ()
@@ -973,7 +935,7 @@
       TransportClosed -> afterP ()
 
 -- | Configure socket after creation
-unsafeConfigurePush :: ZMQTransport
+unsafeConfigurePush :: TransportInternals
                     -> EndPointAddress
                     -> EndPointAddress
                     -> (ZMQ.Socket ZMQ.Push -> IO ())
@@ -988,23 +950,19 @@
         LocalEndPointClosed   -> return ()
     TransportClosed -> return ()
 
-apiNewMulticastGroup :: ZMQParameters -> ZMQTransport -> LocalEndPoint -> IO ( Either (TransportError NewMulticastGroupErrorCode) MulticastGroup)
-apiNewMulticastGroup params zmq lep = withMVar (_transportState zmq) $ \case
+apiNewMulticastGroup :: ZMQParameters -> TransportInternals -> LocalEndPoint -> IO ( Either (TransportError NewMulticastGroupErrorCode) MulticastGroup)
+apiNewMulticastGroup _params zmq lep = withMVar (_transportState zmq) $ \case
   TransportClosed -> return $ Left $ TransportError NewMulticastGroupFailed "Transport is closed."
   TransportValid vt -> modifyMVar (localEndPointState lep) $ \case
     LocalEndPointClosed -> return (LocalEndPointClosed, Left $ TransportError NewMulticastGroupFailed "Transport is closed.")
     LocalEndPointValid vl -> mask_ $ do
       (pub, portPub, rep, portRep, wrkThread) <- mkPublisher vt
 
-      let addr = MulticastAddress $ 
+      let addr = MulticastAddress $
            B8.concat [transportAddress zmq,":",B8.pack (show portPub), ":",B8.pack (show portRep)]
           subAddr = extractPubAddress addr
           repAddr = extractRepAddress addr
 
---      printf "new multicastaddress: %s\n"  (show addr)
---      printf "[%s] subscribe address: %s\n" (show addr) (B8.unpack subAddr)
---      printf "[%s] send address: %s\n" (show addr) (B8.unpack repAddr)
-
       -- subscriber api
       sub <- ZMQ.socket (_transportContext vt) ZMQ.Sub
       ZMQ.connect sub (B8.unpack subAddr)
@@ -1020,13 +978,13 @@
               (ctrl: msg) <- ZMQ.receiveMulti sub
               if B8.null ctrl
               then do opened <- readIORef subscribed
-                      when opened $ atomically $ writeTMChan (_localEndPointChan vl) 
+                      when opened $ atomically $ writeTMChan (_localEndPointChan vl)
                                                $ ReceivedMulticast addr msg
                       return True
               else return False
         return (tid', mgroup)
       return ( LocalEndPointValid vl
-             , Right $ MulticastGroup 
+             , Right $ MulticastGroup
                 { multicastAddress = addr
                 , deleteMulticastGroup = apiDeleteMulticastGroupLocal (multicastGroupState mgroup) lep addr rep repAddr pub sub subAddr (Just subThread) wrkThread
                 , maxMsgSize = Nothing
@@ -1039,16 +997,16 @@
   where
     mkPublisher vt = do
       pub <- ZMQ.socket (_transportContext vt) ZMQ.Pub
-      portPub <- ZMQ.bindFromRangeRandom pub (B8.unpack $ transportAddress zmq) (minPort params) (maxPort params) (maxTries params)
+      portPub <- ZMQ.bindRandomPort pub (B8.unpack $ transportAddress zmq)
       rep <- ZMQ.socket (_transportContext vt) ZMQ.Rep
-      portRep <- ZMQ.bindFromRangeRandom rep (B8.unpack $ transportAddress zmq) (minPort params) (maxPort params) (maxTries params)
+      portRep <- ZMQ.bindRandomPort rep (B8.unpack $ transportAddress zmq)
       wrkThread <- Async.async $ forever $ do
         msg <- ZMQ.receiveMulti rep
         ZMQ.sendMulti pub $ "" :| msg
         ZMQ.send rep [] "OK"
       return (pub,portPub,rep,portRep, wrkThread)
 
-apiResolveMulticastGroup :: ZMQTransport -> LocalEndPoint -> MulticastAddress -> IO (Either (TransportError ResolveMulticastGroupErrorCode) MulticastGroup)
+apiResolveMulticastGroup :: TransportInternals -> LocalEndPoint -> MulticastAddress -> IO (Either (TransportError ResolveMulticastGroupErrorCode) MulticastGroup)
 apiResolveMulticastGroup zmq lep addr = withMVar (_transportState zmq) $ \case
   TransportClosed -> return $ Left $ TransportError ResolveMulticastGroupFailed  "Transport is closed."
   TransportValid vt -> modifyMVar (localEndPointState lep) $ \case
@@ -1070,7 +1028,7 @@
           (ctrl: msg) <- ZMQ.receiveMulti sub
           if B8.null ctrl
           then do opened <- readIORef subscribed
-                  when opened $ atomically $ writeTMChan (_localEndPointChan vl) 
+                  when opened $ atomically $ writeTMChan (_localEndPointChan vl)
                                            $ ReceivedMulticast addr msg
                   return True
           else do apiDeleteMulticastGroupRemote v lep addr req reqAddr sub subAddr Nothing
@@ -1079,13 +1037,13 @@
 
       return ( LocalEndPointValid
                 vl{_localEndPointMulticastGroups = Map.insert addr mgroup (_localEndPointMulticastGroups vl)}
-             , Right $ MulticastGroup 
+             , Right $ MulticastGroup
                 { multicastAddress = addr
                 , deleteMulticastGroup =
                     apiDeleteMulticastGroupRemote (multicastGroupState mgroup) lep addr req reqAddr sub subAddr (Just subThread)
                 , maxMsgSize = Nothing
                 , multicastSend = apiMulticastSendRemote mgroup req
-                , multicastSubscribe = apiMulticastSubscribe mgroup 
+                , multicastSubscribe = apiMulticastSubscribe mgroup
                 , multicastUnsubscribe = apiMulticastUnsubscribe mgroup
                 , multicastClose = apiMulticastClose
                 }
@@ -1095,31 +1053,29 @@
     subAddr = extractPubAddress addr
 
 
-apiDeleteMulticastGroupRemote :: MVar MulticastGroupState -> LocalEndPoint -> MulticastAddress 
+apiDeleteMulticastGroupRemote :: MVar MulticastGroupState -> LocalEndPoint -> MulticastAddress
     -> ZMQ.Socket ZMQ.Req -> ByteString -> ZMQ.Socket ZMQ.Sub -> ByteString -> Maybe (Async.Async ())
     -> IO ()
-apiDeleteMulticastGroupRemote mstate lep addr req reqAddr sub subAddr mtid = mask_ $ do
+apiDeleteMulticastGroupRemote mstate lep addr req _reqAddr sub _subAddr mtid = mask_ $ do
   Foldable.traverse_ (\tid -> Async.cancel tid >> void (Async.waitCatch tid)) mtid
   modifyMVar_ mstate $ \case
     MulticastGroupClosed -> return MulticastGroupClosed
     MulticastGroupValid v -> do
       modifyIORef (_multicastGroupSubscribed v) (const False)
---      ZMQ.disconnect req (B8.unpack reqAddr)
-      ZMQ.close req
+      ZMQ.closeZeroLinger req
       ZMQ.unsubscribe sub ""
---      ZMQ.disconnect sub (B8.unpack subAddr)
-      ZMQ.close sub
+      ZMQ.closeZeroLinger sub
       return MulticastGroupClosed
   modifyMVar_ (localEndPointState lep) $ \case
     LocalEndPointClosed -> return LocalEndPointClosed
     LocalEndPointValid v -> return $ LocalEndPointValid
       v{_localEndPointMulticastGroups = Map.delete addr (_localEndPointMulticastGroups v)}
 
-apiDeleteMulticastGroupLocal :: MVar MulticastGroupState -> LocalEndPoint -> MulticastAddress 
+apiDeleteMulticastGroupLocal :: MVar MulticastGroupState -> LocalEndPoint -> MulticastAddress
     -> ZMQ.Socket ZMQ.Rep -> ByteString -> ZMQ.Socket ZMQ.Pub -> ZMQ.Socket ZMQ.Sub -> ByteString -> Maybe (Async.Async ())
-    -> Async.Async () 
+    -> Async.Async ()
     -> IO ()
-apiDeleteMulticastGroupLocal mstate lep addr rep repAddr pub sub pubAddr mtid wrk = mask_ $ do
+apiDeleteMulticastGroupLocal mstate lep addr rep _repAddr pub sub _pubAddr mtid wrk = mask_ $ do
    Foldable.traverse_ (\tid -> Async.cancel tid >> void (Async.waitCatch tid)) mtid
    void $ Async.cancel wrk >> Async.waitCatch wrk
    modifyMVar_ mstate $ \case
@@ -1128,13 +1084,10 @@
        modifyIORef (_multicastGroupSubscribed v) (const False)
        ZMQ.sendMulti pub $ "C" :| []
        threadDelay 50000
-       ZMQ.unbind rep (B8.unpack repAddr)
-       ZMQ.close rep
+       ZMQ.closeZeroLinger rep
        ZMQ.unsubscribe sub ""
---       ZMQ.disconnect sub (B8.unpack pubAddr)
-       ZMQ.close sub
-       ZMQ.unbind pub (B8.unpack pubAddr)
-       ZMQ.close pub
+       ZMQ.closeZeroLinger sub
+       ZMQ.closeZeroLinger pub
        return MulticastGroupClosed
    modifyMVar_ (localEndPointState lep) $ \case
      LocalEndPointClosed -> return LocalEndPointClosed
@@ -1168,7 +1121,32 @@
   MulticastGroupValid v -> modifyIORef (_multicastGroupSubscribed v) (const False)
 
 apiMulticastClose :: IO ()
-apiMulticastClose = return () 
+apiMulticastClose = return ()
+
+-- $cleanup
+-- Cleanup API is prepared to store cleanup actions for example socket
+-- close for objects that uses network-transport-zeromq resources. Theese
+-- actions will be fired on transport close.
+-- If you need to clean resource beforehead use 'applyCleanupAction'.
+
+-- | Register action that will be fired when transport will be closed.
+registerCleanupAction :: TransportInternals -> IO () -> IO (Maybe Unique)
+registerCleanupAction zmq fn = withMVar (_transportState zmq) $ \case
+  TransportValid v -> registerValidCleanupAction v fn
+  TransportClosed  -> return Nothing
+
+-- | Register action on a locked transport.
+registerValidCleanupAction :: ValidTransportState -> IO () -> IO (Maybe Unique)
+registerValidCleanupAction (ValidTransportState _ _ _ im) fn = Just <$> do
+    u <- newUnique
+    atomicModifyIORef' im (\m -> (IntMap.insert (hashUnique u) fn m, u))
+
+-- | Call cleanup action before transport close.
+applyCleanupAction :: TransportInternals -> Unique -> IO ()
+applyCleanupAction zmq u = withMVar (_transportState zmq) $ \case
+  TransportValid (ValidTransportState _ _ _ im) -> mask_ $
+    traverse_ id =<< atomicModifyIORef' im (\m -> (IntMap.delete (hashUnique u) m, IntMap.lookup (hashUnique u) m))
+  TransportClosed -> return ()
 
 extractRepAddress :: MulticastAddress -> ByteString
 extractRepAddress (MulticastAddress bs) = B8.concat [a,":",p]
diff --git a/src/Network/Transport/ZMQ/Internal.hs b/src/Network/Transport/ZMQ/Internal.hs
new file mode 100644
--- /dev/null
+++ b/src/Network/Transport/ZMQ/Internal.hs
@@ -0,0 +1,62 @@
+-- |
+-- Copyright: (C) 2014 EURL Tweag
+--
+
+{-# LANGUAGE OverloadedStrings #-}
+module Network.Transport.ZMQ.Internal
+  ( bindRandomPort
+  , authManager
+  , closeZeroLinger
+  )
+  where
+
+import           Control.Monad ( forever )
+import           Control.Monad.Catch
+      ( mask
+      , finally
+      )
+import           Control.Concurrent.Async
+import           Data.ByteString ( ByteString )
+import           Data.List.NonEmpty ( NonEmpty(..) )
+
+import           System.ZMQ4
+
+-- | Bind socket to the random port.
+bindRandomPort :: Socket t
+               -> String -- ^ Address
+               -> IO Int
+bindRandomPort sock addr = do
+    bind sock $ addr++":0"
+    fmap (read . last . split (/=':')) $ lastEndpoint sock
+
+-- | One possible password authentification
+authManager :: Context -> ByteString -> ByteString -> IO (Async ())
+authManager ctx user pass = do
+    req <- socket ctx Rep
+    bind req "inproc://zeromq.zap.01"
+    mask $ \restore ->
+      async $ (restore $ forever $ do
+          ("1.0":requestId:_domain:_ipAddress:_identity:mech:xs) <- receiveMulti req
+          case mech of
+            "PLAIN" -> case xs of
+               (pass':user':_)
+                 | user == user' && pass == pass' -> do
+                    sendMulti req $ "1.0" :| [requestId, "200", "OK", "", ""]
+                 | otherwise -> sendMulti req $ "1.0" :| [requestId, "400", "Credentials are not implemented", "", ""]
+               _ -> sendMulti req $ "1.0" :| [requestId, "500", "Method not implemented", "", ""]
+            _ -> sendMulti req $ "1.0" :| [requestId, "500", "Method not implemented", "", ""])
+            `finally` (close req)
+
+-- | Close socket immideately.
+closeZeroLinger :: Socket a -> IO ()
+closeZeroLinger sock = do
+  setLinger (restrict (0::Int)) sock
+  close sock
+
+split :: (Char -> Bool) -> String -> [String]
+split f = go
+  where
+    go [] = []
+    go s  = case span f s of
+                (p,[]) -> [p]
+                (p,_:ss) -> p:go ss
diff --git a/src/Network/Transport/ZMQ/Internal/Types.hs b/src/Network/Transport/ZMQ/Internal/Types.hs
new file mode 100644
--- /dev/null
+++ b/src/Network/Transport/ZMQ/Internal/Types.hs
@@ -0,0 +1,232 @@
+-- |
+-- Copyright: (C) 2014 EURL Tweag
+--
+
+module Network.Transport.ZMQ.Internal.Types
+  ( ZMQParameters(..)
+  , SecurityMechanism(..)
+  , defaultZMQParameters
+    -- * Internal types
+  , TransportInternals(..)
+  , TransportState(..)
+  , ValidTransportState(..)
+    -- ** RemoteEndPoint
+  , RemoteEndPoint(..)
+  , RemoteEndPointState(..)
+  , ValidRemoteEndPoint(..)
+  , ClosingRemoteEndPoint(..)
+    -- ** LocalEndPoint
+  , LocalEndPoint(..)
+  , LocalEndPointState(..)
+  , ValidLocalEndPoint(..)
+    -- ** ZeroMQ connection
+  , ZMQConnection(..)
+  , ZMQConnectionState(..)
+  , ValidZMQConnection(..)
+    -- ** ZeroMQ multicast
+  , ZMQMulticastGroup(..)
+  , MulticastGroupState(..)
+  , ValidMulticastGroup(..)
+    -- * Internal data structures
+  , Counter(..)
+  , nextElement
+  , nextElement'
+  , nextElementM
+  , nextElementM'
+  ) where
+
+import Control.Concurrent.Async
+import Control.Concurrent.MVar
+import Control.Concurrent.STM.TMChan
+import Data.Word
+import Data.ByteString
+import Data.IORef
+import Data.IntMap (IntMap)
+import Data.Map.Strict (Map)
+import qualified Data.Map.Strict  as M
+import           Data.Set
+     ( Set
+     )
+import           System.ZMQ4
+      ( Socket
+      , Push
+      )
+
+import Network.Transport
+
+import qualified System.ZMQ4 as ZMQ
+
+
+-- | Parameters for ZeroMQ connection.
+data ZMQParameters = ZMQParameters
+  { zmqHighWaterMark     :: Word64 -- uint64_t
+  , zmqSecurityMechanism :: Maybe SecurityMechanism
+  }
+
+defaultZMQParameters :: ZMQParameters
+defaultZMQParameters = ZMQParameters
+    { zmqHighWaterMark     = 0
+    , zmqSecurityMechanism = Nothing
+    }
+
+-- | A ZeroMQ "security mechanism".
+data SecurityMechanism
+  = SecurityPlain
+  { plainPassword :: ByteString
+  , plainUsername :: ByteString
+  } -- ^ Clear-text authentication, using a (username, password) pair.
+
+type TransportAddress = ByteString
+
+-- | Transport data type.
+data TransportInternals = TransportInternals
+  { transportAddress :: !TransportAddress
+  -- ^ Transport address (used as identifier).
+  , _transportState  :: !(MVar TransportState)
+  -- ^ Internal state.
+  }
+
+-- | Transport state.
+data TransportState
+  = TransportValid !ValidTransportState         -- ^ Transport is in active state.
+  | TransportClosed                             -- ^ Transport is closed.
+
+-- | Transport state.
+data ValidTransportState = ValidTransportState
+  { _transportContext   :: !ZMQ.Context
+  , _transportEndPoints :: !(Map EndPointAddress LocalEndPoint)
+  , _transportAuth      :: !(Maybe (Async ()))
+  , _transportSockets   :: !(IORef (IntMap (IO ())))
+  -- ^ A set of cleanup actions.
+  }
+
+data LocalEndPoint = LocalEndPoint
+  { localEndPointAddress :: !EndPointAddress
+  , localEndPointState   :: !(MVar LocalEndPointState)
+  , localEndPointPort    :: !Int
+  }
+
+data LocalEndPointState
+  = LocalEndPointValid !ValidLocalEndPoint
+  | LocalEndPointClosed
+
+data ValidLocalEndPoint = ValidLocalEndPoint
+  { _localEndPointChan        :: !(TMChan Event)
+    -- ^ channel for n-t - user communication
+  , _localEndPointConnections :: !(Counter ConnectionId ZMQConnection)
+    -- ^ list of incomming connections
+  , _localEndPointRemotes     :: !(Map EndPointAddress RemoteEndPoint)
+    -- ^ list of remote end points
+  , _localEndPointThread      :: !(Async ())
+    -- ^ thread id
+  , _localEndPointOpened      :: !(IORef Bool)
+    -- ^ is remote endpoint opened
+  , _localEndPointMulticastGroups :: !(Map MulticastAddress ZMQMulticastGroup)
+    -- ^ list of multicast nodes
+  }
+
+data ZMQConnection = ZMQConnection
+  { connectionLocalEndPoint  :: !LocalEndPoint
+  , connectionRemoteEndPoint :: !RemoteEndPoint
+  , connectionReliability    :: !Reliability
+  , connectionState          :: !(MVar ZMQConnectionState)
+  , connectionReady          :: !(MVar ())
+  }
+
+data ZMQConnectionState
+  = ZMQConnectionInit
+  | ZMQConnectionValid !ValidZMQConnection
+  | ZMQConnectionClosed
+  | ZMQConnectionFailed
+
+data ValidZMQConnection = ValidZMQConnection
+  { _connectionSocket         :: !(Maybe (ZMQ.Socket ZMQ.Push))
+  , _connectionId             :: !Word64
+  }
+
+data RemoteEndPoint = RemoteEndPoint
+  { remoteEndPointAddress :: !EndPointAddress
+  , remoteEndPointState   :: !(MVar RemoteEndPointState)
+  , remoteEndPointOpened  :: !(IORef Bool)
+  }
+
+data RemoteEndPointState
+  = RemoteEndPointValid ValidRemoteEndPoint
+  | RemoteEndPointClosed
+  | RemoteEndPointFailed
+  | RemoteEndPointPending (IORef [RemoteEndPointState -> IO RemoteEndPointState])
+  | RemoteEndPointClosing ClosingRemoteEndPoint
+
+data ValidRemoteEndPoint = ValidRemoteEndPoint
+  { _remoteEndPointChan                 :: !(Socket Push)
+  , _remoteEndPointPendingConnections   :: !(Counter ConnectionId ZMQConnection)
+  , _remoteEndPointIncommingConnections :: !(Set ConnectionId)
+  , _remoteEndPointOutgoingCount        :: !Int
+  }
+
+data ClosingRemoteEndPoint = ClosingRemoteEndPoint
+  { _remoteEndPointClosingSocket :: !(Socket Push)
+  , _remoteEndPointDone :: !(MVar ())
+  }
+
+data ZMQMulticastGroup = ZMQMulticastGroup
+  { multicastGroupState         :: MVar MulticastGroupState
+  , multicastGroupClose         :: IO ()
+  }
+
+data MulticastGroupState
+  = MulticastGroupValid ValidMulticastGroup
+  | MulticastGroupClosed
+
+data ValidMulticastGroup = ValidMulticastGroup
+  { _multicastGroupSubscribed :: IORef Bool
+  }
+
+data Counter a b = Counter
+  { counterNext   :: !a
+  , counterValue  :: !(Map a b)
+  }
+
+nextElement :: (Enum a, Ord a)
+            => (b -> IO Bool)
+            -> b
+            -> Counter a b
+            -> IO (Counter a b, (a, b))
+nextElement t e c = nextElement' t (const e) c
+
+nextElement' :: (Enum a, Ord a)
+             => (b -> IO Bool)
+             -> (a -> b)
+             -> Counter a b
+             -> IO (Counter a b, (a,b))
+nextElement' t e c = nextElementM t (return . e) c
+
+nextElementM :: (Enum a, Ord a)
+             => (b -> IO Bool)
+             -> (a -> IO b)
+             -> Counter a b
+             -> IO (Counter a b, (a,b))
+nextElementM t me (Counter n m) =
+    case n' `M.lookup` m of
+      Nothing -> mv >>= \v' -> return (Counter n' (M.insert n' v' m), (n', v'))
+      Just v  -> t v >>= \case
+        True -> mv >>= \v' -> return (Counter n' (M.insert n' v' m), (n', v'))
+        False -> nextElementM t me (Counter n' m)
+  where
+    n' = succ n
+    mv = me n'
+
+nextElementM' :: (Enum a, Ord a)
+              => (b -> IO Bool)
+              -> (a -> IO (b,c))
+              -> Counter a b
+              -> IO (Counter a b, (a,c))
+nextElementM' t me (Counter n m) =
+    case n' `M.lookup` m of
+      Nothing -> mv >>= \(v',r) -> return (Counter n' (M.insert n' v' m), (n', r))
+      Just v  -> t v >>= \case
+        True -> mv >>= \(v',r) -> return (Counter n' (M.insert n' v' m), (n', r))
+        False -> nextElementM' t me (Counter n' m)
+  where
+    n' = succ n
+    mv = me n'
diff --git a/src/Network/Transport/ZMQ/Types.hs b/src/Network/Transport/ZMQ/Types.hs
deleted file mode 100644
--- a/src/Network/Transport/ZMQ/Types.hs
+++ /dev/null
@@ -1,222 +0,0 @@
--- |
--- Module:     Network.Transport.ZMQ.Types
--- Copyright:  (C) 2014, EURL Tweag
--- Licese:     BSD-3
---
-{-# LANGUAGE LambdaCase #-}
-module Network.Transport.ZMQ.Types
-    ( ZMQParameters(..)
-    , ZMQAuthType(..)
-    , defaultZMQParameters
-      -- * Internal types
-    , ZMQTransport(..)
-    , TransportState(..)
-    , ValidTransportState(..)
-      -- ** RemoteEndPoint
-    , RemoteEndPoint(..)
-    , RemoteEndPointState(..)
-    , ValidRemoteEndPoint(..)
-    , ClosingRemoteEndPoint(..)
-      -- ** LocalEndPoint
-    , LocalEndPoint(..)
-    , LocalEndPointState(..)
-    , ValidLocalEndPoint(..)
-      -- ** ZeroMQ connection
-    , ZMQConnection(..)
-    , ZMQConnectionState(..)
-    , ValidZMQConnection(..)
-      -- ** ZeroMQ multicast
-    , ZMQMulticastGroup(..)
-    , MulticastGroupState(..)
-    , ValidMulticastGroup(..)
-      -- * Internal data structures
-    , Counter(..)
-    , nextElement
-    , nextElement'
-    , nextElementM
-    , nextElementM'
-    ) where
-
-import Control.Concurrent.Async
-import Control.Concurrent.MVar
-import Control.Concurrent.STM.TMChan
-import Data.Word
-import Data.ByteString
-import Data.IORef
-import Data.Map.Strict (Map)
-import qualified Data.Map.Strict  as M
-import           Data.Set 
-     ( Set
-     )
-import           System.ZMQ4
-      ( Socket
-      , Push
-      )
-
-import Network.Transport
-
-import qualified System.ZMQ4 as ZMQ
-
-
--- | Parameters for ZeroMQ connection
-data ZMQParameters = ZMQParameters
-      { highWaterMark :: Word64 -- uint64_t
-      , authorizationType :: ZMQAuthType
-      , minPort       :: Int
-      , maxPort       :: Int
-      , maxTries      :: Int
-      }
--- High Watermark
-
-defaultZMQParameters :: ZMQParameters
-defaultZMQParameters = ZMQParameters
-      { highWaterMark = 0
-      , authorizationType = ZMQNoAuth
-      , minPort       = 40000
-      , maxPort       = 60000
-      , maxTries      = 10000
-      }
-
-data ZMQAuthType
-        = ZMQNoAuth
-        | ZMQAuthPlain
-            { zmqAuthPlainPassword :: ByteString
-            , zmqAutnPlainUserName :: ByteString
-            }
-
-type TransportAddress = ByteString
-
--- | Transport data type.
-data ZMQTransport = ZMQTransport
-    { transportAddress :: !TransportAddress
-    -- ^ Transport address (used as identifier).
-    , _transportState  :: !(MVar TransportState)
-    -- ^ Internal state.
-    }
-
--- | Transport state.
-data TransportState
-      = TransportValid !ValidTransportState         -- ^ Transport is in active state.
-      | TransportClosed                             -- ^ Transport is closed.
-
--- | Transport state.
-data ValidTransportState = ValidTransportState
-      { _transportContext   :: !ZMQ.Context
-      , _transportEndPoints :: !(Map EndPointAddress LocalEndPoint)
-      , _transportAuth      :: !(Maybe (Async ()))
-      }
-
-data LocalEndPoint = LocalEndPoint
-      { localEndPointAddress :: !EndPointAddress
-      , localEndPointState   :: !(MVar LocalEndPointState)
-      , localEndPointPort    :: !(Int)
-      }
-
-data LocalEndPointState
-      = LocalEndPointValid !ValidLocalEndPoint
-      | LocalEndPointClosed
-
-data ValidLocalEndPoint = ValidLocalEndPoint
-      { _localEndPointChan        :: !(TMChan Event)
-        -- ^ channel for n-t - user communication
-      , _localEndPointConnections :: !(Counter ConnectionId ZMQConnection)
-        -- ^ list of incomming connections
-      , _localEndPointRemotes     :: !(Map EndPointAddress RemoteEndPoint)
-        -- ^ list of remote end points
-      , _localEndPointThread      :: !(Async ())
-        -- ^ thread id
-      , _localEndPointOpened      :: !(IORef Bool)
-        -- ^ is remote endpoint opened
-      , _localEndPointMulticastGroups :: !(Map MulticastAddress ZMQMulticastGroup)
-        -- ^ list of multicast nodes
-      }
-
-data ZMQConnection = ZMQConnection
-      { connectionLocalEndPoint  :: !LocalEndPoint
-      , connectionRemoteEndPoint :: !RemoteEndPoint
-      , connectionReliability    :: !Reliability
-      , connectionState          :: !(MVar ZMQConnectionState)
-      , connectionReady          :: !(MVar ())
-      }
-
-data ZMQConnectionState
-      = ZMQConnectionInit
-      | ZMQConnectionValid !ValidZMQConnection
-      | ZMQConnectionClosed
-      | ZMQConnectionFailed
-
-data ValidZMQConnection = ValidZMQConnection
-      { _connectionSocket         :: !(Maybe (ZMQ.Socket ZMQ.Push))
-      , _connectionId             :: !Word64
-      }
-
-data RemoteEndPoint = RemoteEndPoint
-      { remoteEndPointAddress :: !EndPointAddress
-      , remoteEndPointState   :: !(MVar RemoteEndPointState) 
-      , remoteEndPointOpened  :: !(IORef Bool)
-      }
-
-data RemoteEndPointState
-      = RemoteEndPointValid ValidRemoteEndPoint
-      | RemoteEndPointClosed
-      | RemoteEndPointFailed
-      | RemoteEndPointPending (IORef [RemoteEndPointState -> IO RemoteEndPointState])
-      | RemoteEndPointClosing ClosingRemoteEndPoint
-
-data ValidRemoteEndPoint = ValidRemoteEndPoint
-      { _remoteEndPointChan                 :: !(Socket Push)
-      , _remoteEndPointPendingConnections   :: !(Counter ConnectionId ZMQConnection)
-      , _remoteEndPointIncommingConnections :: !(Set ConnectionId)
-      , _remoteEndPointOutgoingCount        :: !Int
-      }
-
-data ClosingRemoteEndPoint = ClosingRemoteEndPoint
-     { _remoteEndPointClosingSocket :: !(Socket Push)
-     , _remoteEndPointDone :: !(MVar ())
-     }
-
-data ZMQMulticastGroup = ZMQMulticastGroup
-      { multicastGroupState         :: MVar MulticastGroupState
-      , multicastGroupClose         :: IO ()
-      }
-
-data MulticastGroupState
-        = MulticastGroupValid ValidMulticastGroup
-        | MulticastGroupClosed
-
-data ValidMulticastGroup = ValidMulticastGroup
-      { _multicastGroupSubscribed :: IORef Bool
-      }
-
-data Counter a b = Counter { counterNext   :: !a
-                           , counterValue :: !(Map a b)
-                           }
-
-nextElement :: (Enum a, Ord a) => (b -> IO Bool) -> b -> Counter a b -> IO (Counter a b, (a, b))
-nextElement t e c = nextElement' t (const e) c
-
-nextElement' :: (Enum a, Ord a) => (b -> IO Bool) -> (a -> b) -> Counter a b -> IO (Counter a b, (a,b))
-nextElement' t e c = nextElementM t (return . e) c
-
-nextElementM :: (Enum a, Ord a) => (b -> IO Bool) -> (a -> IO b) -> Counter a b -> IO (Counter a b, (a,b))
-nextElementM t me (Counter n m) =
-    case n' `M.lookup` m of
-      Nothing -> mv >>= \v' -> return (Counter n' (M.insert n' v' m), (n', v'))
-      Just v  -> t v >>= \case
-        True -> mv >>= \v' -> return (Counter n' (M.insert n' v' m), (n', v'))
-        False -> nextElementM t me (Counter n' m)
-  where
-    n' = succ n
-    mv = me n'
-
-nextElementM' :: (Enum a, Ord a) => (b -> IO Bool) -> (a -> IO (b,c)) -> Counter a b -> IO (Counter a b, (a,c))
-nextElementM' t me (Counter n m) =
-    case n' `M.lookup` m of
-      Nothing -> mv >>= \(v',r) -> return (Counter n' (M.insert n' v' m), (n', r))
-      Just v  -> t v >>= \case
-        True -> mv >>= \(v',r) -> return (Counter n' (M.insert n' v' m), (n', r))
-        False -> nextElementM' t me (Counter n' m)
-  where
-    n' = succ n
-    mv = me n'
-
diff --git a/src/System/ZMQ4/Utils.hs b/src/System/ZMQ4/Utils.hs
deleted file mode 100644
--- a/src/System/ZMQ4/Utils.hs
+++ /dev/null
@@ -1,87 +0,0 @@
-{-# LANGUAGE OverloadedStrings #-}
-module System.ZMQ4.Utils
-  ( bindFromRangeRandom
-  , bindFromRangeRandomM
-  , authManager
-  , closeZeroLinger
-  )
-  where
-
-import Control.Monad ( forever )
-import Control.Monad.Catch
-import Control.Concurrent.Async
-import Data.ByteString
-import Data.List.NonEmpty
-import System.Random ( randomRIO )
-import Text.Printf
-
-
-import           System.ZMQ4.Monadic
-      ( ZMQ 
-      )
-import qualified System.ZMQ4.Monadic as M
-import           System.ZMQ4
-      ( Context
-      , errno
-      )
-import qualified System.ZMQ4         as P
-
--- | Bind socket to the random port in a given range.
-bindFromRangeRandomM :: M.Socket z t
-                     -> String -- ^ Address
-                     -> Int    -- ^ Min port
-                     -> Int    -- ^ Max port
-                     -> Int    -- ^ Max tries
-                     -> ZMQ z Int
-bindFromRangeRandomM sock addr mI mA tr = go tr
-    where
-      go 0 = error "!" -- XXX: throw correct error
-      go x = do
-        port   <- M.liftIO $ randomRIO (mI,mA)
-        result <- try $ M.bind sock (printf "%s:%i" addr port)
-        case result of
-            Left e
-              | errno e == -1 -> go (x - 1)
-              | otherwise -> throwM e
-            Right () -> return port
-
-bindFromRangeRandom :: P.Socket t
-                    -> String -- ^ Address
-                    -> Int    -- ^ Min port
-                    -> Int    -- ^ Max port
-                    -> Int    -- ^ Max tries
-                    -> IO Int
-bindFromRangeRandom sock addr mI mA tr = go tr
-    where
-      go 0 = error "!" -- XXX: throw correct error
-      go x = do
-        port   <- randomRIO (mI,mA)
-        result <- try $ P.bind sock (printf "%s:%i" addr port)
-        case result of
-            Left e
-              | errno e == -1 -> go (x - 1)
-              | otherwise -> throwM e
-            Right () -> return port
-
--- | One possible password authentification
-authManager :: P.Context -> ByteString -> ByteString -> IO (Async ())
-authManager ctx user pass = do
-    req <- P.socket ctx P.Rep
-    P.bind req "inproc://zeromq.zap.01"
-    async $ forever $ do
-      ("1.0":requestId:domain:ipAddress:identity:mech:xs) <- P.receiveMulti req
-      case mech of
-        "PLAIN" -> case xs of
-           (pass':user':_)
-             | user == user' && pass == pass' -> do
-                P.sendMulti req $ "1.0" :| [requestId, "200", "OK", "", ""]
-             | otherwise -> P.sendMulti req $ "1.0" :| [requestId, "400", "Credentials are not implemented", "", ""]
-           _ -> P.sendMulti req $ "1.0" :| [requestId, "500", "Method not implemented", "", ""]
-        _ -> P.sendMulti req $ "1.0" :| [requestId, "500", "Method not implemented", "", ""]
-
-
--- | Close socket immideately.
-closeZeroLinger :: P.Socket a -> IO ()
-closeZeroLinger sock = do
-  P.setLinger (P.restrict 0) sock
-  P.close sock
diff --git a/tests/TestAPI.hs b/tests/TestAPI.hs
--- a/tests/TestAPI.hs
+++ b/tests/TestAPI.hs
@@ -1,18 +1,30 @@
 {-# LANGUAGE OverloadedStrings #-}
 module Main where
 
-import Control.Concurrent 
-import Control.Exception
-import Control.Monad ( (<=<), replicateM )
+import Control.Concurrent
+import Control.Monad ( replicateM )
 
+import Data.IORef
+
 import Network.Transport
 import Network.Transport.ZMQ
-import System.Exit
+import Test.Tasty
+import Test.Tasty.HUnit
 
 main :: IO ()
-main = finish <=< trySome $ do
-    putStr "simple: "
-    Right (zmq, transport) <- createTransportEx defaultZMQParameters "127.0.0.1"
+main = defaultMain $
+  testGroup "API tests"
+      [ testCase "simple" test_simple
+      , testCase "connection break" test_connectionBreak
+      , testCase "test multicast" test_multicast
+      , testCase "authentification" test_auth
+      , testCase "connect to non existent host" test_nonexists
+      , testCase "test cleanup actions" test_cleanup
+      ]
+
+test_simple :: IO ()
+test_simple = do
+    Right transport <- createTransport defaultZMQParameters "127.0.0.1"
     Right ep1 <- newEndPoint transport
     Right ep2 <- newEndPoint transport
     Right c1  <- connect ep1 (address ep2) ReliableOrdered defaultConnectHints
@@ -23,18 +35,23 @@
     close c2
     [ConnectionOpened 1 ReliableOrdered _, Received 1 ["321"], ConnectionClosed 1] <- replicateM 3 $ receive ep1
     [ConnectionOpened 1 ReliableOrdered _, Received 1 ["123"], ConnectionClosed 1] <- replicateM 3 $ receive ep2
-    putStrLn "OK"
+    closeTransport transport
 
-    putStr "connection break: "
+test_connectionBreak :: IO ()
+test_connectionBreak = do
+    Right (zmq, transport) <-
+      createTransportExposeInternals defaultZMQParameters "127.0.0.1"
+    Right ep1 <- newEndPoint transport
+    Right ep2 <- newEndPoint transport
     Right ep3 <- newEndPoint transport
 
     Right c21  <- connect ep1 (address ep2) ReliableOrdered defaultConnectHints
     Right c22  <- connect ep2 (address ep1) ReliableOrdered defaultConnectHints
     Right c23  <- connect ep3 (address ep1) ReliableOrdered defaultConnectHints
 
-    ConnectionOpened 2 ReliableOrdered _ <- receive ep2
+    ConnectionOpened 1 ReliableOrdered _ <- receive ep2
+    ConnectionOpened 1 ReliableOrdered _ <- receive ep1
     ConnectionOpened 2 ReliableOrdered _ <- receive ep1
-    ConnectionOpened 3 ReliableOrdered _ <- receive ep1
 
     breakConnection zmq (address ep1) (address ep2)
 
@@ -48,50 +65,59 @@
     ErrorEvent (TransportError (EventConnectionLost _) _ ) <- receive ep1
     Right c24 <- connect ep3 (address ep1) ReliableOrdered defaultConnectHints
     Right ()  <- send c24 ["final"]
-    ConnectionOpened 4 ReliableOrdered _ <- receive ep1
-    Received 4 ["final"] <- receive ep1
-    putStrLn "OK"
-    multicast transport
-  where
-      multicast transport = do 
-        Right ep1 <- newEndPoint transport
-        Right ep2 <- newEndPoint transport
-        putStr "multicast: "
-        Right g1 <- newMulticastGroup ep1
-        multicastSubscribe g1
-        threadDelay 1000000
-        multicastSend g1 ["test"]
-        ReceivedMulticast _ ["test"] <- receive ep1
-        Right g2 <- resolveMulticastGroup ep2 (multicastAddress g1)
-        multicastSubscribe g2
-        threadDelay 100000
-        multicastSend g2 ["test-2"]
-        ReceivedMulticast _ ["test-2"] <- receive ep2
-        ReceivedMulticast _ ["test-2"] <- receive ep1
-        putStrLn "OK"
-        
-        putStr "Auth:"
-        Right tr2 <- createTransport
-                       defaultZMQParameters {authorizationType=ZMQAuthPlain "user" "password"}
-                       "127.0.0.1"
-        Right ep3 <- newEndPoint tr2
-        Right ep4 <- newEndPoint tr2
-        Right c3  <- connect ep3 (address ep4) ReliableOrdered defaultConnectHints
-        Right _   <- send c3 ["4456"]
-        [ConnectionOpened 1 ReliableOrdered _, Received 1 ["4456"]] <- replicateM 2 $ receive ep4
-        Right c4  <- connect ep3 (address ep4) ReliableOrdered defaultConnectHints
-        Right _   <- send c4 ["5567"]
-        [ConnectionOpened 2 ReliableOrdered _, Received 2 ["5567"]] <- replicateM 2 $ receive ep4
-        putStrLn "OK"
+    ConnectionOpened 3 ReliableOrdered _ <- receive ep1
+    Received 3 ["final"] <- receive ep1
+    closeTransport transport
 
-        putStr "Connect to non existing host: "
-        Left (TransportError ConnectFailed _) <- connect ep3 (EndPointAddress "tcp://128.0.0.1:7689") ReliableOrdered defaultConnectHints
-        putStrLn "OK"
-      finish (Left e) = print e >> exitWith (ExitFailure 1)
-      finish Right{} = exitWith ExitSuccess
+test_multicast :: IO ()
+test_multicast = do
+    Right transport <- createTransport defaultZMQParameters "127.0.0.1"
+    Right ep1 <- newEndPoint transport
+    Right ep2 <- newEndPoint transport
+    Right g1 <- newMulticastGroup ep1
+    multicastSubscribe g1
+    threadDelay 1000000
+    multicastSend g1 ["test"]
+    ReceivedMulticast _ ["test"] <- receive ep1
+    Right g2 <- resolveMulticastGroup ep2 (multicastAddress g1)
+    multicastSubscribe g2
+    threadDelay 100000
+    multicastSend g2 ["test-2"]
+    ReceivedMulticast _ ["test-2"] <- receive ep2
+    ReceivedMulticast _ ["test-2"] <- receive ep1
+    return ()
 
-trySome :: IO a -> IO (Either SomeException a)
-trySome = try
+test_auth :: IO ()
+test_auth = do
+    Right tr2 <-
+      createTransport defaultZMQParameters{ zmqSecurityMechanism =
+                                              Just $ SecurityPlain "user" "password" }
+                      "127.0.0.1"
+    Right ep3 <- newEndPoint tr2
+    Right ep4 <- newEndPoint tr2
+    Right c3  <- connect ep3 (address ep4) ReliableOrdered defaultConnectHints
+    Right _   <- send c3 ["4456"]
+    [ConnectionOpened 1 ReliableOrdered _, Received 1 ["4456"]] <- replicateM 2 $ receive ep4
+    Right c4  <- connect ep3 (address ep4) ReliableOrdered defaultConnectHints
+    Right _   <- send c4 ["5567"]
+    [ConnectionOpened 2 ReliableOrdered _, Received 2 ["5567"]] <- replicateM 2 $ receive ep4
+    return ()
 
-forkTry :: IO () -> IO ThreadId
-forkTry = forkIO
+test_nonexists :: IO ()
+test_nonexists = do
+    Right tr <- createTransport defaultZMQParameters "127.0.0.1"
+    Right ep <- newEndPoint tr
+    Left (TransportError ConnectFailed _) <- connect ep (EndPointAddress "tcp://129.0.0.1:7684") ReliableOrdered defaultConnectHints
+    closeTransport tr
+
+test_cleanup :: IO ()
+test_cleanup = do
+    Right (zmq, transport) <-
+      createTransportExposeInternals defaultZMQParameters "127.0.0.1"
+    x <- newIORef (0::Int)
+    Just _ <- registerCleanupAction zmq (modifyIORef x (+1))
+    Just u <- registerCleanupAction zmq (modifyIORef x (+1))
+    applyCleanupAction zmq u
+    closeTransport transport
+    2 <- readIORef x
+    return ()
diff --git a/tests/TestZMQ.hs b/tests/TestZMQ.hs
--- a/tests/TestZMQ.hs
+++ b/tests/TestZMQ.hs
@@ -32,7 +32,7 @@
     , ("CloseTransport",        testCloseTransport newTransport)
     , ("ExceptionOnReceive",    testExceptionOnReceive newTransport)
     , ("SendException",         testSendException newTransport)
-    , ("Kill",                  testKill newTransport 1000)
+    , ("Kill",                  testKill newTransport 100)
     ]
   where
-    numPings = 10000 :: Int
+    numPings = 500 :: Int
diff --git a/tests/test-ch.hs b/tests/test-ch.hs
--- a/tests/test-ch.hs
+++ b/tests/test-ch.hs
@@ -16,13 +16,6 @@
   , defaultZMQParameters
   , breakConnection
   )
-import Control.Concurrent.MVar
-import Control.Concurrent.STM
-import Control.Concurrent.STM.TMChan
-import qualified Data.Map as Map
-import qualified Data.ByteString.Char8 as B
-import Network.Transport
-import Network.Transport.ZMQ.Types
 import Test.Framework (defaultMain)
 
 main :: IO ()