Network refactoring with explicit prefixes

3 files changed, 116 insertions, 54 deletions

diff --git a/src/Main.hs b/src/Main.hs
index adb738e..0d8a3fd 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -50,8 +50,8 @@ withNodePacketLoss node loss inner = do
   where
     resetLoss = do
         tl <- asks $ fromMaybe 0 . M.lookup (nodeName node) . tsNodePacketLoss . snd
-        liftIO $ callOn node $ "tc qdisc replace dev veth0 root netem loss " ++ show (tl * 100) ++ "%"
-        liftIO $ putStrLn $ "tc qdisc replace dev veth0 root netem loss " ++ show (tl * 100) ++ "%"
+        liftIO $ callOn node $ "tc qdisc replace dev veth0 root netem loss " <> T.pack (show (tl * 100)) <> "%"
+        liftIO $ putStrLn $ "tc qdisc replace dev veth0 root netem loss " <> show (tl * 100) <> "%"
 
 atomicallyTest :: STM a -> TestRun a
 atomicallyTest act = do
@@ -67,57 +67,26 @@ atomicallyTest act = do
 withNetwork :: (Network -> TestRun a) -> TestRun a
 withNetwork inner = do
     testDir <- asks $ optTestDir . teOptions . fst
-    net <- liftIO $ do
-        callCommand "ip link add name br0 group 1 type bridge"
-        callCommand "ip addr add 192.168.0.1/24 broadcast 192.168.0.255 dev br0"
-        callCommand "ip link set dev br0 up"
-        callCommand "ip link set dev lo up"
-        Network <$> newMVar [] <*> pure testDir
-
-    res <- withProcess (Left net) (ProcNameTcpdump) (Just softwareTermination)
-        ("tcpdump -i br0 -w '" ++ testDir ++ "/br0.pcap' -U -Z root") $ \_ -> do
-        local (fmap $ \s -> s { tsNetwork = net }) $ inner net
+    inet <- newInternet testDir
+    let net = inetRoot inet
+
+    tcpdump <- liftIO (findExecutable "tcpdump") >>= return . \case
+        Just path -> withProcess (Left net) ProcNameTcpdump (Just softwareTermination)
+            (path ++ " -i br0 -w '" ++ testDir ++ "/br0.pcap' -U -Z root") . const
+        Nothing -> id
 
-    liftIO $ do
-        callCommand $ "ip -all netns del"
-        callCommand $ "ip link del group 1"
+    res <- tcpdump $ do
+        local (fmap $ \s -> s { tsNetwork = net }) $ inner net
 
+    delInternet inet
     return res
 
-createNode :: Expr Network -> Either (TypedVarName Node) (TypedVarName Process) -> (Node -> TestRun a) -> TestRun a
-createNode netexpr tvname inner = do
+withNode :: Expr Network -> Either (TypedVarName Node) (TypedVarName Process) -> (Node -> TestRun a) -> TestRun a
+withNode netexpr tvname inner = do
     net <- eval netexpr
-    node <- liftIO $ do
-        node <- modifyMVar (netNodes net) $ \nodes -> do
-            let nname = nextNodeName (either fromTypedVarName fromTypedVarName tvname) $ map nodeName nodes
-                ip = "192.168.0." ++ show (11 + length nodes)
-                node = Node { nodeName = nname
-                            , nodeIp = T.pack ip
-                            , nodeNetwork = net
-                            , nodeDir = netDir net </> ("erebos_" ++ unpackNodeName nname)
-                            }
-            return $ (node : nodes, node)
-
-        let name = unpackNodeName $ nodeName node
-            dir = nodeDir node
-
-        exists <- doesPathExist dir
-        when exists $ ioError $ userError $ dir ++ " exists"
-        createDirectoryIfMissing True dir
-
-        callCommand $ "ip netns add \""++ name ++ "\""
-        callCommand $ "ip link add \"veth_" ++ name ++ "\" group 1 type veth peer name veth0 netns \"" ++ name ++ "\""
-        callCommand $ "ip link set dev \"veth_" ++ name ++ "\" master br0 up"
-        callOn node $ "ip addr add " ++ T.unpack (nodeIp node) ++ "/24 broadcast 192.168.0.255 dev veth0"
-        callOn node $ "ip link set dev veth0 up"
-        callOn node $ "ip link set dev lo up"
-        return node
-
+    node <- newNode net (either fromTypedVarName fromTypedVarName tvname)
     either (flip withVar node . fromTypedVarName) (const id) tvname $ inner node
 
