T8963: policy-route: trigger domain resolver for domain groups

Signed-off-by: Miaosen Wang <secretandanon@gmail.com>

2 files changed, 88 insertions, 0 deletions

diff --git a/smoketest/scripts/cli/test_policy_route.py b/smoketest/scripts/cli/test_policy_route.py
index 939e4bc53..d39044b4b 100755
--- a/smoketest/scripts/cli/test_policy_route.py
+++ b/smoketest/scripts/cli/test_policy_route.py
@@ -18,6 +18,8 @@ import unittest
 
 from base_vyostest_shim import VyOSUnitTestSHIM
 
+from vyos.utils.process import run
+
 mark = '100'
 conn_mark = '555'
 conn_mark_set = '111'
@@ -53,10 +55,13 @@ class TestPolicyRoute(VyOSUnitTestSHIM.TestCase):
     def tearDown(self):
         self.cli_delete(['policy', 'route'])
         self.cli_delete(['policy', 'route6'])
+        self.cli_delete(['firewall', 'group', 'domain-group', 'smoketest_domain'])
+        self.cli_delete(['system', 'static-host-mapping', 'host-name', 'pbr.example.com'])
         self.cli_commit()
 
         # Verify nftables cleanup
         nftables_search = [
+            ['set D_smoketest_domain'],
             ['set N_smoketest_network'],
             ['set N_smoketest_network1'],
             ['chain VYOS_PBR_smoketest']
@@ -94,6 +99,46 @@ class TestPolicyRoute(VyOSUnitTestSHIM.TestCase):
 
         self.cli_delete(['firewall'])
 
+    def test_pbr_domain_group(self):
+        domain_group = 'smoketest_domain'
+        domain_name = 'pbr.example.com'
+        domain_ip = '192.0.2.5'
+
+        self.cli_set(['system', 'static-host-mapping', 'host-name', domain_name, 'inet', domain_ip])
+        self.cli_commit()
+
+        self.cli_set(['firewall', 'group', 'domain-group', domain_group, 'address', domain_name])
+        self.cli_commit()
+
+        self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'destination', 'group', 'domain-group', domain_group])
+        self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'set', 'mark', mark])
+        self.cli_set(['policy', 'route', 'smoketest', 'interface', interface])
+
+        self.cli_commit()
+
+        result, code = self.wait_for_result(
+            lambda: run(
+                f'sudo nft get element ip vyos_mangle '
+                f'D_{domain_group} {{ {domain_ip} }}'
+            ),
+            0,
+            pause=1,
+            timeout=10
+        )
+        self.assertTrue(
+            result,
+            f'Expected {domain_ip} in D_{domain_group}, last nft exit code {code}'
+        )
+
+        nftables_search = [
+            [f'iifname "{interface}"', 'jump VYOS_PBR_UD_smoketest'],
+            [f'set D_{domain_group}'],
+            [f'elements = {{ {domain_ip} }}'],
+            [f'ip daddr @D_{domain_group}', 'meta mark set']
+        ]
+
+        self.verify_nftables(nftables_search, 'ip vyos_mangle')
+
     def test_pbr_mark(self):
         self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'source', 'address', '172.16.20.10'])
         self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'destination', 'address', '172.16.10.10'])
diff --git a/src/conf_mode/policy_route.py b/src/conf_mode/policy_route.py
index 3cfdad913..63b6ec257 100755
--- a/src/conf_mode/policy_route.py
+++ b/src/conf_mode/policy_route.py
@@ -16,6 +16,7 @@
 
 import os
 
+from glob import glob
 from json import loads
 from sys import exit
 
@@ -25,6 +26,8 @@ from vyos.configdiff import Diff, get_config_diff
 from vyos.template import render
 from vyos.utils.dict import dict_search_args
 from vyos.utils.dict import dict_search_recursive
+from vyos.utils.file import write_file
+from vyos.utils.process import call
 from vyos.utils.process import cmd
 from vyos.utils.process import run
 from vyos.utils.network import get_vrf_tableid
@@ -38,6 +41,8 @@ airbag.enable()
 
 mark_offset = 0x7FFFFFFF
 nftables_conf = '/run/nftables_policy.conf'
+domain_resolver_usage = '/run/use-vyos-domain-resolver-policy-route'
+domain_resolver_usage_glob = '/run/use-vyos-domain-resolver*'
 
 valid_groups = [
     'address_group',
@@ -187,6 +192,42 @@ def generate(policy):
     render(nftables_conf, 'firewall/nftables-policy.j2', policy)
     return None
 
+def domain_group_used(policy):
+    for route in ['route', 'route6']:
+        if route not in policy:
+            continue
+
+        for pol_conf in policy[route].values():
+            if 'rule' not in pol_conf:
+                continue
+
+            for rule_conf in pol_conf['rule'].values():
+                if 'disable' in rule_conf:
+                    continue
+
+                for side in ['destination', 'source']:
+                    if dict_search_args(rule_conf, side, 'group', 'domain_group') is not None:
+                        return True
+
+    return False
+
+def update_domain_resolver(policy):
+    domain_action = None
+
+    if domain_group_used(policy):
+        text = '# Automatically generated by policy_route.py\nThis file indicates that vyos-domain-resolver service is used by policy_route.\n'
+        write_file(domain_resolver_usage, text)
+        domain_action = 'restart'
+    elif os.path.exists(domain_resolver_usage):
+        os.unlink(domain_resolver_usage)
+        domain_action = 'restart'
+
+        if not glob(domain_resolver_usage_glob):
+            domain_action = 'stop'
+
+    if domain_action:
+        call(f'systemctl {domain_action} vyos-domain-resolver.service')
+
 def apply_table_marks(policy):
     for route in ['route', 'route6']:
         if route in policy:
@@ -240,6 +281,8 @@ def apply(policy):
 
     apply_table_marks(policy)
 
+    update_domain_resolver(policy)
+
     if policy['geoip_sets']['name'] or policy['geoip_sets']['ipv6_name']:
         if policy['geoip_updated'] or not geoip_refresh():
             print('Updating GeoIP. Please wait...')