bgp: add BGP capability code 64 (graceful restart)

Signed-off-by: FUJITA Tomonori <fujita.tomonori@lab.ntt.co.jp>

2 files changed, 37 insertions, 0 deletions

diff --git a/ryu/lib/packet/bgp.py b/ryu/lib/packet/bgp.py
index 9cc98399..adeda92c 100644
--- a/ryu/lib/packet/bgp.py
+++ b/ryu/lib/packet/bgp.py
@@ -50,6 +50,7 @@ BGP_OPT_CAPABILITY = 2  # RFC 5492
 BGP_CAP_MULTIPROTOCOL = 1  # RFC 4760
 BGP_CAP_ROUTE_REFRESH = 2  # RFC 2918
 BGP_CAP_CARRYING_LABEL_INFO = 4  # RFC 3107
+BGP_CAP_GRACEFUL_RESTART = 64  # RFC 4724
 BGP_CAP_FOUR_OCTET_AS_NUMBER = 65  # RFC 4893
 BGP_CAP_ENHANCED_ROUTE_REFRESH = 70  # https://tools.ietf.org/html/\
 # draft-ietf-idr-bgp-enhanced-route-refresh-05
@@ -1196,6 +1197,39 @@ class BGPOptParamCapabilityEnhancedRouteRefresh(_OptParamEmptyCapability):
     pass
 
 
+@_OptParamCapability.register_type(BGP_CAP_GRACEFUL_RESTART)
+class BGPOptParamCapabilityGracefulRestart(_OptParamCapability):
+    _CAP_PACK_STR = "!H"
+
+    def __init__(self, flags, time, tuples, **kwargs):
+        super(BGPOptParamCapabilityGracefulRestart, self).__init__(**kwargs)
+        self.flags = flags
+        self.time = time
+        self.tuples = tuples
+
+    @classmethod
+    def parse_cap_value(cls, buf):
+        (restart, ) = struct.unpack_from(cls._CAP_PACK_STR, buffer(buf))
+        buf = buf[2:]
+        l = []
+        while len(buf) > 0:
+            l.append(struct.unpack_from("!HBB", buffer(buf)))
+            buf = buf[4:]
+        return {'flags': restart >> 12, 'time': restart & 0xfff, 'tuples': l}
+
+    def serialize_cap_value(self):
+        buf = bytearray()
+        msg_pack_into(self._CAP_PACK_STR, buf, 0, self.flags << 12 | self.time)
+        tuples = self.tuples
+        i = 0
+        offset = 2
+        for i in self.tuples:
+            afi, safi, flags = i
+            msg_pack_into("!HBB", buf, offset, afi, safi, flags)
+            offset += 4
+        return buf
+
+
 @_OptParamCapability.register_type(BGP_CAP_FOUR_OCTET_AS_NUMBER)
 class BGPOptParamCapabilityFourOctetAsNumber(_OptParamCapability):
     _CAP_PACK_STR = '!I'
diff --git a/ryu/tests/unit/packet/test_bgp.py b/ryu/tests/unit/packet/test_bgp.py
index c2592252..b9201c13 100644
--- a/ryu/tests/unit/packet/test_bgp.py
+++ b/ryu/tests/unit/packet/test_bgp.py
@@ -44,6 +44,9 @@ class Test_bgp(unittest.TestCase):
     def test_open2(self):
         opt_param = [bgp.BGPOptParamCapabilityUnknown(cap_code=200,
                                                       cap_value='hoge'),
+                     bgp.BGPOptParamCapabilityGracefulRestart(flags=0,
+                                                              time=120,
+                                                              tuples=[]),
                      bgp.BGPOptParamCapabilityRouteRefresh(),
                      bgp.BGPOptParamCapabilityCiscoRouteRefresh(),
                      bgp.BGPOptParamCapabilityMultiprotocol(