From 4914893445ae24bd1fa3b4aeea4910eeb412f7de Mon Sep 17 00:00:00 2001
From: Daniele Moro <daniele@opennetworking.org>
Date: Tue, 15 Sep 2020 13:07:07 +0200
Subject: [PATCH] Add support to install clone sessions

---
 utils/p4runtime_lib/helper.py            | 12 ++++++++++++
 utils/p4runtime_lib/simple_controller.py | 24 +++++++++++++++++++++++-
 utils/p4runtime_lib/switch.py            |  4 ++--
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/utils/p4runtime_lib/helper.py b/utils/p4runtime_lib/helper.py
index 8e3e20b..d5d76a3 100644
--- a/utils/p4runtime_lib/helper.py
+++ b/utils/p4runtime_lib/helper.py
@@ -198,3 +198,15 @@ class P4InfoHelper(object):
             r.instance = replica['instance']
             mc_entry.multicast_group_entry.replicas.extend([r])
         return mc_entry
+
+    def buildCloneSessionEntry(self, clone_session_id, replicas, packet_length_bytes=0):
+        clone_entry = p4runtime_pb2.PacketReplicationEngineEntry()
+        clone_entry.clone_session_entry.session_id = clone_session_id
+        clone_entry.clone_session_entry.packet_length_bytes = packet_length_bytes
+        clone_entry.clone_session_entry.class_of_service = 0  # PI currently supports only CoS=0 for clone session entry
+        for replica in replicas:
+            r = p4runtime_pb2.Replica()
+            r.egress_port = replica['egress_port']
+            r.instance = replica['instance']
+            clone_entry.clone_session_entry.replicas.extend([r])
+        return clone_entry
diff --git a/utils/p4runtime_lib/simple_controller.py b/utils/p4runtime_lib/simple_controller.py
index 94a13ab..4e78886 100755
--- a/utils/p4runtime_lib/simple_controller.py
+++ b/utils/p4runtime_lib/simple_controller.py
@@ -135,6 +135,13 @@ def program_switch(addr, device_id, sw_conf_file, workdir, proto_dump_fpath):
                 info(groupEntryToString(entry))
                 insertMulticastGroupEntry(sw, entry, p4info_helper)
 
+        if 'clone_session_entries' in sw_conf:
+            clone_entries = sw_conf['clone_session_entries']
+            info("Inserting %d clone entries..." % len(clone_entries))
+            for entry in clone_entries:
+                info(cloneEntryToString(entry))
+                insertCloneGroupEntry(sw, entry, p4info_helper)
+
     finally:
         sw.shutdown()
 
@@ -205,9 +212,24 @@ def groupEntryToString(rule):
     ports_str = ', '.join(replicas)
     return 'Group {0} => ({1})'.format(group_id, ports_str)
 
+def cloneEntryToString(rule):
+    clone_id = rule["clone_session_id"]
+    if "packet_length_bytes" in rule:
+        packet_length_bytes = str(rule["packet_length_bytes"])+"B"
+    else:
+        packet_length_bytes = "NO_TRUNCATION"
+    replicas = ['%d' % replica["egress_port"] for replica in rule['replicas']]
+    ports_str = ', '.join(replicas)
+    return 'Clone Session {0} => ({1}) ({2})'.format(clone_id, ports_str, packet_length_bytes)
+
 def insertMulticastGroupEntry(sw, rule, p4info_helper):
     mc_entry = p4info_helper.buildMulticastGroupEntry(rule["multicast_group_id"], rule['replicas'])
-    sw.WriteMulticastGroupEntry(mc_entry)
+    sw.WritePREEntry(mc_entry)
+
+def insertCloneGroupEntry(sw, rule, p4info_helper):
+    clone_entry = p4info_helper.buildCloneSessionEntry(rule['clone_session_id'], rule['replicas'],
+                                                       rule.get('packet_length_bytes', 0))
+    sw.WritePREEntry(clone_entry)
 
 
 if __name__ == '__main__':
diff --git a/utils/p4runtime_lib/switch.py b/utils/p4runtime_lib/switch.py
index a7c4f7a..254145c 100644
--- a/utils/p4runtime_lib/switch.py
+++ b/utils/p4runtime_lib/switch.py
@@ -133,13 +133,13 @@ class SwitchConnection(object):
                 yield response
 
 
-    def WriteMulticastGroupEntry(self, mc_entry, dry_run=False):
+    def WritePREEntry(self, pre_entry, dry_run=False):
         request = p4runtime_pb2.WriteRequest()
         request.device_id = self.device_id
         request.election_id.low = 1
         update = request.updates.add()
         update.type = p4runtime_pb2.Update.INSERT
-        update.entity.packet_replication_engine_entry.CopyFrom(mc_entry)
+        update.entity.packet_replication_engine_entry.CopyFrom(pre_entry)
         if dry_run:
             print "P4Runtime Write:", request
         else: