Route tracking in namespaces and helper functions

3 files changed, 97 insertions, 32 deletions

diff --git a/src/Network.hs b/src/Network.hs
index 50dc7a3..d892404 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -146,14 +146,10 @@ newSubnet net vname = atomicallyWithIO $ do
     linkUp    vethSub                 -- sometimes gets stuck with NO-CARRIER for a while.
     linkUp    vethNet
 
-    postpone $ do
-        -- If the new subnet can be split further, routing rule for the whole prefix is needed
-        when (allowsSubnets (netPrefix sub)) $ callOn net $ "ip route add "
-            <> textIpNetwork (netPrefix sub)
-            <> " via " <> textIpAddress bridge
-            <> " dev " <> linkName vethNet
-            <> " src " <> textIpAddress router
-        callOn sub $ "ip route add default via " <> textIpAddress router <> " dev br0 src " <> textIpAddress bridge
+    -- If the new subnet can be split further, routing rule for the whole prefix is needed
+    when (allowsSubnets (netPrefix sub)) $ do
+        addRoute (netPrefix sub) bridge vethNet router
+    addRoute (IpPrefix []) router (netBridge sub) bridge
     return sub { netUpstream = Just vethSub }
 
 newNetwork :: IpPrefix -> FilePath -> WriterT [IO ()] STM Network
@@ -202,12 +198,6 @@ newNode nodeNetwork vname = atomicallyWithIO $ do
     addAddress nodeUpstream $ nodeIp
     linkUp $ nodeUpstream
     linkUp $ loopback node
-    postpone $ callOn node $ "ip route add default via " <> textIpAddress (lanIp 1) <> " dev veth0 src " <> textIpAddress nodeIp
+    addRoute (IpPrefix []) (lanIp 1) nodeUpstream nodeIp
 
     return node
