Merge pull request #944 from sever-sever/T3702

pbr: T3702: Add rules match fwmark

3 files changed, 98 insertions, 6 deletions

diff --git a/interface-definitions/policy-local-route.xml.in b/interface-definitions/policy-local-route.xml.in
index 3769c3748..86445b65d 100644
--- a/interface-definitions/policy-local-route.xml.in
+++ b/interface-definitions/policy-local-route.xml.in
@@ -40,6 +40,18 @@
                   </leafNode>
                 </children>
               </node>
+              <leafNode name="fwmark">
+                <properties>
+                  <help>Match fwmark value</help>
+                  <valueHelp>
+                    <format>u32:1-2147483647</format>
+                    <description>Address to match against</description>
+                  </valueHelp>
+                  <constraint>
+                    <validator name="numeric" argument="--range 1-2147483647"/>
+                  </constraint>
+                </properties>
+              </leafNode>
               <leafNode name="source">
                 <properties>
                   <help>Source address or prefix</help>
diff --git a/smoketest/scripts/cli/test_policy.py b/smoketest/scripts/cli/test_policy.py
index 2d7b78048..485cc274a 100755
--- a/smoketest/scripts/cli/test_policy.py
+++ b/smoketest/scripts/cli/test_policy.py
@@ -1116,5 +1116,58 @@ class TestPolicy(VyOSUnitTestSHIM.TestCase):
 
         self.assertEqual(tmp, original)
 
+    # Test set table for fwmark
+    def test_fwmark_table_id(self):
+        path = base_path + ['local-route']
+
+        fwmk = '24'
+        rule = '101'
+        table = '154'
+
+        self.cli_set(path + ['rule', rule, 'set', 'table', table])
+        self.cli_set(path + ['rule', rule, 'fwmark', fwmk])
+
+        self.cli_commit()
+
+        # Check generated configuration
+
+        # Expected values
+        original = """
+        101:    from all fwmark 0x18 lookup 154
+        """
+        tmp = cmd('ip rule show prio 101')
+        original = original.split()
+        tmp = tmp.split()
+
+        self.assertEqual(tmp, original)
+
+    # Test set table for sources with fwmark
+    def test_fwmark_sources_table_id(self):
+        path = base_path + ['local-route']
+
+        sources = ['203.0.113.11', '203.0.113.12']
+        fwmk = '23'
+        rule = '100'
+        table = '150'
+        for src in sources:
+            self.cli_set(path + ['rule', rule, 'set', 'table', table])
+            self.cli_set(path + ['rule', rule, 'source', src])
+            self.cli_set(path + ['rule', rule, 'fwmark', fwmk])
+
+        self.cli_commit()
+
+        # Check generated configuration
+
+        # Expected values
+        original = """
+        100:	from 203.0.113.11 fwmark 0x17 lookup 150
+        100:	from 203.0.113.12 fwmark 0x17 lookup 150
+        """
+        tmp = cmd('ip rule show prio 100')
+        original = original.split()
+        tmp = tmp.split()
+
+        self.assertEqual(tmp, original)
+
 if __name__ == '__main__':
     unittest.main(verbosity=2)
diff --git a/src/conf_mode/policy-local-route.py b/src/conf_mode/policy-local-route.py
index 013f22665..0b7ceedeb 100755
--- a/src/conf_mode/policy-local-route.py
+++ b/src/conf_mode/policy-local-route.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 #
-# Copyright (C) 2020 VyOS maintainers and contributors
+# Copyright (C) 2020-2021 VyOS maintainers and contributors
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License version 2 or later as
@@ -44,17 +44,26 @@ def get_config(config=None):
     if tmp:
         for rule in (tmp or []):
             src = leaf_node_changed(conf, ['policy', 'local-route', 'rule', rule, 'source'])
+            fwmk = leaf_node_changed(conf, ['policy', 'local-route', 'rule', rule, 'fwmark'])
             if src:
                 dict = dict_merge({'rule_remove' : {rule : {'source' : src}}}, dict)
                 pbr.update(dict)
+            if fwmk:
+                dict = dict_merge({'rule_remove' : {rule : {'fwmark' : fwmk}}}, dict)
+                pbr.update(dict)
 
     # delete policy local-route rule x source x.x.x.x
+    # delete policy local-route rule x fwmark x
     if 'rule' in pbr:
         for rule in pbr['rule']:
             src = leaf_node_changed(conf, ['policy', 'local-route', 'rule', rule, 'source'])
+            fwmk = leaf_node_changed(conf, ['policy', 'local-route', 'rule', rule, 'fwmark'])
             if src:
                 dict = dict_merge({'rule_remove' : {rule : {'source' : src}}}, dict)
                 pbr.update(dict)
+            if fwmk:
+                dict = dict_merge({'rule_remove' : {rule : {'fwmark' : fwmk}}}, dict)
+                pbr.update(dict)
 
     return pbr
 
@@ -65,8 +74,8 @@ def verify(pbr):
 
     if 'rule' in pbr:
         for rule in pbr['rule']:
-            if 'source' not in pbr['rule'][rule]:
-                raise ConfigError('Source address required!')
+            if 'source' not in pbr['rule'][rule] and 'fwmark' not in pbr['rule'][rule]:
+                raise ConfigError('Source address or fwmark is required!')
             else:
                 if 'set' not in pbr['rule'][rule] or 'table' not in pbr['rule'][rule]['set']:
                     raise ConfigError('Table set is required!')
@@ -86,16 +95,34 @@ def apply(pbr):
     # Delete old rule if needed
     if 'rule_remove' in pbr:
         for rule in pbr['rule_remove']:
-            for src in pbr['rule_remove'][rule]['source']:
-                call(f'ip rule del prio {rule} from {src}')
+            if 'source' in pbr['rule_remove'][rule]:
+                for src in pbr['rule_remove'][rule]['source']:
+                    call(f'ip rule del prio {rule} from {src}')
+            if 'fwmark' in  pbr['rule_remove'][rule]:
+                for fwmk in pbr['rule_remove'][rule]['fwmark']:
+                    call(f'ip rule del prio {rule} from all fwmark {fwmk}')
 
     # Generate new config
     if 'rule' in pbr:
         for rule in pbr['rule']:
             table = pbr['rule'][rule]['set']['table']
-            if pbr['rule'][rule]['source']:
+            # Only source in the rule
+            # set policy local-route rule 100 source '203.0.113.1'
+            if 'source' in pbr['rule'][rule] and not 'fwmark' in pbr['rule'][rule]:
                 for src in pbr['rule'][rule]['source']:
                     call(f'ip rule add prio {rule} from {src} lookup {table}')
+            # Only fwmark in the rule
+            # set policy local-route rule 101 fwmark '23'
+            if 'fwmark' in pbr['rule'][rule] and not 'source' in pbr['rule'][rule]:
+                for fwmk in pbr['rule'][rule]['fwmark']:
+                    call(f'ip rule add prio {rule} from all fwmark {fwmk} lookup {table}')
+            # Source and fwmark in the rule
+            # set policy local-route rule 100 source '203.0.113.1'
+            # set policy local-route rule 100 fwmark '23'
+            if 'source' in pbr['rule'][rule] and 'fwmark' in pbr['rule'][rule]:
+                for src in pbr['rule'][rule]['source']:
+                    for fwmk in pbr['rule'][rule]['fwmark']:
+                        call(f'ip rule add prio {rule} from {src} fwmark {fwmk} lookup {table}')
 
     return None