Support to customize sport and dport for UDP traffic

[nfvbench.git] / nfvbench / traffic_gen / trex.py

diff --git a/nfvbench/traffic_gen/trex.py b/nfvbench/traffic_gen/trex.py
index 6c2a304..8aca290 100644 (file)
--- 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(