From 69e9c1225896bc828c852bbe99730e643ebd2ee7 Mon Sep 17 00:00:00 2001
From: Sergi Delgado Segura <sergi.delgado.s@gmail.com>
Date: Wed, 22 Apr 2020 14:50:13 +0200
Subject: [PATCH] gatekeeper reformats update_available_slots to
 add_update_appointment

- add_update_appointment handles the slot count and also keeps a copy of the appointment uuid and slots taken.
- gatekeeper.appointments is turned to a dict uuid:slots so appointment updates can be computed with no further data
- deleting appointment will also allow to update the slot count with no further data from both the Watcher and the Responder
---
 teos/gatekeeper.py                | 41 +++++++++++++++++--------------
 test/teos/unit/test_gatekeeper.py | 33 ++++++++++++-------------
 2 files changed, 38 insertions(+), 36 deletions(-)

diff --git a/teos/gatekeeper.py b/teos/gatekeeper.py
index d0416a7..b333501 100644
--- a/teos/gatekeeper.py
+++ b/teos/gatekeeper.py
@@ -27,9 +27,9 @@ class UserInfo:
         self.available_slots = available_slots
         self.subscription_expiry = subscription_expiry
 
-        # FIXME: this list is currently never wiped
         if not appointments:
-            self.appointments = []
+            # A dictionary of the form uuid:required_slots for each user appointment
+            self.appointments = {}
         else:
             self.appointments = appointments
 
@@ -136,21 +136,20 @@ class Gatekeeper:
         except (InvalidParameter, InvalidKey, SignatureError):
             raise AuthenticationFailure("Wrong message or signature.")
 
-    def update_available_slots(self, user_id, new_appointment, old_appointment=None):
+    def add_update_appointment(self, user_id, uuid, appointment):
         """
-        Updates (add/removes) slots from a user subscription.
+        Adds (or updates) an appointment to a user subscription. The user slots are updated accordingly.
 
-        Slots are removed if a new appointment is given, or an update is given with an appointment bigger than the
-        old one.
+        Slots are taken if a new appointment is given, or an update is given with an appointment bigger than the
+        existing one.
 
-        Slots are added if an update is given but the new appointment is smaller than the old one.
+        Slots are given back if an update is given but the new appointment is smaller than the existing one.
 
         Args:
-            user_id(:obj:`str`): the public key that identifies the user (33-bytes hex str).
-            new_appointment (:obj:`dict`): the summary of new appointment the user is requesting
-                to add.
-            old_appointment (:obj:`dict`): the summary old appointment the user wants to replace.
-                Optional.
+            user_id (:obj:`str`): the public key that identifies the user (33-bytes hex str).
+            uuid (:obj:`str`): the appointment uuid.
+            appointment (:obj:`ExtendedAppointment <teos.extended_appointment.ExtendedAppointment`): the summary of new
+                appointment the user is requesting.
 
         Returns:
             :obj:`int`: the number of remaining appointment slots.
@@ -160,18 +159,21 @@ class Gatekeeper:
         """
 
         self.lock.acquire()
-        if old_appointment:
-            # For updates the difference between the existing appointment and the update is computed.
-            used_slots = ceil(old_appointment.get("size") / ENCRYPTED_BLOB_MAX_SIZE_HEX)
-            required_slots = ceil(new_appointment.get("size") / ENCRYPTED_BLOB_MAX_SIZE_HEX) - used_slots
+        # For updates the difference between the existing appointment and the update is computed.
+        if uuid in self.registered_users[user_id].appointments:
+            used_slots = self.registered_users[user_id].appointments[uuid]
+
         else:
             # For regular appointments 1 slot is reserved per ENCRYPTED_BLOB_MAX_SIZE_HEX block.
-            required_slots = ceil(new_appointment.get("size") / ENCRYPTED_BLOB_MAX_SIZE_HEX)
+            used_slots = 0
 
-        if required_slots <= self.registered_users.get(user_id).available_slots:
+        required_slots = ceil(len(appointment.encrypted_blob) / ENCRYPTED_BLOB_MAX_SIZE_HEX)
+
+        if required_slots - used_slots <= self.registered_users.get(user_id).available_slots:
             # Filling / freeing slots depending on whether this is an update or not, and if it is bigger or smaller than
             # the old appointment.