-callOn :: Node -> String -> IO ()
-callOn node cmd = callCommand $ "ip netns exec \"" ++ unpackNodeName (nodeName node) ++ "\" " ++ cmd
-
 tryMatch :: Regex -> [Text] -> Maybe ((Text, [Text]), [Text])
 tryMatch re (x:xs) | Right (Just (_, _, _, capture)) <- regexMatch re x = Just ((x, capture), xs)
                    | otherwise = fmap (x:) <$> tryMatch re xs
@@ -188,13 +157,13 @@ evalSteps = mapM_ $ \case
             withVar name i $ evalSteps inner
 
     DeclNode name@(TypedVarName vname) net inner -> do
-        createNode net (Left name) $ \node -> do
+        withNode net (Left name) $ \node -> do
             withVar vname node $ evalSteps inner
 
     Spawn tvname@(TypedVarName vname@(VarName tname)) target inner -> do
         case target of
-            Left nname -> createNode RootNetwork (Left nname) go
-            Right (Left net) -> createNode net (Right tvname) go
+            Left nname -> withNode RootNetwork (Left nname) go
+            Right (Left net) -> withNode net (Right tvname) go
             Right (Right node) -> go =<< eval node
       where
         go node = do
diff --git a/src/Network.hs b/src/Network.hs
index 5b386c8..29621fc 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -1,23 +1,46 @@
 module Network (
+    Internet(..),
     Network(..),
     Node(..),
     NodeName(..), textNodeName, unpackNodeName,
     nextNodeName,
+
+    HasNetns(..),
+    callOn,
+
+    newInternet, delInternet,
+    newNode,
 ) where
 
 import Control.Arrow
-import Control.Concurrent
+import Control.Concurrent.STM
+import Control.Monad
+import Control.Monad.IO.Class
 
 import Data.Text (Text)
 import Data.Text qualified as T
+import Data.Word
+
+import System.Directory
+import System.FilePath
+import System.Process
 
 import Test
 
+data Internet = Internet
+    { inetDir :: FilePath
+    , inetRoot :: Network
+    }
+
 data Network = Network
