From 50533a78c412824f2548e85ef8b1cc960023a384 Mon Sep 17 00:00:00 2001
From: Michael Tremer <michael.tremer@ipfire.org>
Date: Tue, 28 Jun 2022 15:24:53 +0000
Subject: [PATCH] users: Implement quota checks

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>
---
 src/buildservice/users.py     | 53 +++++++++++++++++++++++++++++++++++
 src/database.sql              | 28 +++++++++++++++++-
 src/hub/handlers.py           |  8 +++++-
 src/templates/users/show.html | 24 ++++++++++++++++
 4 files changed, 111 insertions(+), 2 deletions(-)

diff --git a/src/buildservice/users.py b/src/buildservice/users.py
index 843a1f1b..229c05e2 100644
--- a/src/buildservice/users.py
+++ b/src/buildservice/users.py
@@ -369,6 +369,55 @@ class User(base.DataObject):
 		return self.backend.sessions._get_sessions("SELECT * FROM sessions \
 			WHERE user_id = %s AND valid_until >= NOW() ORDER BY created_at")
 
+	# Quota
+
+	def get_quota(self):
+		return self.data.quota
+
+	def set_quota(self, quota):
+		self._set_attribute("quota", quota)
+
+	quota = property(get_quota, set_quota)
+
+	def has_exceeded_quota(self, size=None):
+		"""
+			Returns True if this user has exceeded their quota
+		"""
+		# Skip quota check if this user has no quota
+		if not self.quota:
+			return
+
+		return self.disk_usage + (size or 0) >= self.quota
+
+	def check_quota(self, size=None):
+		"""
+			Determines the user's disk usage
+			and raises an exception when the user is over quota.
+		"""
+		# Raise QuotaExceededError if this user is over quota
+		if self.has_exceeded_quota(size=size):
+			raise QuotaExceededError
+
+	@lazy_property
+	def disk_usage(self):
+		"""
+			Returns the total disk usage of this user
+		"""
+		res = self.db.get("""
+			SELECT
+				disk_usage
+			FROM
+				user_disk_usages
+			WHERE
+				user_id = %s""",
+			self.id,
+		)
+
+		if res:
+			return res.disk_usage
+
+		return 0
+
 	# Custom repositories
 
 	@property
@@ -469,3 +518,7 @@ class UserEmail(base.DataObject):
 		logging.debug("Sending email address activation mail to %s" % self.email)
 
 		self.user.send_template("messages/users/email-activation", email=self)
+
+
+class QuotaExceededError(Exception):
+	pass
diff --git a/src/database.sql b/src/database.sql
index 498eac9d..f986b18a 100644
--- a/src/database.sql
+++ b/src/database.sql
@@ -1348,6 +1348,31 @@ ALTER TABLE public.uploads_id_seq OWNER TO pakfire;
 ALTER SEQUENCE public.uploads_id_seq OWNED BY public.uploads.id;
 
 
+--
+-- Name: user_disk_usages; Type: VIEW; Schema: public; Owner: pakfire
+--
+
+CREATE VIEW public.user_disk_usages AS
+ WITH objects AS (
+         SELECT uploads.user_id,
+            uploads.size
+           FROM public.uploads
+          WHERE (uploads.expires_at > CURRENT_TIMESTAMP)
+        UNION ALL
+         SELECT builds.owner_id,
+            packages.size
+           FROM (public.builds
+             LEFT JOIN public.packages ON ((builds.pkg_id = packages.id)))
+          WHERE ((builds.deleted IS FALSE) AND (builds.owner_id IS NOT NULL))
+        )
+ SELECT objects.user_id,
+    sum(objects.size) AS disk_usage
+   FROM objects
+  GROUP BY objects.user_id;
+
+
+ALTER TABLE public.user_disk_usages OWNER TO pakfire;
+
 --
 -- Name: users; Type: TABLE; Schema: public; Owner: pakfire
 --
@@ -1365,7 +1390,8 @@ CREATE TABLE public.users (
     ldap_dn text,
     password_recovery_code text,
     password_recovery_code_expires_at timestamp without time zone,
-    admin boolean DEFAULT false NOT NULL
+    admin boolean DEFAULT false NOT NULL,
+    quota bigint
 );
 
 
diff --git a/src/hub/handlers.py b/src/hub/handlers.py
index 38d2b481..b7d30164 100644
--- a/src/hub/handlers.py
+++ b/src/hub/handlers.py
@@ -149,7 +149,13 @@ class UploadHandler(BaseHandler):
 		# Fetch file size
 		self.filesize = self.get_argument_int("size")
 
-		# XXX check quota
+		# Check quota
+		if isinstance(self.current_user, users.User):
+			try:
+				self.current_user.check_quota(self.filesize)
+			except users.QuotaExceededError as e:
+				raise tornado.web.HTTPError(400,
+					"Quota exceeded for %s" % self.current_user) from e
 
 		# Count how many bytes have been received
 		self.bytes_read = 0
diff --git a/src/templates/users/show.html b/src/templates/users/show.html
index 495a2120..d1b0744c 100644
--- a/src/templates/users/show.html
+++ b/src/templates/users/show.html
@@ -48,4 +48,28 @@
 	<a class="success expanded button" href="/users/{{ user.name }}/repos/create">
 		{{ _("Create Repository") }}
 	</a>
+
+	<div class="callout">
+		<div class="grid-x grid-padding-x">
+			{% if user.quota %}
+				<div class="cell large-4 offset-8">
+					<h6>{{ _("Quota") }}</h6>
+
+					{% set p = min(user.disk_usage * 100 / user.quota, 100) %}
+
+					<div class="{% if user.has_exceeded_quota() %}alert{% end %} progress"
+							role="progressbar"
+							aria-valuenow="{{ "%.0f" % p }}"
+							aria-valuemin="0"
+							aria-valuemax="100">
+						<span class="progress-meter" style="width: {{ "%.0f%%" % p }}">
+							<span class="progress-meter-text">
+								{{ format_size(user.disk_usage) }}/{{ format_size(user.quota) }}
+							</span>
+						</span>
+					</div>
+				</div>
+			{% end %}
+		</div>
+	</div>
 {% end block %}
-- 
2.47.2