fix: apply quota with writeStream

Signed-off-by: Robin Appelman <robin@icewind.nl>

Author: icewind1991

lib/private/Files/Storage/Wrapper/Quota.php

@@ -11,6 +11,8 @@
 use OC\SystemConfig;
 use OCP\Files\Cache\ICacheEntry;
 use OCP\Files\FileInfo;
+use OCP\Files\GenericFileException;
+use OCP\Files\NotEnoughSpaceException;
 use OCP\Files\Storage\IStorage;
 
 class Quota extends Wrapper {
@@ -214,4 +216,31 @@ public function touch(string $path, ?int $mtime = null): bool {
 	public function enableQuota(bool $enabled): void {
 		$this->enabled = $enabled;
 	}
+
+	#[\Override]
+	public function writeStream(string $path, $stream, ?int $size = null): int {
+		if (!$this->hasQuota()) {
+			return parent::writeStream($path, $stream, $size);
+		}
+
+		$free = $this->free_space($path);
+		if ($this->shouldApplyQuota($path) && $free == 0) {
+			throw new NotEnoughSpaceException();
+		}
+
+		if ($size !== null) {
+			if ($size < $free) {
+				return parent::writeStream($path, $stream, $size);
+			} else {
+				throw new NotEnoughSpaceException();
+			}
+		} else {
+			// force fallback through `fopen` to handle the quota
+			try {
+				return parent::writeStreamFallback($path, $stream);
+			} catch (GenericFileException) {
+				throw new NotEnoughSpaceException();
+			}
+		}
+	}
 }

lib/private/Files/Storage/Wrapper/Wrapper.php

@@ -392,6 +392,13 @@ public function writeStream(string $path, $stream, ?int $size = null): int {
 			return $storage->writeStream($path, $stream, $size);
 		}
 
+		return $this->writeStreamFallback($path, $stream);
+	}
+
+	/**
+	 * @param resource $stream
+	 */
+	protected function writeStreamFallback(string $path, $stream): int {
 		$target = $this->fopen($path, 'w');
 		if ($target === false) {
 			throw new GenericFileException('Failed to open ' . $path);

tests/lib/Files/Storage/Wrapper/QuotaTest.php

@@ -235,4 +235,25 @@ public function testNoFopenQuotaZero(): void {
 		$fh = $instance->fopen('files/test.txt', 'w');
 		$this->assertFalse($fh);
 	}
+
+	public function testNoWriteStreamQuota(): void {
+		$instance = $this->getLimitedStorage(5.0);
+		$stream = fopen('php://temp', 'w+');
+		fwrite($stream, 'foo');
+		rewind($stream);
+		$instance->writeStream('files/test.txt', $stream);
+
+		$stream = fopen('php://temp', 'w+');
+		fwrite($stream, 'foobar');
+		rewind($stream);
+		$this->expectException(Files\NotEnoughSpaceException::class);
+		$instance->writeStream('files/test.txt', $stream);
+	}
+
+	public function testNoWriteStreamQuotaZero(): void {
+		$instance = $this->getLimitedStorage(0.0);
+		$stream = fopen('php://temp', 'w+');
+		$this->expectException(Files\NotEnoughSpaceException::class);
+		$instance->writeStream('files/test.txt', $stream);
+	}
 }