Refactor discovery state to custom data types

1 files changed, 81 insertions, 65 deletions

diff --git a/src/Erebos/Discovery.hs b/src/Erebos/Discovery.hs
index f156c85..63400cb 100644
--- a/src/Erebos/Discovery.hs
+++ b/src/Erebos/Discovery.hs
@@ -3,7 +3,7 @@
 module Erebos.Discovery (
     DiscoveryService(..),
     DiscoveryAttributes(..),
-    DiscoveryConnection(..)
+    DiscoveryConnection(..),
 ) where
 
 import Control.Concurrent
@@ -90,12 +90,13 @@ instance Storable DiscoveryService where
             DiscoveryConnectionRequest conn -> storeConnection "request" conn
             DiscoveryConnectionResponse conn -> storeConnection "response" conn
 
-        where storeConnection ctype conn = do
-                  storeText "connection" $ ctype
-                  storeRawRef "source" $ dconnSource conn
-                  storeRawRef "target" $ dconnTarget conn
-                  storeMbText "address" $ dconnAddress conn
-                  storeMbRef "ice-info" $ dconnIceInfo conn
+      where
+        storeConnection ctype DiscoveryConnection {..} = do
+            storeText "connection" $ ctype
+            storeRawRef "source" dconnSource
+            storeRawRef "target" dconnTarget
+            storeMbText "address" dconnAddress
+            storeMbRef "ice-info" dconnIceInfo
 
     load' = loadRec $ msum
             [ do
@@ -120,14 +121,15 @@ instance Storable DiscoveryService where
             , loadConnection "request" DiscoveryConnectionRequest
             , loadConnection "response" DiscoveryConnectionResponse
             ]
-        where loadConnection ctype ctor = do
-                  ctype' <- loadText "connection"
-                  guard $ ctype == ctype'
-                  return . ctor =<< DiscoveryConnection
-                      <$> loadRawRef "source"
-                      <*> loadRawRef "target"
-                      <*> loadMbText "address"
-                      <*> loadMbRef "ice-info"
+      where
+        loadConnection ctype ctor = do
+            ctype' <- loadText "connection"
+            guard $ ctype == ctype'
+            dconnSource <- loadRawRef "source"
+            dconnTarget <- loadRawRef "target"
+            dconnAddress <- loadMbText "address"
+            dconnIceInfo <- loadMbRef "ice-info"
+            return $ ctor DiscoveryConnection {..}
 
 data DiscoveryPeer = DiscoveryPeer
     { dpPriority :: Int
@@ -135,22 +137,38 @@ data DiscoveryPeer = DiscoveryPeer
     , dpAddress :: [ Text ]
 #ifdef ENABLE_ICE_SUPPORT
     , dpIceSession :: Maybe IceSession
+#else
+    , dpIceSession :: Maybe ()
 #endif
     }
 
