From a36666e2fe8d61a161daa1836ed84fb2e389dd67 Mon Sep 17 00:00:00 2001
From: Dan Smith <dansmith@redhat.com>
Date: Tue, 27 Apr 2021 07:09:06 -0700
Subject: [PATCH] Add image_count_total quota enforcement

This makes us enforce a quota on the total number of (non-deleted)
images owned by a user.

Partially-implements: blueprint glance-unified-quotas

Change-Id: I8af124d9307263cd8289d0701fb9a745d13b1d56
---
 glance/api/v2/images.py                   |  1 +
 glance/quota/keystone.py                  | 13 +++++++++
 glance/tests/functional/v2/test_images.py | 35 ++++++++++++++++++++---
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/glance/api/v2/images.py b/glance/api/v2/images.py
index d3f0a4f0bd..16a93ea89a 100644
--- a/glance/api/v2/images.py
+++ b/glance/api/v2/images.py
@@ -94,6 +94,7 @@ class ImagesController(object):
         image_factory = self.gateway.get_image_factory(req.context)
         image_repo = self.gateway.get_repo(req.context)
         try:
+            ks_quota.enforce_image_count_total(req.context, req.context.owner)
             image = image_factory.new_image(extra_properties=extra_properties,
                                             tags=tags, **image)
             image_repo.add(image)
diff --git a/glance/quota/keystone.py b/glance/quota/keystone.py
index a8d422b390..7f6cdd2304 100644
--- a/glance/quota/keystone.py
+++ b/glance/quota/keystone.py
@@ -29,6 +29,7 @@ limit.opts.register_opts(CONF)
 
 QUOTA_IMAGE_SIZE_TOTAL = 'image_size_total'
 QUOTA_IMAGE_STAGING_TOTAL = 'image_stage_total'
+QUOTA_IMAGE_COUNT_TOTAL = 'image_count_total'
 
 
 def _enforce_some(context, project_id, quota_value_fns, deltas):
@@ -112,3 +113,15 @@ def enforce_image_staging_total(context, project_id, delta=0):
         context, project_id, QUOTA_IMAGE_STAGING_TOTAL,
         lambda: db.user_get_staging_usage(context, project_id) // units.Mi,
         delta=delta)
+
+
+def enforce_image_count_total(context, project_id):
+    """Enforce the image_count_total quota.
+
+    This enforces the total count of non-deleted images owned by the
+    supplied project_id.
+    """
+    _enforce_one(
+        context, project_id, QUOTA_IMAGE_COUNT_TOTAL,
+        lambda: db.user_get_image_count(context, project_id),
+        delta=1)
diff --git a/glance/tests/functional/v2/test_images.py b/glance/tests/functional/v2/test_images.py
index da4963c168..eaaa9b4c5b 100644
--- a/glance/tests/functional/v2/test_images.py
+++ b/glance/tests/functional/v2/test_images.py
@@ -7058,7 +7058,8 @@ class TestKeystoneQuotas(functional.SynchronousAPIBase):
 
     def test_upload(self):
         # Set a quota of 5MiB
-        self.set_limit({'image_size_total': 5})
+        self.set_limit({'image_size_total': 5,
+                        'image_count_total': 10})
         self.start_server()
 
         # First upload of 3MiB is good
@@ -7081,7 +7082,8 @@ class TestKeystoneQuotas(functional.SynchronousAPIBase):
 
     def test_import(self):
         # Set a quota of 5MiB
-        self.set_limit({'image_size_total': 5})
+        self.set_limit({'image_size_total': 5,
+                        'image_count_total': 10})
         self.start_server()
 
         # First upload of 3MiB is good
@@ -7103,7 +7105,8 @@ class TestKeystoneQuotas(functional.SynchronousAPIBase):
 
     def test_import_would_go_over(self):
         # Set a quota limit of 5MiB
-        self.set_limit({'image_size_total': 5})
+        self.set_limit({'image_size_total': 5,
+                        'image_count_total': 10})
         self.start_server()
 
         # First upload of 3MiB is good
@@ -7142,6 +7145,7 @@ class TestKeystoneQuotas(functional.SynchronousAPIBase):
     def test_copy(self):
         # Set a size quota of 5MiB, with more staging quota than we need.
         self.set_limit({'image_size_total': 5,
+                        'image_count_total': 10,
                         'image_stage_total': 15})
         self.start_server()
 
@@ -7183,7 +7187,8 @@ class TestKeystoneQuotas(functional.SynchronousAPIBase):
     def test_stage(self):
         # Set a quota of 5MiB
         self.set_limit({'image_size_total': 15,
-                        'image_stage_total': 5})
+                        'image_stage_total': 5,
+                        'image_count_total': 10})
         self.start_server()
 
         # Stage 6MiB, which is allowed to complete, but leaves us over
@@ -7213,3 +7218,25 @@ class TestKeystoneQuotas(functional.SynchronousAPIBase):
         # Stage should now succeed because we have freed up quota
         self._create_and_stage(
             data_iter=test_utils.FakeData(6 * units.Mi))
+
+    def test_create(self):
+        # Set a quota of 2 images
+        self.set_limit({'image_size_total': 15,
+                        'image_count_total': 2})
+        self.start_server()
+
+        # Create one image
+        image_id = self._create().json['id']
+
+        # Create a second. This leaves us *at* quota
+        self._create()
+
+        # Attempt to create a third is rejected as OverLimit
+        resp = self._create()
+        self.assertEqual(413, resp.status_code)
+
+        # Delete one image, which should put us under quota
+        self.api_delete('/v2/images/%s' % image_id)
+
+        # Now we can create that third image
+        self._create()