wlb: T114: Add firewall group support for WAN load balancer

8 files changed, 157 insertions, 2 deletions

diff --git a/data/config-mode-dependencies/vyos-1x.json b/data/config-mode-dependencies/vyos-1x.json
index ccfc022f4..d186c52bb 100644
--- a/data/config-mode-dependencies/vyos-1x.json
+++ b/data/config-mode-dependencies/vyos-1x.json
@@ -5,7 +5,7 @@
     },
     "firewall": {
         "conntrack": ["system_conntrack"],
-        "group_resync": ["system_conntrack", "nat", "policy_route"]
+        "group_resync": ["system_conntrack", "nat", "policy_route", "load-balancing_wan"]
     },
     "interfaces_bonding": {
         "ethernet": ["interfaces_ethernet"]
diff --git a/data/templates/load-balancing/nftables-wlb.j2 b/data/templates/load-balancing/nftables-wlb.j2
index 8afd0c162..6d71676c4 100644
--- a/data/templates/load-balancing/nftables-wlb.j2
+++ b/data/templates/load-balancing/nftables-wlb.j2
@@ -1,5 +1,7 @@
 #!/usr/sbin/nft -f
 
+{% import 'firewall/nftables-defines.j2' as group_tmpl %}
+
 {% if first_install is not vyos_defined %}
 delete table ip vyos_wanloadbalance
 {% endif %}
@@ -61,4 +63,6 @@ table ip vyos_wanloadbalance {
         meta mark set {{ state.mark }} ct mark set {{ state.mark }} counter accept
     }
 {% endfor %}
+
+{{ group_tmpl.groups(firewall_group, False, True) }}
 }
diff --git a/interface-definitions/load-balancing_wan.xml.in b/interface-definitions/load-balancing_wan.xml.in
index f80440411..24be00111 100644
--- a/interface-definitions/load-balancing_wan.xml.in
+++ b/interface-definitions/load-balancing_wan.xml.in
@@ -203,6 +203,7 @@
                   <help>Destination</help>
                 </properties>
                 <children>
+                  #include <include/firewall/source-destination-group-ipv4.xml.i>
                   #include <include/ipv4-address-prefix-range.xml.i>
                   #include <include/port-port-range.xml.i>
                 </children>
@@ -373,6 +374,7 @@
                   <help>Source information</help>
                 </properties>
                 <children>
+                  #include <include/firewall/source-destination-group-ipv4.xml.i>
                   #include <include/ipv4-address-prefix-range.xml.i>
                   #include <include/port-port-range.xml.i>
                 </children>
