From 9703b1da41837c48b527ee791440defa3026cdf1 Mon Sep 17 00:00:00 2001
From: =?utf8?q?John=20Molakvo=C3=A6?= <skjnldsv@protonmail.com>
Date: Wed, 10 Jan 2024 13:11:06 +0100
Subject: [PATCH] fix(quota): automatically detect chunk quota with
 `OC-Total-Length` header
MIME-Version: 1.0
Content-Type: text/plain; charset=utf8
Content-Transfer-Encoding: 8bit

Signed-off-by: John Molakvoæ <skjnldsv@protonmail.com>
---
 apps/dav/lib/Connector/Sabre/QuotaPlugin.php | 83 ++++++++++++--------
 1 file changed, 50 insertions(+), 33 deletions(-)

diff --git a/apps/dav/lib/Connector/Sabre/QuotaPlugin.php b/apps/dav/lib/Connector/Sabre/QuotaPlugin.php
index 1eed3d0d778..687b05e86cb 100644
--- a/apps/dav/lib/Connector/Sabre/QuotaPlugin.php
+++ b/apps/dav/lib/Connector/Sabre/QuotaPlugin.php
@@ -31,6 +31,7 @@
 namespace OCA\DAV\Connector\Sabre;
 
 use OCA\DAV\Upload\FutureFile;
+use OCA\DAV\Upload\UploadFolder;
 use OCP\Files\StorageNotAvailableException;
 use Sabre\DAV\Exception\InsufficientStorage;
 use Sabre\DAV\Exception\ServiceUnavailable;
@@ -90,6 +91,19 @@ class QuotaPlugin extends \Sabre\DAV\ServerPlugin {
 	 * @param bool $modified modified
 	 */
 	public function beforeCreateFile($uri, $data, INode $parent, $modified) {
+		$request = $this->server->httpRequest;
+		if ($parent instanceof UploadFolder && $request->getHeader('Destination')) {
+			// If chunked upload and Total-Length header is set, use that
+			// value for quota check. This allows us to also check quota while
+			// uploading chunks and not only when the file is assembled.
+			$length = $request->getHeader('OC-Total-Length');
+			$destinationPath = $this->server->calculateUri($request->getHeader('Destination'));
+			$quotaPath = $this->getPathForDestination($destinationPath);
+			if ($quotaPath && is_numeric($length)) {
+				return $this->checkQuota($quotaPath, (int)$length);
+			}
+		}
+
 		if (!$parent instanceof Node) {
 			return;
 		}
@@ -114,29 +128,20 @@ class QuotaPlugin extends \Sabre\DAV\ServerPlugin {
 	}
 
 	/**
-	 * Check if we're moving a Futurefile in which case we need to check
+	 * Check if we're moving a FutureFile in which case we need to check
 	 * the quota on the target destination.
-	 *
-	 * @param string $source source path
-	 * @param string $destination destination path
 	 */
-	public function beforeMove($source, $destination) {
-		$sourceNode = $this->server->tree->getNodeForPath($source);
+	public function beforeMove(string $sourcePath, string $destinationPath): bool {
+		$sourceNode = $this->server->tree->getNodeForPath($sourcePath);
 		if (!$sourceNode instanceof FutureFile) {
-			return;
+			return true;
 		}
 
-		// get target node for proper path conversion
-		if ($this->server->tree->nodeExists($destination)) {
-			$destinationNode = $this->server->tree->getNodeForPath($destination);
-			$path = $destinationNode->getPath();
-		} else {
-			$parent = dirname($destination);
-			if ($parent === '.') {
-				$parent = '';
-			}
-			$parentNode = $this->server->tree->getNodeForPath($parent);
-			$path = $parentNode->getPath();
+		try {
+			// The final path is not known yet, we check the quota on the parent
+			$path = $this->getPathForDestination($destinationPath);
+		} catch (\Exception $e) {
+			return true;
 		}
 
 		return $this->checkQuota($path, $sourceNode->getSize());
@@ -151,26 +156,36 @@ class QuotaPlugin extends \Sabre\DAV\ServerPlugin {
 			return true;
 		}
 
+		try {
+			$path = $this->getPathForDestination($destinationPath);
+		} catch (\Exception $e) {
+			return true;
+		}
+
+		return $this->checkQuota($path, $sourceNode->getSize());
+	}
+
+	private function getPathForDestination(string $destinationPath): string {
 		// get target node for proper path conversion
 		if ($this->server->tree->nodeExists($destinationPath)) {
 			$destinationNode = $this->server->tree->getNodeForPath($destinationPath);
 			if (!$destinationNode instanceof Node) {
-				return true;
-			}
-			$path = $destinationNode->getPath();
-		} else {
-			$parent = dirname($destinationPath);
-			if ($parent === '.') {
-				$parent = '';
-			}
-			$parentNode = $this->server->tree->getNodeForPath($parent);
-			if (!$parentNode instanceof Node) {
-				return true;
+				throw new \Exception('Invalid destination node');
 			}
-			$path = $parentNode->getPath();
+			return $destinationNode->getPath();
 		}
 
-		return $this->checkQuota($path, $sourceNode->getSize());
+		$parent = dirname($destinationPath);
+		if ($parent === '.') {
+			$parent = '';
+		}
+
+		$parentNode = $this->server->tree->getNodeForPath($parent);
+		if (!$parentNode instanceof Node) {
+			throw new \Exception('Invalid destination node');
+		}
+
+		return $parentNode->getPath();
 	}
 
 
@@ -182,7 +197,7 @@ class QuotaPlugin extends \Sabre\DAV\ServerPlugin {
 	 * @throws InsufficientStorage
 	 * @return bool
 	 */
-	public function checkQuota($path, $length = null) {
+	public function checkQuota(string $path, $length = null) {
 		if ($length === null) {
 			$length = $this->getLength();
 		}
@@ -194,7 +209,7 @@ class QuotaPlugin extends \Sabre\DAV\ServerPlugin {
 			}
 			$req = $this->server->httpRequest;
 
-			// If chunked upload
+			// If LEGACY chunked upload
 			if ($req->getHeader('OC-Chunked')) {
 				$info = \OC_FileChunking::decodeName($newName);
 				$chunkHandler = $this->getFileChunking($info);
@@ -210,12 +225,14 @@ class QuotaPlugin extends \Sabre\DAV\ServerPlugin {
 
 			$freeSpace = $this->getFreeSpace($path);
 			if ($freeSpace >= 0 && $length > $freeSpace) {
+				// If LEGACY chunked upload, clean up
 				if (isset($chunkHandler)) {
 					$chunkHandler->cleanup();
 				}
 				throw new InsufficientStorage("Insufficient space in $path, $length required, $freeSpace available");
 			}
 		}
+
 		return true;
 	}
 
-- 
2.39.5