+data DiscoveryPeerState = DiscoveryPeerState
+#ifdef ENABLE_ICE_SUPPORT
+    { dpsIceConfig :: Maybe IceConfig
+#else
+    { dpsIceConfig :: Maybe ()
+#endif
+    }
+
+data DiscoveryGlobalState = DiscoveryGlobalState
+    { dgsPeers :: Map RefDigest DiscoveryPeer
+    }
+
 instance Service DiscoveryService where
     serviceID _ = mkServiceID "dd59c89c-69cc-4703-b75b-4ddcd4b3c23c"
 
     type ServiceAttributes DiscoveryService = DiscoveryAttributes
     defaultServiceAttributes _ = defaultDiscoveryAttributes
 
-#ifdef ENABLE_ICE_SUPPORT
-    type ServiceState DiscoveryService = Maybe IceConfig
-    emptyServiceState _ = Nothing
-#endif
+    type ServiceState DiscoveryService = DiscoveryPeerState
+    emptyServiceState _ = DiscoveryPeerState
+        { dpsIceConfig = Nothing
+        }
 
-    type ServiceGlobalState DiscoveryService = Map RefDigest DiscoveryPeer
-    emptyServiceGlobalState _ = M.empty
+    type ServiceGlobalState DiscoveryService = DiscoveryGlobalState
+    emptyServiceGlobalState _ = DiscoveryGlobalState
+        { dgsPeers = M.empty
+        }
 
     serviceHandler msg = case fromStored msg of
         DiscoverySelf addrs priority -> do
@@ -171,15 +189,14 @@ instance Service DiscoveryService where
 
                 | otherwise -> return Nothing
 
-            forM_ (idDataF =<< unfoldOwners pid) $ \s ->
-                svcModifyGlobal $ M.insertWith insertHelper (refDigest $ storedRef s) DiscoveryPeer
-                    { dpPriority = fromMaybe 0 priority
-                    , dpPeer = Just peer
-                    , dpAddress = addrs
-#ifdef ENABLE_ICE_SUPPORT
-                    , dpIceSession = Nothing
-#endif
-                    }
+            forM_ (idDataF =<< unfoldOwners pid) $ \sdata -> do
+                let dp = DiscoveryPeer
+                        { dpPriority = fromMaybe 0 priority
+                        , dpPeer = Just peer
+                        , dpAddress = addrs
+                        , dpIceSession = Nothing
+                        }
+                svcModifyGlobal $ \s -> s { dgsPeers = M.insertWith insertHelper (refDigest $ storedRef sdata) dp $ dgsPeers s }
             attrs <- asks svcAttributes
             replyPacket $ DiscoveryAcknowledged matchedAddrs
                 (discoveryStunServer attrs)
@@ -207,22 +224,23 @@ instance Service DiscoveryService where
             cfg <- liftIO $ iceCreateConfig
                 (toIceServer stunServer stunPort)
                 (toIceServer turnServer turnPort)
-            svcSet cfg
+            svcModify $ \s -> s { dpsIceConfig = cfg }
 #endif
             return ()
 
         DiscoverySearch ref -> do
-            dpeer <- M.lookup (refDigest ref) <$> svcGetGlobal
+            dpeer <- M.lookup (refDigest ref) . dgsPeers <$> svcGetGlobal
             replyPacket $ DiscoveryResult ref $ maybe [] dpAddress dpeer
 
         DiscoveryResult ref [] -> do
             svcPrint $ "Discovery: " ++ show (refDigest ref) ++ " not found"
 
         DiscoveryResult ref addrs -> do
+            let dgst = refDigest ref
             -- TODO: check if we really requested that
             server <- asks svcServer
             self <- svcSelf
-            mbIceConfig <- svcGet
+            mbIceConfig <- dpsIceConfig <$> svcGet
             discoveryPeer <- asks svcPeer
             let runAsService = runPeerService @DiscoveryService discoveryPeer
 
@@ -240,12 +258,13 @@ instance Service DiscoveryService where
                             Left err -> putStrLn $ "Discovery: failed to send connection request: " ++ err
 
                     runAsService $ do
-                        svcModifyGlobal $ M.insert (refDigest ref) DiscoveryPeer
-                            { dpPriority = 0
-                            , dpPeer = Nothing
-                            , dpAddress = []
-                            , dpIceSession = Just ice
-                            }
+                        let dp = DiscoveryPeer
+                                { dpPriority = 0
+                                , dpPeer = Nothing
+                                , dpAddress = []
+                                , dpIceSession = Just ice
+                                }
+                        svcModifyGlobal $ \s -> s { dgsPeers = M.insert dgst dp $ dgsPeers s }
 #else
                 -> do
                     return ()
@@ -256,14 +275,13 @@ instance Service DiscoveryService where
                         getAddrInfo (Just $ defaultHints { addrSocketType = Datagram }) (Just ipaddr) (Just port)
                     peer <- serverPeer server (addrAddress saddr)
                     runAsService $ do
-                        svcModifyGlobal $ M.insert (refDigest ref) DiscoveryPeer
-                            { dpPriority = 0
-                            , dpPeer = Just peer
-                            , dpAddress = []
-#ifdef ENABLE_ICE_SUPPORT
-                            , dpIceSession = Nothing
-#endif
-                        }
+                        let dp = DiscoveryPeer
+                                { dpPriority = 0
+                                , dpPeer = Just peer
+                                , dpAddress = []
+                                , dpIceSession = Nothing
+                                }
+                        svcModifyGlobal $ \s -> s { dgsPeers = M.insert dgst dp $ dgsPeers s }
 
                 | otherwise -> do
                     runAsService $ do
@@ -273,31 +291,30 @@ instance Service DiscoveryService where
             self <- svcSelf
             let rconn = emptyConnection (dconnSource conn) (dconnTarget conn)
             if refDigest (dconnTarget conn) `elem` (map (refDigest . storedRef) $ idDataF =<< unfoldOwners self)
-               then do
+              then if
 #ifdef ENABLE_ICE_SUPPORT
-                    -- request for us, create ICE sesssion
+                -- request for us, create ICE sesssion
+                | Just prinfo <- dconnIceInfo conn -> do
                     server <- asks svcServer
                     peer <- asks svcPeer
-                    svcGet >>= \case
+                    dpsIceConfig <$> svcGet >>= \case
                         Just config -> do
                             liftIO $ void $ iceCreateSession config PjIceSessRoleControlled $ \ice -> do
                                 rinfo <- iceRemoteInfo ice
                                 res <- runExceptT $ sendToPeer peer $ DiscoveryConnectionResponse rconn { dconnIceInfo = Just rinfo }
                                 case res of
-                                    Right _ -> do
-                                        case dconnIceInfo conn of
-                                            Just prinfo -> iceConnect ice prinfo $ void $ serverPeerIce server ice
-                                            Nothing -> putStrLn $ "Discovery: connection request without ICE remote info"
+                                    Right _ -> iceConnect ice prinfo $ void $ serverPeerIce server ice
                                     Left err -> putStrLn $ "Discovery: failed to send connection response: " ++ err
                         Nothing -> do
                             svcPrint $ "Discovery: ICE request from peer without ICE configuration"
-#else
-                    return ()
 #endif
 
-               else do
+                | otherwise -> do
+                    svcPrint $ "Discovery: unsupported connection request"
+
+              else do
                     -- request to some of our peers, relay
-                    mbdp <- M.lookup (refDigest $ dconnTarget conn) <$> svcGetGlobal
+                    mbdp <- M.lookup (refDigest $ dconnTarget conn) . dgsPeers <$> svcGetGlobal
                     case mbdp of
                         Nothing -> replyPacket $ DiscoveryConnectionResponse rconn
                         Just dp
@@ -307,29 +324,28 @@ instance Service DiscoveryService where
 
         DiscoveryConnectionResponse conn -> do
             self <- svcSelf
-            dpeers <- svcGetGlobal
+            dpeers <- dgsPeers <$> svcGetGlobal
             if refDigest (dconnSource conn) `elem` (map (refDigest . storedRef) $ idDataF =<< unfoldOwners self)
                then do
                     -- response to our request, try to connect to the peer
-#ifdef ENABLE_ICE_SUPPORT
                     server <- asks svcServer
                     if  | Just addr <- dconnAddress conn
                         , [ipaddr, port] <- words (T.unpack addr) -> do
                             saddr <- liftIO $ head <$>
                                 getAddrInfo (Just $ defaultHints { addrSocketType = Datagram }) (Just ipaddr) (Just port)
                             peer <- liftIO $ serverPeer server (addrAddress saddr)
-                            svcModifyGlobal $ M.insert (refDigest $ dconnTarget conn) $
-                                DiscoveryPeer 0 (Just peer) [] Nothing
+                            svcModifyGlobal $ \s -> s
+                                { dgsPeers = M.insert (refDigest $ dconnTarget conn)
+                                    (DiscoveryPeer 0 (Just peer) [] Nothing) $ dgsPeers s }
 
+#ifdef ENABLE_ICE_SUPPORT
                         | Just dp <- M.lookup (refDigest $ dconnTarget conn) dpeers
                         , Just ice <- dpIceSession dp
                         , Just rinfo <- dconnIceInfo conn -> do
                             liftIO $ iceConnect ice rinfo $ void $ serverPeerIce server ice
+#endif
 
                         | otherwise -> svcPrint $ "Discovery: connection request failed"
-#else
-                    return ()
-#endif
                else do
                     -- response to relayed request
                     case M.lookup (refDigest $ dconnSource conn) dpeers of