diff --git a/python/vyos/wanloadbalance.py b/python/vyos/wanloadbalance.py
index 158ee76c8..afe005731 100644
--- a/python/vyos/wanloadbalance.py
+++ b/python/vyos/wanloadbalance.py
@@ -67,6 +67,45 @@ def nft_rule(rule_conf, rule_id, local=False, exclude=False, limit=False, weight
                 port = port[1:]
             output.append(f'th {prefix}port {operator} {{ {port} }}')
 
+        if 'group' in direction_conf:
+                group = direction_conf['group']
+                if 'address_group' in group:
+                    group_name = group['address_group']
+                    operator = ''
+                    exclude = group_name[0] == "!"
+                    if exclude:
+                        operator = '!='
+                        group_name = group_name[1:]
+                    output.append(f'ip {prefix}addr {operator} @A_{group_name}')
+                if 'network_group' in group:
+                    group_name = group['network_group']
+                    operator = ''
+                    if group_name[0] == "!":
+                        operator = '!='
+                        group_name = group_name[1:]
+                    output.append(f'ip {prefix}addr {operator} @N_{group_name}')
+                # Generate firewall group domain-group
+                if 'domain_group' in group:
+                    group_name = group['domain_group']
+                    operator = ''
+                    if group_name[0] == '!':
+                        operator = '!='
+                        group_name = group_name[1:]
+                    output.append(f'ip {prefix}addr {operator} @D_{group_name}')
+                if 'port_group' in group:
+                    proto = rule_conf['protocol']
+                    group_name = group['port_group']
+
+                    if proto == 'tcp_udp':
+                        proto = 'th'
+
+                    operator = ''
+                    if group_name[0] == '!':
+                        operator = '!='
+                        group_name = group_name[1:]
+
+                    output.append(f'{proto} {prefix}port {operator} @P_{group_name}')
+
     if 'source_based_routing' not in rule_conf and not restore_mark:
         output.append('ct state new')
 
diff --git a/smoketest/scripts/cli/test_load-balancing_wan.py b/smoketest/scripts/cli/test_load-balancing_wan.py
index 904dabd04..2b262a006 100755
--- a/smoketest/scripts/cli/test_load-balancing_wan.py
+++ b/smoketest/scripts/cli/test_load-balancing_wan.py
@@ -369,5 +369,62 @@ echo "$ifname - $state" > {hook_output_path}
         with open(hook_output_path, 'r') as f:
             self.assertIn('eth0 - FAILED', f.read())
 
+    def test_firewall_groups(self):
+        isp1_iface = 'eth0'
+        isp2_iface = 'eth1'
+        lan_iface = 'eth2'
+
+        network_group1 = 'NET1'
+        network_group2 = 'NET2'
+        port_group = 'PORT1'
+
+        self.cli_set(['interfaces', 'ethernet', isp1_iface, 'address', '203.0.113.2/30'])
+        self.cli_set(['interfaces', 'ethernet', isp2_iface, 'address', '192.0.2.2/30'])
+        self.cli_set(['interfaces', 'ethernet', lan_iface, 'address', '198.51.100.2/30'])
+
+        self.cli_set(['firewall', 'group', 'network-group', network_group1, 'network', '10.0.0.0/8'])
+        self.cli_set(['firewall', 'group', 'network-group', network_group2, 'network', '198.51.100.0/24'])
+        self.cli_set(['firewall', 'group', 'port-group', port_group, 'port', '53'])
+
+        self.cli_set(base_path + ['wan', 'interface-health', isp1_iface, 'failure-count', '1'])
+        self.cli_set(base_path + ['wan', 'interface-health', isp1_iface, 'nexthop', '203.0.113.2'])
+        self.cli_set(base_path + ['wan', 'interface-health', isp1_iface, 'success-count', '1'])
+        self.cli_set(base_path + ['wan', 'interface-health', isp2_iface, 'failure-count', '1'])
+        self.cli_set(base_path + ['wan', 'interface-health', isp2_iface, 'nexthop', '192.0.2.2'])
+        self.cli_set(base_path + ['wan', 'interface-health', isp2_iface, 'success-count', '1'])
+        self.cli_set(base_path + ['wan', 'rule', '5', 'exclude'])
+        self.cli_set(base_path + ['wan', 'rule', '5', 'inbound-interface', 'eth*'])
+        self.cli_set(base_path + ['wan', 'rule', '5', 'destination', 'group', 'network-group', network_group1])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'failover'])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'inbound-interface', lan_iface])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'protocol', 'udp'])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'source', 'group', 'network-group', network_group2])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'source', 'group', 'port-group', port_group])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'destination', 'address', '192.0.2.0/24'])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'destination', 'group', 'port-group', port_group])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'interface', isp1_iface])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'interface', isp1_iface, 'weight', '10'])
+        self.cli_set(base_path + ['wan', 'rule', '10', 'interface', isp2_iface])
+
+        # commit changes
+        self.cli_commit()
+
+        time.sleep(5)
+
+        nftables_search = [
+            ['iifname "eth*"', f'ip daddr @N_{network_group1}', 'return'],
+            [
+                f'iifname "{lan_iface}"',
+                f'ip saddr @N_{network_group2}',
+                f'udp sport @P_{port_group}',
+                'ip daddr 192.0.2.0/24',
+                f'udp dport @P_{port_group}',
+                f'jump wlb_mangle_isp_{isp1_iface}',
+            ],
+        ]
+
+        self.verify_nftables_chain(nftables_search, 'ip vyos_wanloadbalance', 'wlb_mangle_prerouting')
+
+
 if __name__ == '__main__':
     unittest.main(verbosity=2)
diff --git a/src/conf_mode/load-balancing_wan.py b/src/conf_mode/load-balancing_wan.py
index dc7f1c838..8efb33f47 100755
--- a/src/conf_mode/load-balancing_wan.py
+++ b/src/conf_mode/load-balancing_wan.py
@@ -18,6 +18,7 @@ from sys import exit
 
 from vyos.config import Config
 from vyos.configdep import set_dependents, call_dependents
+from vyos.utils.dict import dict_search_args
 from vyos.utils.process import cmd
 from vyos import ConfigError
 from vyos import airbag
