T5014: nat: add source and destination nat options for configuring load balance within a single rule.

6 files changed, 156 insertions, 2 deletions

diff --git a/interface-definitions/include/firewall/firewall-hashing-parameters.xml.i b/interface-definitions/include/firewall/firewall-hashing-parameters.xml.i
new file mode 100644
index 000000000..7f34de3ba
--- /dev/null
+++ b/interface-definitions/include/firewall/firewall-hashing-parameters.xml.i
@@ -0,0 +1,35 @@
+<!-- include start from firewall/firewall-hashing-parameters.xml.i -->
+<leafNode name="hash">
+  <properties>
+    <help>Define the parameters of the packet header to apply the hashing</help>
+    <completionHelp>
+      <list>source-address destination-address source-port destination-port random</list>
+    </completionHelp>
+    <valueHelp>
+      <format>source-address</format>
+      <description>Use source IP address for hashing</description>
+    </valueHelp>
+    <valueHelp>
+      <format>destination-address</format>
+      <description>Use destination IP address for hashing</description>
+    </valueHelp>
+    <valueHelp>
+      <format>source-port</format>
+      <description>Use source port for hashing</description>
+    </valueHelp>
+    <valueHelp>
+      <format>destination-port</format>
+      <description>Use destination port for hashing</description>
+    </valueHelp>
+    <valueHelp>
+      <format>random</format>
+      <description>Do not use information from ip header. Use random value.</description>
+    </valueHelp>
+    <constraint>
+      <regex>(source-address|destination-address|source-port|destination-port|random)</regex>
+    </constraint>
+    <multi/>
+  </properties>
+  <defaultValue>random</defaultValue>
+</leafNode>
+<!-- include end -->
\ No newline at end of file
diff --git a/interface-definitions/include/firewall/nat-balance.xml.i b/interface-definitions/include/firewall/nat-balance.xml.i
new file mode 100644
index 000000000..ac60a2545
--- /dev/null
+++ b/interface-definitions/include/firewall/nat-balance.xml.i
@@ -0,0 +1,28 @@
+<!-- include start from firewall/nat-balance.xml.i -->
+<tagNode name="member">
+  <properties>
+    <help>Translated IP address</help>
+    <valueHelp>
+      <format>ipv4</format>
+      <description>IPv4 address to match</description>
+    </valueHelp>
+    <constraint>
+      <validator name="ipv4-address"/>
+    </constraint>
+  </properties>
+  <children>
+    <leafNode name="weight">
+      <properties>
+        <help>Set probability for this output value</help>
+        <valueHelp>
+          <format>u32:1-100</format>
+          <description>Set probability for this output value</description>
+        </valueHelp>
+        <constraint>
+          <validator name="numeric" argument="--allow-range --range 1-100"/>
+        </constraint>
+      </properties>
+    </leafNode>
+  </children>
+</tagNode>
+<!-- include end -->
\ No newline at end of file
diff --git a/interface-definitions/include/nat-rule.xml.i b/interface-definitions/include/nat-rule.xml.i
index 7b3b8804e..fa7625c7d 100644
--- a/interface-definitions/include/nat-rule.xml.i
+++ b/interface-definitions/include/nat-rule.xml.i
@@ -25,6 +25,15 @@
     </node>
     #include <include/generic-disable-node.xml.i>
     #include <include/nat-exclude.xml.i>
+    <node name="balance">
+      <properties>
+        <help>Apply NAT balance</help>
+      </properties>
+      <children>
+        #include <include/firewall/firewall-hashing-parameters.xml.i>
+        #include <include/firewall/nat-balance.xml.i>
+      </children>
+    </node>
     <leafNode name="log">
       <properties>
         <help>NAT rule logging</help>
diff --git a/python/vyos/nat.py b/python/vyos/nat.py
index 5b8d5d1a3..9978993a7 100644
--- a/python/vyos/nat.py
+++ b/python/vyos/nat.py
@@ -90,6 +90,39 @@ def parse_nat_rule(rule_conf, rule_id, nat_type, ipv6=False):
         if options:
             translation_str += f' {",".join(options)}'
 
+        if 'member' in rule_conf['balance']:
+            hash_input_items = []
+            current_prob = 0
+            nat_map = []
+
+            for trans_addr, addr in rule_conf['balance']['member'].items():
+                item_prob = int(addr['weight'])
+                upper_limit = current_prob + item_prob - 1
+                hash_val = str(current_prob) + '-' + str(upper_limit)
+                element = hash_val + " : " + trans_addr
+                nat_map.append(element)
+                current_prob = current_prob + item_prob
+
+            elements = ' , '.join(nat_map)
+
+            if 'hash' in rule_conf['balance'] and 'random' in rule_conf['balance']['hash']:
+                translation_str += ' numgen random mod 100 map ' + '{ ' + f'{elements}' + ' }'
+            else:
+                for input_param in rule_conf['balance']['hash']:
+                    if input_param == 'source-address':
+                        param = 'ip saddr'
+                    elif input_param == 'destination-address':
+                        param = 'ip daddr'
+                    elif input_param == 'source-port':
+                        prot = rule_conf['protocol']
+                        param = f'{prot} sport'
+                    elif input_param == 'destination-port':
+                        prot = rule_conf['protocol']
+                        param = f'{prot} dport'
+                    hash_input_items.append(param)
+                hash_input = ' . '.join(hash_input_items)
+                translation_str += f' jhash ' + f'{hash_input}' + ' mod 100 map ' + '{ ' + f'{elements}' + ' }'
+
     for target in ['source', 'destination']:
         if target not in rule_conf:
             continue
