1 files changed, 32 insertions, 13 deletions
diff --git a/src/Erebos/Network/Protocol.hs b/src/Erebos/Network/Protocol.hs
index c340503..025f52c 100644
--- a/src/Erebos/Network/Protocol.hs
+++ b/src/Erebos/Network/Protocol.hs
@@ -3,6 +3,7 @@ module Erebos.Network.Protocol (
     transportToObject,
     TransportHeader(..),
     TransportHeaderItem(..),
+    ServiceID(..),
     SecurityRequirement(..),
 
     WaitingRef(..),
@@ -22,7 +23,8 @@ module Erebos.Network.Protocol (
     connSetChannel,
     connClose,
 
-    RawStreamReader, RawStreamWriter,
+    RawStreamReader(..), RawStreamWriter(..),
+    StreamPacket(..),
     connAddWriteStream,
     connAddReadStream,
     readStreamToList,
@@ -36,6 +38,7 @@ import Control.Applicative
 import Control.Concurrent
 import Control.Concurrent.Async
 import Control.Concurrent.STM
+import Control.Exception
 import Control.Monad
 import Control.Monad.Except
 import Control.Monad.Trans
@@ -68,9 +71,9 @@ import Erebos.Flow
 import Erebos.Identity
 import Erebos.Network.Channel
 import Erebos.Object
-import Erebos.Service
 import Erebos.Storable
 import Erebos.Storage
+import Erebos.UUID (UUID)
 
 
 protocolVersion :: Text
@@ -107,6 +110,9 @@ data TransportHeaderItem
     | StreamOpen Word8
     deriving (Eq, Show)
 
+newtype ServiceID = ServiceID UUID
+    deriving (Eq, Ord, Show, StorableUUID)
+
 newtype Cookie = Cookie ByteString
     deriving (Eq, Show)
 
@@ -283,7 +289,11 @@ connAddWriteStream conn@Connection {..} = do
     runExceptT $ do
         ((streamNumber, stream), outStreams') <- doInsert 1 outStreams
         lift $ writeTVar cOutStreams outStreams'
-        return (StreamOpen streamNumber, sFlowIn stream, go cGlobalState streamNumber stream)
+        return
+            ( StreamOpen streamNumber
+            , RawStreamWriter (fromIntegral streamNumber) (sFlowIn stream)
+            , go cGlobalState streamNumber stream
+            )
 
   where
     go gs@GlobalState {..} streamNumber stream = do
@@ -356,14 +366,21 @@ connAddReadStream Connection {..} streamNumber = do
             sNextSequence <- newTVar 0
             sWaitingForAck <- newTVar 0
             let stream = Stream {..}
-            return (stream, (streamNumber, stream) : streams)
-    (stream, inStreams') <- doInsert inStreams
+            return ( streamNumber, stream, (streamNumber, stream) : streams )
+    ( num, stream, inStreams' ) <- doInsert inStreams
     writeTVar cInStreams inStreams'
-    return $ sFlowOut stream
+    return $ RawStreamReader (fromIntegral num) (sFlowOut stream)
+
 
+data RawStreamReader = RawStreamReader
+    { rsrNum :: Int
+    , rsrFlow :: Flow StreamPacket Void
+    }
 
-type RawStreamReader = Flow StreamPacket Void
-type RawStreamWriter = Flow Void StreamPacket
+data RawStreamWriter = RawStreamWriter
+    { rswNum :: Int
+    , rswFlow :: Flow Void StreamPacket
+    }
 
 data Stream = Stream
     { sState :: TVar StreamState
@@ -398,7 +415,7 @@ streamClosed Connection {..} snum = atomically $ do
     modifyTVar' cOutStreams $ filter ((snum /=) . fst)
 
 readStreamToList :: RawStreamReader -> IO (Word64, [(Word64, BC.ByteString)])
-readStreamToList stream = readFlowIO stream >>= \case
+readStreamToList stream = readFlowIO (rsrFlow stream) >>= \case
     StreamData sq bytes -> fmap ((sq, bytes) :) <$> readStreamToList stream
     StreamClosed sqEnd  -> return (sqEnd, [])
 
@@ -420,10 +437,10 @@ writeByteStringToStream :: RawStreamWriter -> BL.ByteString -> IO ()
 writeByteStringToStream stream = go 0
   where
     go seqNum bstr
-        | BL.null bstr = writeFlowIO stream $ StreamClosed seqNum
+        | BL.null bstr = writeFlowIO (rswFlow stream) $ StreamClosed seqNum
         | otherwise    = do
             let (cur, rest) = BL.splitAt 500 bstr -- TODO: MTU
-            writeFlowIO stream $ StreamData seqNum (BL.toStrict cur)
+            writeFlowIO (rswFlow stream) $ StreamData seqNum (BL.toStrict cur)
             go (seqNum + 1) rest
 
 
@@ -512,8 +529,10 @@ erebosNetworkProtocol initialIdentity gLog gDataFlow gControlFlow = do
 
             race_ (waitTill next) waitForUpdate
 
-    race_ signalTimeouts $ forever $ join $ atomically $
-        passUpIncoming gs <|> processIncoming gs <|> processOutgoing gs
+    race_ signalTimeouts $ forever $ do
+        io <- atomically $ do
+            passUpIncoming gs <|> processIncoming gs <|> processOutgoing gs
+        catch io $ \(e :: SomeException) -> atomically $ gLog $ "exception during network protocol handling: " <> show e
 
 
 getConnection :: GlobalState addr -> addr -> STM (Connection addr)