Network subnets and routing

4 files changed, 96 insertions, 19 deletions

diff --git a/src/Main.hs b/src/Main.hs
index 0d8a3fd..b49fe09 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -64,22 +64,28 @@ atomicallyTest act = do
         Left e -> throwError e
         Right x -> return x
 
-withNetwork :: (Network -> TestRun a) -> TestRun a
-withNetwork inner = do
+withInternet :: (Network -> TestRun a) -> TestRun a
+withInternet inner = do
     testDir <- asks $ optTestDir . teOptions . fst
     inet <- newInternet testDir
-    let net = inetRoot inet
+    res <- withNetwork (inetRoot inet) $ \net -> do
+        local (fmap $ \s -> s { tsNetwork = net }) $ inner net
+    delInternet inet
+    return res
 
+withSubnet :: Network -> Maybe (TypedVarName Network) -> (Network -> TestRun a) -> TestRun a
+withSubnet parent tvname inner = do
+    net <- newSubnet parent (fromTypedVarName <$> tvname)
+    withNetwork net inner
+
+withNetwork :: Network -> (Network -> TestRun a) -> TestRun a
+withNetwork net inner = do
     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
+            (path ++ " -i br0 -w '" ++ netDir net ++ "/br0.pcap' -U -Z root") . const
         Nothing -> id
 
-    res <- tcpdump $ do
-        local (fmap $ \s -> s { tsNetwork = net }) $ inner net
-
-    delInternet inet
-    return res
+    tcpdump $ inner net
 
 withNode :: Expr Network -> Either (TypedVarName Node) (TypedVarName Process) -> (Node -> TestRun a) -> TestRun a
 withNode netexpr tvname inner = do
@@ -156,6 +162,11 @@ evalSteps = mapM_ $ \case
         forM_ value $ \i -> do
             withVar name i $ evalSteps inner
 
+    Subnet name@(TypedVarName vname) parentExpr inner -> do
+        parent <- eval parentExpr
+        withSubnet parent (Just name) $ \net -> do
+            withVar vname net $ evalSteps inner
+
     DeclNode name@(TypedVarName vname) net inner -> do
         withNode net (Left name) $ \node -> do
             withVar vname node $ evalSteps inner
@@ -243,7 +254,7 @@ runTest out opts test = do
     oldHandler <- installHandler processStatusChanged (CatchInfo sigHandler) Nothing
 
     res <- runExceptT $ flip runReaderT (tenv, tstate) $ fromTestRun $ do
-        withNetwork $ \_ -> evalSteps (testSteps test)
+        withInternet $ \_ -> evalSteps (testSteps test)
 
     void $ installHandler processStatusChanged oldHandler Nothing
 
diff --git a/src/Network.hs b/src/Network.hs
index 29621fc..ec0b380 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -9,6 +9,7 @@ module Network (
     callOn,
 
     newInternet, delInternet,
+    newSubnet,
     newNode,
 ) where
 
@@ -27,6 +28,29 @@ import System.Process
 
 import Test
 
+{-
+NETWORK STRUCTURE
+=================
+
+Local network (namespace "s<PREFIX>", e.g. "s1_2"):
+
+  (upstream, if any)                 (to subnets, if any and prefix length < 24)
+         ↑                           veth_sX_1 (IP: prefix.1(.0)*.254)
+       veth0                         veth_sX_2 (IP: prefix.2(.0)*.254) → veth0 in subnet namespace
+         |                           veth_sX_3 (IP: prefix.3(.0)*.254)
+        br0 (IP: prefix(.0)*.1/24)   ...
+       / | \
+  veth2 ... veth253
+    ↓    ↓    ↓
+     (to nodes)
+
+Node (namespace "s<PREFIX>:<NODE>", e.g. "s1_2:p0"):
+
+     (upstream)
+         ↑
+       veth0 (IP: prefix.N/24)
+-}
+
 data Internet = Internet
     { inetDir :: FilePath
     , inetRoot :: Network
@@ -35,6 +59,7 @@ data Internet = Internet
 data Network = Network
     { netPrefix :: [Word8]
     , netNodes :: TVar [Node]
+    , netSubnets :: TVar [(Word8, Network)]
     , netDir :: FilePath
     }
 
@@ -77,7 +102,7 @@ callOn n cmd = callCommand $ T.unpack $ "ip netns exec \"" <> netnsName n <> "\"
 instance ExprType Network where
     textExprType _ = T.pack "network"
     textExprValue n = "s:" <> textNetworkName n
-    emptyVarValue = Network [] undefined undefined
+    emptyVarValue = Network [] undefined undefined undefined
 
 instance ExprType Node where
     textExprType _ = T.pack "node"
@@ -89,6 +114,9 @@ instance ExprType Node where
         ]
 
 
+nextPrefix :: [Word8] -> [Word8] -> Word8
+nextPrefix _ used = maximum (0 : used) + 1
+
 makeIpAddress :: [Word8] -> Word8 -> Text
 makeIpAddress prefix num = T.intercalate "." $ map (T.pack . show) $ prefix ++ replicate (3 - length prefix) 0 ++ [num]
 
@@ -105,46 +133,76 @@ delInternet :: MonadIO m => Internet -> m ()
 delInternet _ = liftIO $ do
     callCommand $ "ip -all netns delete"
 