@@ -25,6 +26,13 @@ airbag.enable()
 
 service = 'vyos-wan-load-balance.service'
 
+valid_groups = [
+    'address_group',
+    'domain_group',
+    'network_group',
+    'port_group'
+]
+
 def get_config(config=None):
     if config:
         conf = config
@@ -38,6 +46,10 @@ def get_config(config=None):
                               get_first_key=True,
                               with_recursive_defaults=True)
 
+    if lb:
+        lb['firewall_group'] = conf.get_config_dict(['firewall', 'group'], key_mangling=('-', '_'), get_first_key=True,
+                                        no_tag_node_value_mangle=True)
+
     # prune limit key if not set by user
     for rule in lb.get('rule', []):
         if lb.from_defaults(['rule', rule, 'limit']):
@@ -89,6 +101,43 @@ def verify(lb):
 
             for direction in ['source', 'destination']:
                 if direction in rule_conf:
+                    side_conf = rule_conf[direction]
+
+                    if 'group' in side_conf:
+                        if len({'address_group', 'network_group', 'domain_group'} & set(side_conf['group'])) > 1:
+                            raise ConfigError('Only one address-group, network-group or domain-group can be specified')
+
+                        for group in valid_groups:
+                            if group in side_conf['group']:
+                                group_name = side_conf['group'][group]
+                                error_group = group.replace("_", "-")
+
+                                if group in ['address_group', 'network_group', 'domain_group']:
+                                    if 'address' in side_conf:
+                                        raise ConfigError(f'{error_group} and address cannot both be defined')
+
+                                if group in ['port_group']:
+                                    if 'port' in side_conf:
+                                        raise ConfigError(f'{error_group} and port cannot both be defined')
+
+                                if group_name and group_name[0] == '!':
+                                    group_name = group_name[1:]
+
+                                group_obj = dict_search_args(lb['firewall_group'], group, group_name)
+
+                                if group_obj is None:
+                                    raise ConfigError(f'Invalid {error_group} "{group_name}" on load-balancing wan rule')
+
+                                if not group_obj:
+                                    Warning(f'{error_group} "{group_name}" has no members!')
+
+                    if dict_search_args(side_conf, 'group', 'port_group'):
+                        if 'protocol' not in rule_conf:
+                            raise ConfigError('Protocol must be defined if specifying a port-group')
+
+                        if rule_conf['protocol'] not in ['tcp', 'udp', 'tcp_udp']:
+                            raise ConfigError('Protocol must be tcp, udp, or tcp_udp when specifying a port-group')
+
                     if 'port' in rule_conf[direction]:
                         if 'protocol' not in rule_conf:
                             raise ConfigError(f'Protocol required to specify port on load-balancing wan rule {rule_id}')
diff --git a/src/helpers/vyos-load-balancer.py b/src/helpers/vyos-load-balancer.py
index 6b260bf2c..5df83b31e 100755
--- a/src/helpers/vyos-load-balancer.py
+++ b/src/helpers/vyos-load-balancer.py
@@ -160,6 +160,9 @@ def get_config():
     lb = conf.get_config_dict(base, key_mangling=('-', '_'),
                             get_first_key=True, with_recursive_defaults=True)
 
+    lb['firewall_group'] = conf.get_config_dict(['firewall', 'group'], key_mangling=('-', '_'), get_first_key=True,
+                                    no_tag_node_value_mangle=True)
+
     # prune limit key if not set by user
     for rule in lb.get('rule', []):
         if lb.from_defaults(['rule', rule, 'limit']):
@@ -279,7 +282,7 @@ if __name__ == '__main__':
                         if state_changed and state['failure_count'] >= int(health_conf['failure_count']):
                             state['state'] = False
                             state['state_changed'] = True
-                    
+
                     #Force state changed to trigger the first write
                     if init == True:
                         state['state_changed'] = True
diff --git a/src/services/vyos-domain-resolver b/src/services/vyos-domain-resolver
index 17dae38e0..e1a52c93a 100755
--- a/src/services/vyos-domain-resolver
+++ b/src/services/vyos-domain-resolver
@@ -48,6 +48,7 @@ ipv4_tables = {
     'ip vyos_mangle',
     'ip vyos_filter',
     'ip vyos_nat',
+    'ip vyos_wanloadbalance',
     'ip raw'
 }