From 492568c72273ee2e7f98bc520e46fa01c2959851 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Roman=20Smr=C5=BE?= <roman.smrz@seznam.cz>
Date: Wed, 26 Apr 2023 22:07:06 +0200
Subject: Route tracking in namespaces and helper functions

---
 src/Network/Ip.hs | 101 +++++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 88 insertions(+), 13 deletions(-)

(limited to 'src/Network')

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)
-- 
cgit v1.2.3