-
-atomicallyWithIO :: MonadIO m => WriterT [IO ()] STM a -> m a
-atomicallyWithIO act = liftIO $ do
-    (x, fin) <- atomically $ runWriterT act
-    sequence_ fin
-    return x
diff --git a/src/Network/Ip.hs b/src/Network/Ip.hs
index e3d95cb..ae8a2d2 100644
--- a/src/Network/Ip.hs
+++ b/src/Network/Ip.hs
@@ -11,6 +11,8 @@ module Network.Ip (
     lanSubnet,
 
     MonadPIO(..),
+    MonadSTM(..),
+    atomicallyWithIO,
 
     NetworkNamespace,
     HasNetns(..),
@@ -25,12 +27,18 @@ module Network.Ip (
     addAddress,
     setMaster,
     linkUp, linkDown,
+
+    Route(..),
+    addRoute,
 ) where
 
+import Control.Concurrent.STM
 import Control.Monad.Writer
 
+import Data.Function
 import Data.Text (Text)
 import Data.Text qualified as T
+import Data.Typeable
 import Data.Word
 
 import System.Process
@@ -74,35 +82,71 @@ instance Monad m => MonadPIO (WriterT [IO ()] m) where
     postpone = tell . (:[])
 
 
-newtype NetworkNamespace = NetworkNamespace
+class Monad m => MonadSTM m where
+    liftSTM :: STM a -> m a
+
+instance MonadSTM STM where
+    liftSTM = id
+
+instance MonadSTM m => MonadSTM (WriterT [IO ()] m) where
+    liftSTM = lift . liftSTM
+
+
+atomicallyWithIO :: MonadIO m => WriterT [IO ()] STM a -> m a
+atomicallyWithIO act = liftIO $ do
+    (x, fin) <- atomically $ runWriterT act
+    sequence_ fin
+    return x
+
+
+data NetworkNamespace = NetworkNamespace
     { netnsName :: Text
+    , netnsRoutesConfigured :: TVar [Route]
+    , netnsRoutesActive :: TVar [Route]
     }
-    deriving (Eq, Ord)
+
+instance Eq NetworkNamespace where
+    (==) = (==) `on` netnsName
+
+instance Ord NetworkNamespace where
+    compare = compare `on` netnsName
 
 class HasNetns a where getNetns :: a -> NetworkNamespace
 instance HasNetns NetworkNamespace where getNetns = id
 
-addNetworkNamespace :: MonadPIO m => Text -> m NetworkNamespace
-addNetworkNamespace name = do
-    postpone $ callCommand $ T.unpack $ "ip netns add \"" <> name <> "\""
-    return $ NetworkNamespace
-        { netnsName = name
-        }
+addNetworkNamespace :: (MonadPIO m, MonadSTM m) => Text -> m NetworkNamespace
+addNetworkNamespace netnsName = do
+    postpone $ callCommand $ T.unpack $ "ip netns add \"" <> netnsName <> "\""
+    netnsRoutesConfigured <- liftSTM $ newTVar []
+    netnsRoutesActive <- liftSTM $ newTVar []
+    return $ NetworkNamespace {..}
 
 textNetnsName :: NetworkNamespace -> Text
 textNetnsName = netnsName
 
 callOn :: HasNetns a => a -> Text -> IO ()
 callOn n cmd = callCommand $ T.unpack $ "ip netns exec \"" <> ns <> "\" " <> cmd
-    where NetworkNamespace ns = getNetns n
+    where ns = textNetnsName $ getNetns n
 
 
 data Link a = Link
     { linkName :: Text
     , linkNetns :: NetworkNamespace
     }
+    deriving (Eq)
+
+data SomeLink = forall a. Typeable a => SomeLink (Link a)
+
+instance Eq SomeLink where
+    SomeLink a == SomeLink b
+        | Just b' <- cast b = a == b'
+        | otherwise = False
+
+liftSomeLink :: (forall a. Link a -> b) -> SomeLink -> b
+liftSomeLink f (SomeLink x) = f x
 
 instance HasNetns (Link a) where getNetns = linkNetns
+instance HasNetns SomeLink where getNetns = liftSomeLink linkNetns
 
 data Loopback
 
@@ -135,10 +179,41 @@ setMaster link bridge = postpone $ do
     when (getNetns link /= getNetns bridge) $ fail "link and bridge in different network namespaces"
     callOn link $ "ip link set dev \"" <> linkName link <> "\" master \"" <> linkName bridge <> "\""
 
-linkUp :: MonadPIO m => Link a -> m ()
+linkUp :: (Typeable a, MonadPIO m, MonadSTM m) => Link a -> m ()
 linkUp link = do
-    postpone $ callOn link $ "ip link set dev \"" <> linkName link <> "\" up"
-
-linkDown :: MonadPIO m => Link a -> m ()
+    routes <- liftSTM $ filter ((== SomeLink link) . routeDev) <$> readTVar (netnsRoutesConfigured (getNetns link))
+    liftSTM $ modifyTVar (netnsRoutesActive (getNetns link)) $ (routes ++)
+    postpone $ do
+        callOn link $ "ip link set dev \"" <> linkName link <> "\" up"
+        -- add back routes that were automatically removed by kernel when the link went down
+        mapM_ applyRoute routes
+
+linkDown :: (Typeable a, MonadPIO m, MonadSTM m) => Link a -> m ()
 linkDown link = do
+    -- routes using this device will be automatically removed by kernel
+    liftSTM $ modifyTVar (netnsRoutesActive (getNetns link)) $ filter ((/= SomeLink link) . routeDev)
     postpone $ callOn link $ "ip link set dev \"" <> linkName link <> "\" down"
+
+
+data Route = Route
+    { routePrefix :: IpPrefix
+    , routeVia :: IpAddress
+    , routeDev :: SomeLink
+    , routeSrc :: IpAddress
+    }
+
+addRoute :: Typeable a => IpPrefix -> IpAddress -> Link a -> IpAddress -> WriterT [IO ()] STM ()
+addRoute routePrefix routeVia link routeSrc = do
+    let routeDev = SomeLink link
+        route = Route {..}
+    lift $ do
+        modifyTVar (netnsRoutesConfigured (getNetns link)) (route:)
+        modifyTVar (netnsRoutesActive (getNetns link)) (route:)
+    postpone $ applyRoute route
+
+applyRoute :: Route -> IO ()
+applyRoute route = callOn (routeDev route) $ "ip route add "
+    <> textIpNetwork (routePrefix route)
+    <> " via " <> textIpAddress (routeVia route)
+    <> " dev " <> linkName `liftSomeLink` (routeDev route)
+    <> " src " <> textIpAddress (routeSrc route)
diff --git a/src/Run.hs b/src/Run.hs
index 1a0b981..58df868 100644
--- a/src/Run.hs
+++ b/src/Run.hs
@@ -215,9 +215,9 @@ withDisconnectedUp link inner = do
       then inner
       else do
         local (fmap $ \s -> s { tsDisconnectedUp = S.insert netns $ tsDisconnectedUp s }) $ do
-            linkDown link
+            atomicallyWithIO $ linkDown link
             x <- inner
-            linkUp link
+            atomicallyWithIO $ linkUp link
             return x
 
 withDisconnectedBridge :: Link Bridge -> TestRun a -> TestRun a
@@ -228,9 +228,9 @@ withDisconnectedBridge bridge inner = do
       then inner
       else do
         local (fmap $ \s -> s { tsDisconnectedBridge = S.insert netns $ tsDisconnectedBridge s }) $ do
-            linkDown bridge
+            atomicallyWithIO $ linkDown bridge
             x <- inner
-            linkUp bridge
+            atomicallyWithIO $ linkUp bridge
             return x
 
 withNodePacketLoss :: Node -> Scientific -> TestRun a -> TestRun a