diff --git a/smoketest/scripts/cli/test_nat.py b/smoketest/scripts/cli/test_nat.py
index 02fa03f7b..673be9905 100755
--- a/smoketest/scripts/cli/test_nat.py
+++ b/smoketest/scripts/cli/test_nat.py
@@ -231,5 +231,42 @@ class TestNAT(VyOSUnitTestSHIM.TestCase):
 
         self.verify_nftables(nftables_search, 'ip vyos_static_nat')
 
+    def test_nat_balance(self):
+        ifname = 'eth0'
+        member_1 = '198.51.100.1'
+        weight_1 = '10'
+        member_2 = '198.51.100.2'
+        weight_2 = '90'
+        member_3 = '192.0.2.1'
+        weight_3 = '35'
+        member_4 = '192.0.2.2'
+        weight_4 = '65'
+        dst_port = '443'
+
+        self.cli_set(dst_path + ['rule', '1', 'inbound-interface', ifname])
+        self.cli_set(dst_path + ['rule', '1', 'protocol', 'tcp'])
+        self.cli_set(dst_path + ['rule', '1', 'destination', 'port', dst_port])
+        self.cli_set(dst_path + ['rule', '1', 'balance', 'hash', 'source-address'])
+        self.cli_set(dst_path + ['rule', '1', 'balance', 'hash', 'source-port'])
+        self.cli_set(dst_path + ['rule', '1', 'balance', 'hash', 'destination-address'])
+        self.cli_set(dst_path + ['rule', '1', 'balance', 'hash', 'destination-port'])
+        self.cli_set(dst_path + ['rule', '1', 'balance', 'member', member_1, 'weight', weight_1])
+        self.cli_set(dst_path + ['rule', '1', 'balance', 'member', member_2, 'weight', weight_2])
+
+        self.cli_set(src_path + ['rule', '1', 'outbound-interface', ifname])
+        self.cli_set(src_path + ['rule', '1', 'balance', 'hash', 'random'])
+        self.cli_set(src_path + ['rule', '1', 'balance', 'member', member_3, 'weight', weight_3])
+        self.cli_set(src_path + ['rule', '1', 'balance', 'member', member_4, 'weight', weight_4])
+
+        self.cli_commit()
+
+        nftables_search = [
+            [f'iifname "{ifname}"', f'tcp dport {dst_port}', f'dnat to jhash ip saddr . tcp sport . ip daddr . tcp dport mod 100 map', f'0-9 : {member_1}, 10-99 : {member_2}'],
+            [f'oifname "{ifname}"', f'snat to numgen random mod 100 map', f'0-34 : {member_3}, 35-99 : {member_4}']
+        ]
+
+        self.verify_nftables(nftables_search, 'ip vyos_nat')
+
+
 if __name__ == '__main__':
     unittest.main(verbosity=2)
diff --git a/src/conf_mode/nat.py b/src/conf_mode/nat.py
index 5f4b658f8..dea833cf1 100755
--- a/src/conf_mode/nat.py
+++ b/src/conf_mode/nat.py
@@ -125,6 +125,18 @@ def verify_rule(config, err_msg, groups_dict):
                 if config['protocol'] not in ['tcp', 'udp', 'tcp_udp']:
                     raise ConfigError('Protocol must be tcp, udp, or tcp_udp when specifying a port-group')
 
+    if 'balance' in config:
+        for item in ['source-port', 'destination-port']:
+            if item in config['balance']['hash'] and config['protocol'] not in ['tcp', 'udp']:
+                raise ConfigError('Protocol must be tcp or udp when specifying hash ports')
+        count = 0
+        if 'member' in config['balance']:
+            for member in config['balance']['member']:
+                weight = config['balance']['member'][member]['weight']
+                count = count +  int(weight)
+            if count != 100:
+                Warning(f'Sum of weight for nat balance rule is not 100. You may get unexpected behaviour')
+
 def get_config(config=None):
     if config:
         conf = config
@@ -198,7 +210,7 @@ def verify(nat):
                 Warning(f'rule "{rule}" interface "{config["outbound_interface"]}" does not exist on this system')
 
             if not dict_search('translation.address', config) and not dict_search('translation.port', config):
-                if 'exclude' not in config:
+                if 'exclude' not in config and 'member' not in config['balance']:
                     raise ConfigError(f'{err_msg} translation requires address and/or port')
 
             addr = dict_search('translation.address', config)
@@ -222,7 +234,7 @@ def verify(nat):
                 Warning(f'rule "{rule}" interface "{config["inbound_interface"]}" does not exist on this system')
 
             if not dict_search('translation.address', config) and not dict_search('translation.port', config):
-                if 'exclude' not in config:
+                if 'exclude' not in config and 'member' not in config['balance']:
                     raise ConfigError(f'{err_msg} translation requires address and/or port')
 
             # common rule verification