fix: add_path capability distinguishes afi/safi

yabgp/core/protocol.py

@@ -261,8 +261,9 @@ def _update_received(self, timestamp, msg):
         #     LOG.info(time.time())
 
         """Called when a BGP Update message was received."""
-        result = Update().parse(
-            timestamp, msg, self.fourbytesas, self.add_path_ipv4_receive, self.add_path_ipv4_send)
+        # TODO: Need to convert `self.add_path_ipv4_receive` and `self.add_path_ipv4_send` into a unified
+        #  `afi_add_path` format.
+        result = Update().parse(timestamp, msg, self.fourbytesas, afi_add_path={})
         if result['sub_error']:
             msg = {
                 'attr': result['attr'],

yabgp/message/attribute/mpreachnlri.py

@@ -67,11 +67,15 @@ class MpReachNLRI(Attribute):
     FLAG = AttributeFlag.OPTIONAL + AttributeFlag.EXTENDED_LENGTH
 
     @classmethod
-    def parse(cls, value, add_path=False):
+    def parse(cls, value, afi_add_path=None):
         """parse
         """
         try:
             afi, safi, nexthop_length = struct.unpack('!HBB', value[0:4])
+            if afi_add_path:
+                add_path = afi_add_path.get(bgp_cons.AFI_SAFI_DICT.get((afi, safi)), False)
+            else:
+                add_path = False
             nexthop_bin = value[4:4 + nexthop_length]
             nlri_bin = value[5 + nexthop_length:]
         except Exception:

yabgp/message/attribute/mpunreachnlri.py

@@ -58,9 +58,13 @@ class MpUnReachNLRI(Attribute):
     FLAG = AttributeFlag.OPTIONAL + AttributeFlag.EXTENDED_LENGTH
 
     @classmethod
-    def parse(cls, value, add_path=False):
+    def parse(cls, value, afi_add_path=None):
         try:
             afi, safi = struct.unpack('!HB', value[0:3])
+            if afi_add_path:
+                add_path = afi_add_path.get(bgp_cons.AFI_SAFI_DICT.get((afi, safi)), False)
+            else:
+                add_path = False
         except Exception:
             raise excep.UpdateMessageError(sub_error=bgp_cons.ERR_MSG_UPDATE_ATTR_LEN,
                                            data='')

yabgp/message/update.py

@@ -144,15 +144,14 @@ def __init__(self):
         """
 
     @classmethod
-    def parse(cls, t, msg_hex, asn4=False, add_path_remote=False, add_path_local=False):
+    def parse(cls, t, msg_hex, asn4=False, afi_add_path=None):
 
         """
         Parse BGP Update message
         :param t: timestamp
         :param msg_hex: raw message
         :param asn4: support 4 bytes AS or not
-        :param add_path_remote: if the remote peer can send add path NLRI
-        :param add_path_local: if the local can send add path NLRI
+        :param afi_add_path: support add-path or not in each afi/safi
         :return: message after parsing.
         """
         results = {
@@ -171,11 +170,13 @@ def parse(cls, t, msg_hex, asn4=False, add_path_remote=False, add_path_local=Fal
         attribute_data = msg_hex[withdraw_len + 4:withdraw_len + 4 + attr_len]
         nlri_data = msg_hex[withdraw_len + 4 + attr_len:]
         try:
+            add_path = afi_add_path.get('ipv4', False) if afi_add_path else False
+
             # parse withdraw prefixes
-            results['withdraw'] = cls.parse_prefix_list(withdraw_prefix_data, add_path_remote)
+            results['withdraw'] = cls.parse_prefix_list(withdraw_prefix_data, add_path)
 
             # parse nlri
-            results['nlri'] = cls.parse_prefix_list(nlri_data, add_path_remote)
+            results['nlri'] = cls.parse_prefix_list(nlri_data, add_path)
         except Exception as e:
             LOG.error(e)
             error_str = traceback.format_exc()
@@ -184,7 +185,7 @@ def parse(cls, t, msg_hex, asn4=False, add_path_remote=False, add_path_local=Fal
             results['err_data'] = ''
         try:
             # parse attributes
-            results['attr'] = cls.parse_attributes(attribute_data, asn4, add_path_remote)
+            results['attr'] = cls.parse_attributes(attribute_data, asn4, afi_add_path)
         except excep.UpdateMessageError as e:
             LOG.error(e)
             results['sub_error'] = e.sub_error
@@ -278,13 +279,13 @@ def parse_prefix_list(data, addpath=False):
         return prefixes
 
     @staticmethod
-    def parse_attributes(data, asn4=False, add_path=False):
+    def parse_attributes(data, asn4=False, afi_add_path=None):
         """
         Parses an RFC4271 encoded blob of BGP attributes into a list
 
         :param data:
         :param asn4: support 4 bytes asn or not
-        :param add_path: support add_path or not
+        :param afi_add_path: support add path or not in afi/safi
         :return:
         """
         attributes = {}
@@ -361,13 +362,13 @@ def parse_attributes(data, asn4=False, add_path=False):
                     decode_value = LargeCommunity.parse(value=attr_value)
 
                 elif type_code == bgp_cons.BGPTYPE_MP_REACH_NLRI:
-                    decode_value = MpReachNLRI.parse(value=attr_value, add_path=add_path)
+                    decode_value = MpReachNLRI.parse(value=attr_value, afi_add_path=afi_add_path)
                     if decode_value['nlri'][0] and type(decode_value['nlri'][0]) is dict:
                         if decode_value['nlri'][0].get("protocol_id"):
                             bgpls_pro_id = decode_value['nlri'][0]["protocol_id"]
 
                 elif type_code == bgp_cons.BGPTYPE_MP_UNREACH_NLRI:
-                    decode_value = MpUnReachNLRI.parse(value=attr_value, add_path=add_path)
+                    decode_value = MpUnReachNLRI.parse(value=attr_value, afi_add_path=afi_add_path)
 
                 elif type_code == bgp_cons.BGPTYPE_EXTENDED_COMMUNITY:
                     decode_value = ExtCommunity.parse(value=attr_value)

yabgp/tests/unit/message/attribute/test_mpreachnlri.py

@@ -23,6 +23,11 @@ class TestMpReachNLRI(unittest.TestCase):
     def setUp(self):
         self.maxDiff = None
 
+    def test_ipv4_unicast_path_id_parse(self):
+        data_bin = b'\x00\x01\x01\x04\n\x18%7\x00\x00\x00\x00\x02 \x05\x05\x05\x05'
+        data_hoped = {'afi_safi': (1, 1), 'nexthop': '10.24.37.55', 'nlri': [{'prefix': '5.5.5.5/32', 'path_id': 2}]}
+        self.assertEqual(data_hoped, MpReachNLRI.parse(data_bin, {'ipv4': True}))
+
     def test_ipv4_mpls_vpn_parse(self):
         data_bin = b'\x80\x0e\x21\x00\x01\x80\x0c\x00\x00\x00\x00\x00\x00\x00\x00\x02\x02\x02\x02' \
                    b'\x00\x78\x00\x01\x91\x00\x00\x00\x64\x00\x00\x00\x64\xaa\x00\x00\x00'

yabgp/tests/unit/message/attribute/test_mpunreachnlri.py

@@ -23,6 +23,11 @@ class TestMpUnReachNLRI(unittest.TestCase):
     def setUp(self):
         self.maxDiff = None
 
+    def test_ipv4_unicast_path_id_parse(self):
+        data_bin = b'\x00\x01\x01\x00\x00\x00\x02 \x05\x05\x05\x05'
+        data_hoped = {'afi_safi': (1, 1), 'withdraw': [{'prefix': '5.5.5.5/32', 'path_id': 2}]}
+        self.assertEqual(data_hoped, MpUnReachNLRI.parse(data_bin, {'ipv4': True}))
+
     def test_ipv4_mpls_vpn_parse(self):
         data_bin = b'\x80\x0f\x12\x00\x01\x80\x70\x80\x00\x00\x00\x00\x00\x02\x00\x00\x00\x02\xc0\xa8\xc9'
         data_hoped = {'afi_safi': (1, 128),

yabgp/tests/unit/message/test_update.py

@@ -107,7 +107,7 @@ def test_parse_ipv4_addpath_update(self):
                   b'\x01\x80\x04\x04\x00\x00\x00\x00\x40\x05\x04\x00\x00\x00\x64\x80\x0a\x04\x0a\x00\x22' \
                   b'\x04\x80\x09\x04\x0a\x00\x0f\x01\x00\x00\x00\x01\x20\x05\x05\x05\x05\x00\x00\x00\x01' \
                   b'\x20\xc0\xa8\x01\x05'
-        update = Update.parse(None, msg_hex, True, True)
+        update = Update.parse(None, msg_hex, True, {'ipv4': True})
         attributes = {1: 0, 2: [(2, [64511])], 3: '10.0.14.1', 4: 0, 5: 100, 9: '10.0.15.1', 10: ['10.0.34.4']}
         self.assertEqual(attributes, update['attr'])
         self.assertEqual([], update['withdraw'])
@@ -117,7 +117,7 @@ def test_parse_ipv4_addpath_update(self):
 
     def test_parse_ipv4_addpath_withdraw(self):
         msg_hex = b'\x00\x09\x00\x00\x00\x01\x20\x63\x63\x63\x63\x00\x00'
-        update = Update.parse(None, msg_hex, True, True)
+        update = Update.parse(None, msg_hex, True, {'ipv4': True})
         self.assertEqual([{'path_id': 1, 'prefix': '99.99.99.99/32'}], update['withdraw'])
 
     def test_parse_and_construct_ipv6_unicast_update(self):