-            self.registered_users.get(user_id).available_slots -= required_slots
+            self.registered_users.get(user_id).appointments[uuid] = required_slots
+            self.registered_users.get(user_id).available_slots -= required_slots - used_slots
         else:
             self.lock.release()
             raise NotEnoughSlots()
@@ -189,6 +191,7 @@ class Gatekeeper:
         Returns:
             :obj:`list`: a list of appointment uuids that will expire at ``block_height``.
         """
+
         expired_appointments = []
         # Avoiding dictionary changed size during iteration
         for user_id in list(self.registered_users.keys()):
diff --git a/test/teos/unit/test_gatekeeper.py b/test/teos/unit/test_gatekeeper.py
index f878aca..a688fce 100644
--- a/test/teos/unit/test_gatekeeper.py
+++ b/test/teos/unit/test_gatekeeper.py
@@ -128,48 +128,47 @@ def test_identify_user_wrong(gatekeeper):
         gatekeeper.authenticate_user(message, signature.encode())
 
 
-def test_update_available_slots(gatekeeper):
-    # update_available_slots should decrease the slot count if a new appointment is added
+def test_add_update_appointment(gatekeeper):
+    # add_update_appointment should decrease the slot count if a new appointment is added
     # let's add a new user
     sk, pk = generate_keypair()
     user_id = Cryptographer.get_compressed_pk(pk)
     gatekeeper.add_update_user(user_id)
 
-    # And now update the slots given an appointment
+    # And now update add a new appointment
     appointment, _ = generate_dummy_appointment()
-    gatekeeper.update_available_slots(user_id, appointment.get_summary())
+    appointment_uuid = get_random_value_hex(16)
+    remaining_slots = gatekeeper.add_update_appointment(user_id, appointment_uuid, appointment)
 
     # This is a standard size appointment, so it should have reduced the slots by one
-    assert gatekeeper.registered_users[user_id].available_slots == config.get("DEFAULT_SLOTS") - 1
+    assert appointment_uuid in gatekeeper.registered_users[user_id].appointments
+    assert remaining_slots == config.get("DEFAULT_SLOTS") - 1
 
-    # Updates can leave the count as it, decrease it, or increase it, depending on the appointment size (modulo
+    # Updates can leave the count as is, decrease it, or increase it, depending on the appointment size (modulo
     # ENCRYPTED_BLOB_MAX_SIZE_HEX)
 
     # Appointments of the same size leave it as is
     appointment_same_size, _ = generate_dummy_appointment()
-    remaining_slots = gatekeeper.update_available_slots(
-        user_id, appointment.get_summary(), appointment_same_size.get_summary()
-    )
+    remaining_slots = gatekeeper.add_update_appointment(user_id, appointment_uuid, appointment)
+    assert appointment_uuid in gatekeeper.registered_users[user_id].appointments
     assert remaining_slots == config.get("DEFAULT_SLOTS") - 1
 
     # Bigger appointments decrease it
     appointment_x2_size = appointment_same_size
     appointment_x2_size.encrypted_blob = "A" * (ENCRYPTED_BLOB_MAX_SIZE_HEX + 1)
-    remaining_slots = gatekeeper.update_available_slots(
-        user_id, appointment_x2_size.get_summary(), appointment.get_summary()
-    )
+    remaining_slots = gatekeeper.add_update_appointment(user_id, appointment_uuid, appointment_x2_size)
+    assert appointment_uuid in gatekeeper.registered_users[user_id].appointments
     assert remaining_slots == config.get("DEFAULT_SLOTS") - 2
 
-    # Smaller appointments increase it (using the same data but flipped)
-    remaining_slots = gatekeeper.update_available_slots(
-        user_id, appointment.get_summary(), appointment_x2_size.get_summary()
-    )
+    # Smaller appointments increase it
+    remaining_slots = gatekeeper.add_update_appointment(user_id, appointment_uuid, appointment)
     assert remaining_slots == config.get("DEFAULT_SLOTS") - 1
 
     # If the appointment needs more slots than there's free, it should fail
     gatekeeper.registered_users[user_id].available_slots = 1
+    appointment_uuid = get_random_value_hex(16)
     with pytest.raises(NotEnoughSlots):
-        gatekeeper.update_available_slots(user_id, appointment_x2_size.get_summary())
+        gatekeeper.add_update_appointment(user_id, appointment_uuid, appointment_x2_size)
 
 
 def test_get_expired_appointments(gatekeeper):