1 files changed, 14 insertions, 6 deletions

diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index d94f028f..b775e7d3 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -116,6 +116,7 @@ import Data.Text as Text
 import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
 import Data.Serialize as S
 import Network.DHT.Types
+import Network.DatagramServer.Types
 
 
 import Data.Torrent as Torrent
@@ -343,10 +344,12 @@ locFromCS cs = case getCallStack cs of
 -- | Run DHT session. You /must/ properly close session using
 -- 'closeNode' function, otherwise socket or other scarce resources may
 -- leak.
-newNode :: ( Address ip
+newNode :: forall raw dht ip u.
+           ( Address ip
            , FiniteBits (NodeId dht)
            , Serialize (NodeId dht)
            , Kademlia dht
+           , WireFormat raw dht
            )
           => -- [NodeHandler] -- ^ handlers to run on accepted queries;
              Options       -- ^ various dht options;
@@ -363,13 +366,18 @@ newNode opts naddr logger mbid = do
     nodeAddr = toSockAddr naddr
     initNode = do
       s      <- getInternalState
-      (_, m) <- allocate (newManager rpcOpts (logt logger) nodeAddr []) closeManager
+      (myId, infovar, getst) <- liftIO $ do
+        (i, ctx) <- initializeServerState (Proxy :: Proxy (dht raw)) mbid
+        var <- atomically (newTVar Nothing)
+        let getst dest = do
+              info <- atomically . readTVar $ var
+              return ( maybe i myNodeId info, ctx)
+        return (i, var, getst)
+      (_, m) <- allocate (newManager rpcOpts (logt logger) nodeAddr getst []) closeManager
       liftIO $ do
         dta <- initializeDHTData
-        myId   <- maybe genNodeId return mbid
-        node   <- Node opts myId s m
-                <$> atomically (newTVar Nothing)
-                <*> newTVarIO S.empty
+        node   <- Node opts myId s m infovar
+                <$> newTVarIO S.empty
                 <*> pure dta
                 <*> pure logger
         return node