From: Michael Tremer <michael.tremer@ipfire.org>
Date: Mon, 10 Feb 2025 10:50:31 +0000 (+0000)
Subject: registry: Add a simple in-memory cache for indexes
X-Git-Url: http://git.ipfire.org/gitweb/gitweb.cgi?a=commitdiff_plain;h=30b1aabb7ed375809c3e0b3422c582f8471ff528;p=pbs.git

registry: Add a simple in-memory cache for indexes

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>
---

diff --git a/src/web/registry.py b/src/web/registry.py
index a87b873c..b26bd93a 100644
--- a/src/web/registry.py
+++ b/src/web/registry.py
@@ -21,7 +21,7 @@
 # This file implements the OCI Registry API
 #   https://github.com/opencontainers/distribution-spec/blob/main/spec.md#api
 
-import itertools
+import hashlib
 import json
 import tornado.web
 
@@ -58,6 +58,15 @@ class BaseHandler(base.BaseHandler):
 	"""
 		A base handler for any registry stuff
 	"""
+	def hash(self, data):
+		if isinstance(data, str):
+			data = data.encode()
+
+		h = hashlib.new("sha256")
+		h.update(data)
+
+		return "sha256:%s" % h.hexdigest()
+
 	async def get_blob(self, distro, digest):
 		"""
 			Returns a blob if we have it.
@@ -143,6 +152,11 @@ class IndexHandler(BaseHandler):
 		pass
 
 
+# Docker is pulling the same content more than once (god knows why). After the
+# ManifestLabelHandler has responded, the client will ask for the same file again
+# by its digest. Therefore we simply cache them and will send them again upon request.
+CACHE = {}
+
 class ManifestLabelHandler(BaseHandler):
 	async def head(self, *args, **kwargs):
 		return await self.get(*args, **kwargs, send_body=False)
@@ -188,6 +202,9 @@ class ManifestLabelHandler(BaseHandler):
 		self.set_header("Content-Type", "application/vnd.oci.image.index.v1+json")
 		self.set_header("Content-Length", len(index))
 
+		# Hash and cache
+		CACHE[self.hash(index)] = index
+
 		# Send the response
 		if send_body:
 			self.finish(index)
@@ -198,6 +215,10 @@ class ManifestHandler(BaseHandler):
 		return await self.get(*args, **kwargs, send_body=False)
 
 	async def get(self, distro_slug, digest, send_body=True):
+		# If we have this item in the cache, we serve it straight away
+		if digest in CACHE:
+			return self._get_cached_manifest(digest, send_body=send_body)
+
 		# Fetch the distribution
 		distro = await self.backend.distros.get_by_slug(distro_slug)
 		if not distro:
@@ -215,6 +236,19 @@ class ManifestHandler(BaseHandler):
 		if send_body:
 			await self.stream_blob(blob)
 
+	def _get_cached_manifest(self, digest, send_body=True):
+		try:
+			manifest = CACHE[digest]
+		except KeyError:
+			raise ManifestUnknownError
+
+		# Set Content-Type
+		self.set_header("Content-Type", "application/vnd.oci.image.index.v1+json")
+		self.set_header("Content-Length", len(index))
+
+		if send_body:
+			self.finish(manifest)
+
 
 class BlobHandler(BaseHandler):
 	async def head(self, *args, **kwargs):