Disconnect commands

4 files changed, 69 insertions, 3 deletions

diff --git a/src/Parser.hs b/src/Parser.hs
index aafba2c..a550a3f 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -666,6 +666,21 @@ testGuard = command "guard" $ Guard
     <$> cmdLine
     <*> param ""
 
+testDisconnectNode :: TestParser [TestStep]
+testDisconnectNode = command "disconnect_node" $ DisconnectNode
+    <$> paramOrContext ""
+    <*> innerBlock
+
+testDisconnectNodes :: TestParser [TestStep]
+testDisconnectNodes = command "disconnect_nodes" $ DisconnectNodes
+    <$> paramOrContext ""
+    <*> innerBlock
+
+testDisconnectUpstream :: TestParser [TestStep]
+testDisconnectUpstream = command "disconnect_upstream" $ DisconnectUpstream
+    <$> paramOrContext ""
+    <*> innerBlock
+
 testPacketLoss :: TestParser [TestStep]
 testPacketLoss = command "packet_loss" $ PacketLoss
     <$> param ""
@@ -702,6 +717,9 @@ testStep = choice
     , testSend
     , testExpect
     , testGuard
+    , testDisconnectNode
+    , testDisconnectNodes
+    , testDisconnectUpstream
     , testPacketLoss
     , testWait
     ]
diff --git a/src/Run.hs b/src/Run.hs
index 3591926..5b0ac2e 100644
--- a/src/Run.hs
+++ b/src/Run.hs
@@ -12,6 +12,7 @@ import Control.Monad.Reader
 
 import Data.Map qualified as M
 import Data.Maybe
+import Data.Set qualified as S
 import Data.Scientific
 import Data.Text (Text)
 import qualified Data.Text as T
@@ -60,6 +61,8 @@ runTest out opts test = do
             { tsNetwork = error "network not initialized"
             , tsVars = []
             , tsNodePacketLoss = M.empty
+            , tsDisconnectedUp = S.empty
+            , tsDisconnectedBridge = S.empty
             }
 
     let sigHandler SignalInfo { siginfoSpecific = chld } = do
@@ -149,6 +152,18 @@ evalSteps = mapM_ $ \case
     Guard line expr -> do
         testStepGuard line expr
 
+    DisconnectNode node inner -> do
+        n <- eval node
+        withDisconnectedUp n $ evalSteps inner
+
+    DisconnectNodes net inner -> do
+        n <- eval net
+        withDisconnectedBridge n $ evalSteps inner
+
+    DisconnectUpstream net inner -> do
+        n <- eval net
+        withDisconnectedUp n $ evalSteps inner
+
     PacketLoss loss node inner -> do
         l <- eval loss
         n <- eval node
@@ -190,16 +205,42 @@ withNode netexpr tvname inner = do
     node <- newNode net (either fromTypedVarName fromTypedVarName tvname)
     either (flip withVar node . fromTypedVarName) (const id) tvname $ inner node
 
+withDisconnectedUp :: HasNetns n => n -> TestRun a -> TestRun a
+withDisconnectedUp n inner = do
+    let netns = netnsName n
+    disconnected <- asks $ S.member netns . tsDisconnectedUp . snd
+    if disconnected
+      then inner
+      else do
+        local (fmap $ \s -> s { tsDisconnectedUp = S.insert netns $ tsDisconnectedUp s }) $ do
+            liftIO $ callOn n $ "ip link set veth0 down"
+            x <- inner
+            liftIO $ callOn n $ "ip link set veth0 up"
+            return x
+
+withDisconnectedBridge :: HasNetns n => n -> TestRun a -> TestRun a
+withDisconnectedBridge n inner = do
+    let netns = netnsName n
+    disconnected <- asks $ S.member netns . tsDisconnectedBridge . snd
+    if disconnected
+      then inner
+      else do
+        local (fmap $ \s -> s { tsDisconnectedBridge = S.insert netns $ tsDisconnectedBridge s }) $ do
+            liftIO $ callOn n $ "ip link set br0 down"
+            x <- inner
+            liftIO $ callOn n $ "ip link set br0 up"
+            return x
+
 withNodePacketLoss :: Node -> Scientific -> TestRun a -> TestRun a
 withNodePacketLoss node loss inner = do
-    x <- local (fmap $ \s -> s { tsNodePacketLoss = M.insertWith (\l l' -> 1 - (1 - l) * (1 - l')) (nodeName node) loss $ tsNodePacketLoss s }) $ do
+    x <- local (fmap $ \s -> s { tsNodePacketLoss = M.insertWith (\l l' -> 1 - (1 - l) * (1 - l')) (netnsName node) loss $ tsNodePacketLoss s }) $ do
         resetLoss
         inner
     resetLoss
     return x
   where
     resetLoss = do
-        tl <- asks $ fromMaybe 0 . M.lookup (nodeName node) . tsNodePacketLoss . snd
+        tl <- asks $ fromMaybe 0 . M.lookup (netnsName node) . tsNodePacketLoss . snd
         liftIO $ callOn node $ "tc qdisc replace dev veth0 root netem loss " <> T.pack (show (tl * 100)) <> "%"
 
 
diff --git a/src/Run/Monad.hs b/src/Run/Monad.hs
index 221f6d7..5b55897 100644
--- a/src/Run/Monad.hs
+++ b/src/Run/Monad.hs
@@ -16,11 +16,13 @@ import Control.Monad.Except
 import Control.Monad.Reader
 
 import Data.Map (Map)
+import Data.Set (Set)
 import Data.Scientific
 import qualified Data.Text as T
 
 import {-# SOURCE #-} GDB
 import {-# SOURCE #-} Network
+import Network.Ip
 import Output
 import {-# SOURCE #-} Process
 import Test
@@ -39,7 +41,9 @@ data TestEnv = TestEnv
 data TestState = TestState
     { tsNetwork :: Network
     , tsVars :: [(VarName, SomeVarValue)]
-    , tsNodePacketLoss :: Map NodeName Scientific
+    , tsDisconnectedUp :: Set NetworkNamespace
+    , tsDisconnectedBridge :: Set NetworkNamespace
+    , tsNodePacketLoss :: Map NetworkNamespace Scientific
     }
 
 data TestOptions = TestOptions
diff --git a/src/Test.hs b/src/Test.hs
index 0d96902..af1bace 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -42,6 +42,9 @@ data TestStep = forall a. ExprType a => Let SourceLine (TypedVarName a) (Expr a)
               | Send (Expr Process) (Expr Text)
               | Expect SourceLine (Expr Process) (Expr Regex) [TypedVarName Text] [TestStep]
               | Guard SourceLine (Expr Bool)
+              | DisconnectNode (Expr Node) [TestStep]
+              | DisconnectNodes (Expr Network) [TestStep]
+              | DisconnectUpstream (Expr Network) [TestStep]
               | PacketLoss (Expr Scientific) (Expr Node) [TestStep]
               | Wait