-    { netNodes :: MVar [Node]
+    { netPrefix :: [Word8]
+    , netNodes :: TVar [Node]
     , netDir :: FilePath
     }
 
+textNetworkName :: Network -> Text
+textNetworkName n = T.intercalate "_" (map (T.pack . show) (netPrefix n))
+
 data Node = Node
     { nodeName :: NodeName
     , nodeIp :: Text
@@ -43,10 +66,18 @@ nextNodeName (VarName tname) = go 0
                                   | otherwise       = go n ns
 
 
+class HasNetns a where netnsName :: a -> Text
+instance HasNetns Network where netnsName n = "s" <> textNetworkName n
+instance HasNetns Node where netnsName n = netnsName (nodeNetwork n) <> ":" <> textNodeName (nodeName n)
+
+callOn :: HasNetns a => a -> Text -> IO ()
+callOn n cmd = callCommand $ T.unpack $ "ip netns exec \"" <> netnsName n <> "\" " <> cmd
+
+
 instance ExprType Network where
     textExprType _ = T.pack "network"
-    textExprValue _ = T.pack "s:0"
-    emptyVarValue = Network undefined undefined
+    textExprValue n = "s:" <> textNetworkName n
+    emptyVarValue = Network [] undefined undefined
 
 instance ExprType Node where
     textExprType _ = T.pack "node"
@@ -56,3 +87,64 @@ instance ExprType Node where
     recordMembers = map (first T.pack)
         [ ("ip", RecordSelector $ nodeIp)
         ]
+
+
+makeIpAddress :: [Word8] -> Word8 -> Text
+makeIpAddress prefix num = T.intercalate "." $ map (T.pack . show) $ prefix ++ replicate (3 - length prefix) 0 ++ [num]
+
+newInternet :: MonadIO m => FilePath -> m Internet
+newInternet dir = do
+    inet <- liftIO $ atomically $ do
+        Internet
+            <$> pure dir
+            <*> newNetwork [1] dir
+    initNetwork $ inetRoot inet
+    return inet
+
+delInternet :: MonadIO m => Internet -> m ()
+delInternet _ = liftIO $ do
+    callCommand $ "ip -all netns delete"
+
+newNetwork :: [Word8] -> FilePath -> STM Network
+newNetwork prefix dir = do
+    Network
+        <$> pure prefix
+        <*> newTVar []
+        <*> pure dir
+
+initNetwork :: MonadIO m => Network -> m ()
+initNetwork net = liftIO $ do
+    callCommand $ T.unpack $ "ip netns add \"" <> netnsName net <> "\""
+    callOn net $ "ip link add name br0 type bridge"
+    callOn net $ "ip addr add " <> makeIpAddress (netPrefix net) 1 <> " broadcast " <> makeIpAddress (netPrefix net) 255 <> " dev br0"
+    callOn net $ "ip link set dev br0 up"
+    callOn net $ "ip link set dev lo up"
+
+newNode :: MonadIO m => Network -> VarName -> m Node
+newNode net vname = liftIO $ do
+    node <- atomically $ do
+        nodes <- readTVar (netNodes net)
+        let nname = nextNodeName vname $ map nodeName nodes
+            node = Node { nodeName = nname
+                        , nodeIp = makeIpAddress (netPrefix net) (fromIntegral $ 2 + length nodes)
+                        , nodeNetwork = net
+                        , nodeDir = netDir net </> ("node_" ++ unpackNodeName nname)
+                        }
+        writeTVar (netNodes net) (node : nodes)
+        return node
+
+    let name = textNodeName $ nodeName node
+        dir = nodeDir node
+
+    exists <- doesPathExist dir
+    when exists $ ioError $ userError $ dir ++ " exists"
+    createDirectoryIfMissing True dir
+
+    callCommand $ T.unpack $ "ip netns add \"" <> netnsName node <> "\""
+    callOn net  $ "ip link add \"veth_" <> name <> "\" type veth peer name veth0 netns \"" <> netnsName node <> "\""
+    callOn net  $ "ip link set dev \"veth_" <> name <> "\" master br0 up"
+    callOn node $ "ip addr add " <> nodeIp node <> "/24 broadcast " <> makeIpAddress (netPrefix net) 255 <> " dev veth0"
+    callOn node $ "ip link set dev veth0 up"
+    callOn node $ "ip link set dev lo up"
+
+    return node
diff --git a/src/Process.hs b/src/Process.hs
index a93b464..a90a734 100644
--- a/src/Process.hs
+++ b/src/Process.hs
@@ -88,7 +88,8 @@ lineReadingLoop process h act =
 
 spawnOn :: Either Network Node -> ProcName -> Maybe Signal -> String -> TestRun Process
 spawnOn target pname killWith cmd = do
-    let prefix = either (const "") (\node -> "ip netns exec \"" ++ unpackNodeName (nodeName node) ++ "\" ") target
+    let netns = either netnsName netnsName target
+    let prefix = T.unpack $ "ip netns exec \"" <> netns <> "\" "
     (Just hin, Just hout, Just herr, handle) <- liftIO $ createProcess (shell $ prefix ++ cmd)
         { std_in = CreatePipe, std_out = CreatePipe, std_err = CreatePipe
         , env = Just [("EREBOS_DIR", either netDir nodeDir target)]