Support to customize sport and dport for UDP traffic

Change-Id: I40f28f28dd105f18ab69c5b5a56086ff3a06f2f6
Signed-off-by: Yichen Wang <yicwang@cisco.com>

3 files changed, 21 insertions, 3 deletions

diff --git a/nfvbench/cfg.default.yaml b/nfvbench/cfg.default.yaml
index 795ed5d..8de983f 100644
--- a/nfvbench/cfg.default.yaml
+++ b/nfvbench/cfg.default.yaml
@@ -152,12 +152,16 @@ traffic_generator:
     # `tg_gateway_ip_addrs__step`: step for generating traffic generator gateway sequences. default is 0.0.0.1
     # `gateway_ip_addrs`: base IPs of router gateways on both networks, quantity depends on chain count
     # `gateway_ip_addrs_step`: step for generating router gateway sequences. default is 0.0.0.1
+    # `udp_src_port`: the source port for sending UDP traffic, default is picked by TRex (53)
+    # `udp_dst_port`: the destination port for sending UDP traffic, default is picked by TRex (53)
     ip_addrs: ['10.0.0.0/8', '20.0.0.0/8']
     ip_addrs_step: 0.0.0.1
     tg_gateway_ip_addrs: ['1.1.0.100', '2.2.0.100']
     tg_gateway_ip_addrs_step: 0.0.0.1
     gateway_ip_addrs: ['1.1.0.2', '2.2.0.2']
     gateway_ip_addrs_step: 0.0.0.1
+    udp_src_port:
+    udp_dst_port:
 
     # Traffic Generator Profiles
     # In case you have multiple testbeds or traffic generators,
@@ -334,4 +338,4 @@ debug: false
 
 # Module and class name of factory which will be used to provide classes dynamically for other components.
 factory_module: 'nfvbench.factory'
-factory_class: 'BasicFactory'
\ No newline at end of file
+factory_class: 'BasicFactory'
diff --git a/nfvbench/traffic_client.py b/nfvbench/traffic_client.py
index 8bfcd76..319dc0b 100644
--- a/nfvbench/traffic_client.py
+++ b/nfvbench/traffic_client.py
@@ -84,7 +84,8 @@ class Device(object):
 
     def __init__(self, port, pci, switch_port=None, vtep_vlan=None, ip=None, tg_gateway_ip=None,
                  gateway_ip=None, ip_addrs_step=None, tg_gateway_ip_addrs_step=None,
-                 gateway_ip_addrs_step=None, chain_count=1, flow_count=1, vlan_tagging=False):
+                 gateway_ip_addrs_step=None, udp_src_port=None, udp_dst_port=None,
+                 chain_count=1, flow_count=1, vlan_tagging=False):
         self.chain_count = chain_count
         self.flow_count = flow_count
         self.dst = None
@@ -111,6 +112,8 @@ class Device(object):
         self.tg_gateway_ip_list = self.expand_ip(self.tg_gateway_ip,
                                                  self.tg_gateway_ip_addrs_step,
                                                  self.chain_count)
+        self.udp_src_port = udp_src_port
+        self.udp_dst_port = udp_dst_port
 
     def set_mac(self, mac):
         if mac is None:
@@ -151,6 +154,8 @@ class Device(object):
                 'ip_dst_addr_max': self.dst.ip_list[max_idx],
                 'ip_dst_count': ip_dst_count,
                 'ip_addrs_step': self.ip_addrs_step,
+                'udp_src_port': self.udp_src_port,
+                'udp_dst_port': self.udp_dst_port,
                 'mac_discovery_gw': self.gateway_ip_list[chain_idx],
                 'ip_src_tg_gw': self.tg_gateway_ip_list[chain_idx],
                 'ip_dst_tg_gw': self.dst.tg_gateway_ip_list[chain_idx],
@@ -253,6 +258,8 @@ class RunningTrafficProfile(object):
             'gateway_ip_addrs_step': self.gateway_ip_addrs_step,
             'tg_gateway_ip': generator_config.tg_gateway_ip_addrs[0],
             'tg_gateway_ip_addrs_step': self.tg_gateway_ip_addrs_step,
+            'udp_src_port': generator_config.udp_src_port,
+            'udp_dst_port': generator_config.udp_dst_port,
             'vlan_tagging': self.vlan_tagging
         }
         dst_config = {
@@ -264,6 +271,8 @@ class RunningTrafficProfile(object):
             'gateway_ip_addrs_step': self.gateway_ip_addrs_step,
             'tg_gateway_ip': generator_config.tg_gateway_ip_addrs[1],
             'tg_gateway_ip_addrs_step': self.tg_gateway_ip_addrs_step,
+            'udp_src_port': generator_config.udp_src_port,
+            'udp_dst_port': generator_config.udp_dst_port,
             'vlan_tagging': self.vlan_tagging
         }
 
diff --git a/nfvbench/traffic_gen/trex.py b/nfvbench/traffic_gen/trex.py
index 6c2a304..8aca290 100644
--- a/nfvbench/traffic_gen/trex.py
+++ b/nfvbench/traffic_gen/trex.py
@@ -139,7 +139,12 @@ class TRex(AbstractTrafficGenerator):
         if stream_cfg['vlan_tag'] is not None:
             pkt_base /= Dot1Q(vlan=stream_cfg['vlan_tag'])
 
-        pkt_base /= IP() / UDP()
+        udp_args = {}
+        if stream_cfg['udp_src_port']:
+            udp_args['sport'] = int(stream_cfg['udp_src_port'])
+        if stream_cfg['udp_dst_port']:
+            udp_args['dport'] = int(stream_cfg['udp_dst_port'])
+        pkt_base /= IP() / UDP(**udp_args)
 
         if stream_cfg['ip_addrs_step'] == 'random':
             src_fv = STLVmFlowVarRepetableRandom(