Repository: hydroshare/hydroshare

Issue:
Check user quota before uploading files to iRods userZone
The goal is to validate a user's quota before allowing them to upload to the iRods userZone.
I'm not confident the best approach here. It might look like a check for file headers before the file upload is initiated

This is a key component for [Quota enforcement](https://github.com/hydroshare/hydroshare/issues/5228)


Discussion:
Chris helped me brainstorm this and one option is to have Hydroshare execute (something like [how we create userZone account](https://github.com/hydroshare/hydroshare/blob/2.12/irods/create_user.sh#L1) or how we make icommand calls to [transfer from the uzone](https://github.com/hydroshare/hydroshare/blob/2.12/irods_browser_app/views.py#L108)) an [irule](https://docs.irods.org/4.2.11/icommands/user/#irule) that prevents `iput` in the userZone, but still allows all other functionality. This way uploads would be disallowed for users over quota but they would still be able to delete/zip to free up some space.

However, this would effectively be very similar to what we already have in HS quota management -- where the devil is in the details of syncing quota between the 2 systems. It has the potential to make something that is already fairly complicated much more so...

Resolving patch:
diff --git a/hs_core/hydroshare/resource.py b/hs_core/hydroshare/resource.py
--- a/hs_core/hydroshare/resource.py
+++ b/hs_core/hydroshare/resource.py
@@ -22,7 +22,6 @@
 from hs_access_control.models import ResourceAccess, UserResourcePrivilege, PrivilegeCodes
 from hs_labels.models import ResourceLabels
 from theme.models import UserQuota
-from theme.enums import QuotaStatus
 from theme.utils import get_quota_data
 from django_irods.icommands import SessionException
 from django_irods.storage import IrodsStorage
@@ -127,7 +126,7 @@ def update_quota_usage(username):
     username: the name of the user that needs to update quota usage for.
     :return: raise ValidationError if quota cannot be updated.
     """
-    from hs_core.tasks import send_user_notification_at_quota_grace_start
+    from hs_core.tasks import set_user_quota_in_userzone, send_user_notification_at_quota_grace_start
     hs_internal_zone = "hydroshare"
     uq = UserQuota.objects.filter(user__username=username, zone=hs_internal_zone).first()
     if uq is None:
@@ -139,21 +138,23 @@ def update_quota_usage(username):
     uz, dz = get_quota_usage_from_irods(username, raise_on_error=False)
 
     # subtract out published resources from the datazone
-    quota_data = get_quota_data(uq)
-    qmsg = quota_data["qmsg"]
+    original_quota_data = get_quota_data(uq)
+    qmsg = original_quota_data["qmsg"]
     published_percent = qmsg.published_resource_percent
     user = User.objects.get(username=username)
     published_size = get_storage_usage(user, flag="published")
     dz -= published_size * (1 - published_percent)
     uq.update_used_value(uz, dz)
 
-    if quota_data["enforce_quota"]:
+    if original_quota_data["enforce_quota"]:
+        updated_quota_data = get_quota_data(uq)
         # if enforcing quota, take steps to send messages
-        percent = quota_data["percent"]
+        percent = updated_quota_data["percent"]
+        remaining = updated_quota_data["remaining"]
+        set_user_quota_in_userzone.apply_async((user.pk, remaining))
         if percent < qmsg.soft_limit_percent:
             # No need for further action
             return
-        quota_status = quota_data["status"]
         today = datetime.date.today()
         if percent >= qmsg.soft_limit_percent:
             if percent >= 100 and percent < qmsg.hard_limit_percent:
@@ -170,19 +171,6 @@ def update_quota_usage(username):
                 # reset grace period now that the user is below quota soft limit
                 uq.grace_period_ends = None
                 uq.save()
-        # TODO #5329 implement toggle userzone
-        if quota_status == QuotaStatus.ENFORCEMENT:
-            # quota needs to be enforced
-            pass
-        elif quota_status == QuotaStatus.GRACE_PERIOD:
-            # quota is in grace period
-            pass
-        elif quota_status == QuotaStatus.WARNING:
-            # return quota warning message
-            pass
-        elif quota_status == QuotaStatus.INFO:
-            # return quota informational message
-            pass
 
 
 def res_has_web_reference(res):
diff --git a/hs_core/tasks.py b/hs_core/tasks.py
--- a/hs_core/tasks.py
+++ b/hs_core/tasks.py
@@ -556,6 +556,64 @@ def send_user_notification_at_quota_grace_start(user_pk):
             logger.debug("Failed to send quota warning email: " + str(ex))
 
 
+@celery_app.task(ignore_result=True, base=HydroshareTask)
+def set_user_quota_in_userzone(user_pk, quota=0):
+    """
+    Toggles the irods managed quota for a user in the iRODS user zone.
+
+    Args:
+        user_pk (int): The primary key of the user.
+        quota (int): The quota size in bytes.
+
+    Returns:
+        str: JSON-encoded string containing the success message if the operation is successful,
+             otherwise an error message.
+
+    Raises:
+        Exception: If an error occurs while toggling the upload permission.
+    """
+    from hs_core.views.utils import run_ssh_command
+    try:
+        user = User.objects.get(pk=user_pk)
+        script = settings.LINUX_ADMIN_USER_SET_QUOTA_IN_USER_ZONE_CMD
+        exec_cmd = "{0} {1} {2}".format(
+            script,
+            user.username,
+            quota,
+        )
+        output = run_ssh_command(
+            host=settings.HS_USER_ZONE_HOST,
+            uname=settings.LINUX_ADMIN_USER_FOR_HS_USER_ZONE,
+            pwd=settings.LINUX_ADMIN_USER_PWD_FOR_HS_USER_ZONE,
+            exec_cmd=exec_cmd,
+        )
+        for out_str in output:
+            if "ERROR:" in out_str.upper():
+                # there is an error from icommand run, report the error
+                return json.dumps(
+                    {
+                        "error": "iRODS server failed to set quota for this iRODS account {0}. "
+                        "If this issue persists, please notify help@cuahsi.org.".format(
+                            user.username
+                        )
+                    },
+                )
+
+        message = f"iRODS quota for user {user.username} was set successfully"
+        return json.dumps(
+            {
+                "success": message,
+            },
+        )
+    except Exception as ex:
+        return json.dumps(
+            {
+                "error": str(ex)
+                + " - iRODS server failed to set quota in userzone for user {user.username}."
+            },
+        )
+
+
 @celery_app.task(ignore_result=True, base=HydroshareTask)
 def check_geoserver_registrations(resources, run_async=True):
     # Check to ensure resources have updated web services registrations
diff --git a/theme/models.py b/theme/models.py
--- a/theme/models.py
+++ b/theme/models.py
@@ -283,6 +283,11 @@ class Meta:
     def used_percent(self):
         return self.used_value * 100.0 / self.allocated_value
 
+    @property
+    def remaining(self):
+        delta = self.allocated_value - self.used_value
+        return delta if delta > 0 else 0
+
     @property
     def used_value(self):
         uz, dz = self.get_used_value_by_zone(refresh_from_irods=False)
@@ -523,6 +528,7 @@ def update_user_quota_on_quota_request(sender, instance, **kwargs):
     """
     Increment the allocated_value for a UserQuota object uppon approval of a QuotaRequest
     """
+    from hs_core.tasks import set_user_quota_in_userzone
     if kwargs.get('created'):
         # it is a new QuotaRequest instance, no need to check further
         return
@@ -538,6 +544,7 @@ def update_user_quota_on_quota_request(sender, instance, **kwargs):
 
         qr.quota.save()
         notify_user_of_quota_action(qr)
+        set_user_quota_in_userzone.apply_async((qr.request_from.pk, qr.quota.remaining,))
     except QuotaRequest.DoesNotExist:
         logger.warning(
             f"QuotaRequest for {instance.pk} does not exist when trying to update it"
@@ -547,8 +554,9 @@ def update_user_quota_on_quota_request(sender, instance, **kwargs):
 @receiver(models.signals.pre_save, sender=UserQuota)
 def reset_grace_period_on_allocation_change(sender, instance, **kwargs):
     """
-    Reset the pending UserQuota grace perioud when the allocated_value is modified in the UserQuota
+    Reset the pending UserQuota grace period when the allocated_value is modified in the UserQuota
     """
+    from hs_core.tasks import set_user_quota_in_userzone
     if instance.id is None:  # new object will be created
         pass
     else:
@@ -556,6 +564,7 @@ def reset_grace_period_on_allocation_change(sender, instance, **kwargs):
         if previous.allocated_value != instance.allocated_value:
             # allocated_value is being updated
             instance.grace_period_ends = None
+            set_user_quota_in_userzone.apply_async((instance.user.pk, instance.remaining,))
 
 
 @receiver(models.signals.pre_save, sender=QuotaMessage)


Test patch:
diff --git a/hs_core/tests/api/native/test_irods_user_zone_federation.py b/hs_core/tests/api/native/test_irods_user_zone_federation.py
--- a/hs_core/tests/api/native/test_irods_user_zone_federation.py
+++ b/hs_core/tests/api/native/test_irods_user_zone_federation.py
@@ -6,7 +6,7 @@
 
 from hs_core.models import BaseResource
 from hs_core import hydroshare
-from hs_core.tasks import replicate_resource_bag_to_user_zone_task
+from hs_core.tasks import replicate_resource_bag_to_user_zone_task, set_user_quota_in_userzone
 from hs_core.testing import TestCaseCommonUtilities
 from hs_core.hydroshare.resource import update_quota_usage
 
@@ -161,3 +161,39 @@ def test_quota_update_in_fed_zones(self):
 
         # delete test resources
         hydroshare.resource.delete_resource(res.short_id)
+
+    def test_set_user_quota_in_userzone(self):
+        """
+        Test case for setting user quota in user zone.
+
+        This test case verifies that setting the user quota in the user zone works as expected.
+        It creates a test resource in the HydroShare zone, adds files to the resource, and then
+        replicates the resource bag to the user zone. It then attempts to replicate the resource
+        bag again, but with a user quota that is not adequate. The test expects an exception to be
+        raised in this case.
+        """
+        super(TestUserZoneIRODSFederation, self).assert_federated_irods_available()
+
+        res = hydroshare.resource.create_resource(
+            resource_type='CompositeResource',
+            owner=self.user,
+            title='My Test Resource in HydroShare Zone'
+        )
+        fed_test_file1_full_path = '/{zone}/home/testuser/{fname}'.format(
+            zone=settings.HS_USER_IRODS_ZONE, fname=self.file_one)
+        fed_file2_full_path = '/{zone}/home/testuser/{fname}'.format(
+            zone=settings.HS_USER_IRODS_ZONE, fname=self.file_two)
+        hydroshare.add_resource_files(
+            res.short_id,
+            source_names=[fed_test_file1_full_path, fed_file2_full_path])
+
+        # ensure replicate resource bag to user zone is successful
+        res.setAVU('bag_modified', 'false')
+        replicate_resource_bag_to_user_zone_task(res.short_id, self.user.username)
+
+        # test setting user quota in user zone
+        result = set_user_quota_in_userzone(self.user.pk, 1)
+        self.assertTrue('success' in result, msg='setting user quota in user zone failed')
+        # This should raise an exception as the user quota is not adequate to replicate the resource
+        with self.assertRaises(Exception):
+            replicate_resource_bag_to_user_zone_task(res.short_id, self.user.username)