uploads: Rewrite the whole thing

[pbs.git] / src / buildservice / uploads.py

#!/usr/bin/python3

import asyncio
import hashlib
import hmac
import logging
import os
import shutil

from . import base
from . import builders
from . import users
from .constants import *
from .decorators import *

# Setup logging
log = logging.getLogger("pbs.uploads")

MAX_BUFFER_SIZE = 1 * 1024 * 1024 # 1 MiB

supported_digest_algos = (
	"blake2b512",
)

class UnsupportedDigestException(ValueError):
	pass

class Uploads(base.Object):
	def _get_upload(self, query, *args):
		res = self.db.get(query, *args)

		if res:
			return Upload(self.backend, res.id, data=res)

	def _get_uploads(self, query, *args):
		res = self.db.query(query, *args)

		for row in res:
			yield Upload(self.backend, row.id, data=row)

	def __iter__(self):
		uploads = self._get_uploads("SELECT * FROM uploads \
			ORDER BY created_at DESC")

		return iter(uploads)

	def get_by_uuid(self, uuid):
		return self._get_upload("""
			SELECT
				*
			FROM
				uploads
			WHERE
				uuid = %s
			AND
				expires_at > CURRENT_TIMESTAMP
			""", uuid,
		)

	async def create(self, filename, size, digest_algo, digest, owner=None):
		builder = None
		user = None

		# Check if the digest algorithm is supported
		if not digest_algo in supported_digest_algos:
			raise UnsupportedDigestException(digest_algo)

		# Check if the digest is set
		elif not digest:
			raise ValueError("Empty digest")

		# Check uploader type
		if isinstance(owner, builders.Builder):
			builder = owner
		elif isinstance(owner, users.User):
			user = owner

		# Check quota for users
		if user:
			# This will raise an exception if the quota has been exceeded
			user.check_storage_quota(size)

		# Allocate a new temporary file
		upload = self._get_upload("""
			INSERT INTO
				uploads
			(
				filename,
				size,
				builder_id,
				user_id,
				digest_algo,
				digest
			)
			VALUES
			(
				%s,
				%s,
				%s,
				%s,
				%s,
				%s
			)
			RETURNING *""",
			filename,
			size,
			builder,
			user,
			digest_algo,
			digest,
		)

		# Return the newly created upload object
		return upload

	async def create_from_local(self, path, filename=None, **kwargs):
		"""
			Imports a file from the local filesystem
		"""
		# Collect attributes
		if filename is None:
			filename = os.path.basename(path)

		# Fetch the size
		size = os.path.getsize(path)

		# Create the new upload object
		upload = await self.create(filename=filename, size=size, **kwargs)

		# Import the data
		with open(path, "rb") as f:
			await upload.copyfrom(f)

		return upload

	async def cleanup(self):
		# Find all expired uploads
		uploads = self._get_uploads("""
			SELECT
				*
			FROM
				uploads
			WHERE
				expires_at <= CURRENT_TIMESTAMP
			ORDER BY
				created_at
		""")

		# Delete them all
		for upload in uploads:
			with self.db.transaction():
				await upload.delete()


class Upload(base.DataObject):
	table = "uploads"

	def __str__(self):
		return "%s" % self.uuid

	@property
	def uuid(self):
		return self.data.uuid

	@property
	def filename(self):
		return self.data.filename

	@property
	def path(self):
		return self.data.path

	@property
	def size(self):
		return self.data.size

	@property
	def digest_algo(self):
		return self.data.digest_algo

	@property
	def digest(self):
		return self.data.digest

	# Builder

	def get_builder(self):
		if self.data.builder_id:
			return self.backend.builders.get_by_id(self.data.builder_id)

	def set_builder(self, builder):
		self._set_attribute("builder_id", builder.id)

	builder = lazy_property(get_builder, set_builder)

	# User

	def get_user(self):
		if self.data.user_id:
			return self.backend.users.get_by_id(self.data.user_id)

	def set_user(self, user):
		self._set_attribute("user_id", user.id)

	user = lazy_property(get_user, set_user)

	def has_perm(self, who):
		"""
			Returns True if "who" has permissions to use this upload
		"""
		if self.builder == who or self.user == who:
			return True

		# No permission
		return False

	async def delete(self):
		log.info("Deleting upload %s (%s)" % (self, self.path))

		# Remove the uploaded data
		await self.backend.unlink(self.path)

		# Delete the upload from the database
		self.db.execute("DELETE FROM uploads WHERE id = %s", self.id)

	@property
	def created_at(self):
		return self.data.created_at

	@property
	def expires_at(self):
		return self.data.expires_at

	async def copyfrom(self, src):
		"""
			Copies the content of this upload from the source file descriptor
		"""
		# Reset the source file handle
		src.seek(0)

		# Setup the hash function
		h = hashlib.new(self.digest_algo)

		async with self.backend.tempfile(delete=False) as f:
			try:
				while True:
					buffer = src.read(1024 ** 2)
					if not buffer:
						break

					# Put the buffer into the hash function
					h.update(buffer)

					# Write the buffer to disk
					await f.write(buffer)

				# How many bytes did we write?
				received_size = await f.tell()

				# Get the digest
				computed_digest = h.digest()

				# Check if the filesize matches
				if not received_size == self.size:
					raise ValueError("File size differs")

				# Check that the digest matches
				if not hmac.compare_digest(computed_digest, self.digest):
					raise ValueError("Invalid digest")

			# If there has been any kind of exception, we want to delete the temporary file
			except Exception as e:
				await self.backend.unlink(f.name)

				raise e

		# Store the path
		self._set_attribute("path", f.name)

	async def copyinto(self, dst):
		"""
			Copies the content of this upload into the destination file descriptor.
		"""
		return await asyncio.to_thread(self._copyinto, dst)

	def _copyinto(self, dst):
		assert os.path.exists(self.path), "Upload has been deleted - %s" % self

		# Open the source file and copy it into the destination file descriptor
		with open(self.path, "rb") as src:
			shutil.copyfileobj(src, dst)