+newSubnet :: MonadIO m => Network -> Maybe VarName -> m Network
+newSubnet net vname = do
+    sub <- liftIO $ atomically $ do
+        pref <- nextPrefix (netPrefix net) . map fst <$> readTVar (netSubnets net)
+        sub <- newNetwork
+            (netPrefix net ++ [pref])
+            (netDir net </> maybe (T.unpack $ netnsName net) (("sub_"++) . unpackVarName) vname)
+        modifyTVar (netSubnets net) ((pref, sub) :)
+        return sub
+    initNetwork sub
+    liftIO $ do
+        callOn net $ "ip link add \"veth_" <> netnsName sub <> "\" type veth peer name veth0 netns \"" <> netnsName sub <> "\""
+        callOn net $ "ip addr add dev \"veth_" <> netnsName sub <> "\" " <> makeIpAddress (netPrefix sub) 254 <> "/24"
+        callOn net $ "ip link set dev \"veth_" <> netnsName sub <> "\" up"
+
+        -- If the new subnet can be split further, routing rule for the whole prefix is needed
+        when (length (netPrefix sub) < 3) $ callOn net $ "ip route add "
+            <> makeIpAddress (netPrefix sub) 0 <> "/" <> (T.pack $ show $ length (netPrefix sub) * 8)
+            <> " via " <> makeIpAddress (netPrefix sub) 1
+            <> " dev \"veth_" <> netnsName sub
+            <> "\" src " <> makeIpAddress (netPrefix sub) 254
+
+        callOn sub $ "ip link set dev veth0 master br0 up"
+        callOn sub $ "ip route add default via " <> makeIpAddress (netPrefix sub) 254 <> " dev br0 src " <> makeIpAddress (netPrefix sub) 1
+    return sub
+
 newNetwork :: [Word8] -> FilePath -> STM Network
 newNetwork prefix dir = do
     Network
         <$> pure prefix
         <*> newTVar []
+        <*> newTVar []
         <*> pure dir
 
 initNetwork :: MonadIO m => Network -> m ()
 initNetwork net = liftIO $ do
+    createDirectoryIfMissing True $ netDir net
     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"
+    callOn net $ "ip route add " <> makeIpAddress (netPrefix net) 0 <> "/24 dev br0 src " <> makeIpAddress (netPrefix net) 1
 
 newNode :: MonadIO m => Network -> VarName -> m Node
 newNode net vname = liftIO $ do
-    node <- atomically $ do
+    (node, idx) <- atomically $ do
         nodes <- readTVar (netNodes net)
         let nname = nextNodeName vname $ map nodeName nodes
+            idx = fromIntegral $ 2 + length nodes
             node = Node { nodeName = nname
-                        , nodeIp = makeIpAddress (netPrefix net) (fromIntegral $ 2 + length nodes)
+                        , nodeIp = makeIpAddress (netPrefix net) idx
                         , nodeNetwork = net
                         , nodeDir = netDir net </> ("node_" ++ unpackNodeName nname)
                         }
         writeTVar (netNodes net) (node : nodes)
-        return node
-
-    let name = textNodeName $ nodeName node
-        dir = nodeDir node
+        return (node, idx)
 
+    let dir = nodeDir node
     exists <- doesPathExist dir
     when exists $ ioError $ userError $ dir ++ " exists"
     createDirectoryIfMissing True dir
 
+    let veth = T.pack $ "veth" <> show idx
     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 net  $ "ip link add " <> veth <> " type veth peer name veth0 netns \"" <> netnsName node <> "\""
+    callOn net  $ "ip link set dev " <> veth <> " 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"
+    callOn node $ "ip route add default via " <> makeIpAddress (netPrefix net) 1 <> " dev veth0 src " <> nodeIp node
 
     return node
diff --git a/src/Parser.hs b/src/Parser.hs
index b79931b..aafba2c 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -630,6 +630,12 @@ testWith = do
         modify $ \s -> s { testContext = ctx }
         testBlock indent
 
+testSubnet :: TestParser [TestStep]
+testSubnet = command "subnet" $ Subnet
+    <$> param ""
+    <*> paramOrContext "of"
+    <*> innerBlock
+
 testNode :: TestParser [TestStep]
 testNode = command "node" $ DeclNode
     <$> param ""
@@ -690,6 +696,7 @@ testStep = choice
     , forStatement
     , testLocal
     , testWith
+    , testSubnet
     , testNode
     , testSpawn
     , testSend
diff --git a/src/Test.hs b/src/Test.hs
index 6077b92..0d96902 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -36,6 +36,7 @@ data Test = Test
 
 data TestStep = forall a. ExprType a => Let SourceLine (TypedVarName a) (Expr a) [TestStep]
               | forall a. ExprType a => For SourceLine (TypedVarName a) (Expr [a]) [TestStep]
+              | Subnet (TypedVarName Network) (Expr Network) [TestStep]
               | DeclNode (TypedVarName Node) (Expr Network) [TestStep]
               | Spawn (TypedVarName Process) (Either (TypedVarName Node) (Either (Expr Network) (Expr Node))) [TestStep]
               | Send (Expr Process) (Expr Text)