[PATCH v2] crypto: zstd - fix segmented acomp streaming paths

[Wesley Atwell <atwellwea@gmail.com>]

The zstd acomp implementation does not correctly handle segmented
source and destination walks.

The compression path advances the destination walk by the full
segment length rather than the bytes actually produced, and it only
calls zstd_end_stream() once even though the streaming API requires it
to be called until it returns 0. With segmented destinations this can
leave buffered output behind and misaccount the walk progress.

The decompression path has the same destination accounting issue, and
it stops when the source walk is exhausted even if
zstd_decompress_stream() has not yet reported that the frame is fully
decoded and flushed. That can report success too early for segmented
requests and incomplete frames.

Fix both streaming paths by advancing destination segments by actual
output bytes, refilling destination segments as needed, draining
zstd_end_stream() until completion, and continuing to flush buffered
decompression output after the source walk is exhausted. Return
-EINVAL if decompression cannot finish once the input has been fully
consumed.

Fixes: f5ad93ffb541 ("crypto: zstd - convert to acomp")
Assisted-by: Codex:GPT-5
Signed-off-by: Wesley Atwell <atwellwea@gmail.com>
---
Changes in v2:
- always finalize acomp walk mappings in the direct one-shot paths
- add mapped src/dst cleanup on streaming error exits
- reacquire a destination segment in decompression before resuming after a
  full output-chunk completion

Local validation:
- built bzImage with CONFIG_CRYPTO_SELFTESTS=y and
  CONFIG_CRYPTO_SELFTESTS_FULL=y
- exercised segmented zstd acomp requests using temporary local testmgr
  scaffolding
- booted under virtme and verified zstd-generic selftest passed in
  /proc/crypto

diff --git a/crypto/zstd.c b/crypto/zstd.c
index 556f5d2bdd5f..3fb04acf6b7f 100644
--- a/crypto/zstd.c
+++ b/crypto/zstd.c
@@ -94,18 +94,32 @@ static int zstd_compress_one(struct acomp_req *req, struct zstd_ctx *ctx,
 	return 0;
 }
 
+static int zstd_acomp_next_dst(struct acomp_walk *walk, zstd_out_buffer *outbuf)
+{
+	unsigned int dcur = acomp_walk_next_dst(walk);
+
+	if (!dcur)
+		return -ENOSPC;
+
+	outbuf->pos = 0;
+	outbuf->dst = walk->dst.virt.addr;
+	outbuf->size = dcur;
+
+	return 0;
+}
+
 static int zstd_compress(struct acomp_req *req)
 {
 	struct crypto_acomp_stream *s;
-	unsigned int pos, scur, dcur;
+	unsigned int scur;
 	unsigned int total_out = 0;
-	bool data_available = true;
+	bool src_mapped = false;
+	bool dst_mapped = false;
 	zstd_out_buffer outbuf;
 	struct acomp_walk walk;
 	zstd_in_buffer inbuf;
 	struct zstd_ctx *ctx;
-	size_t pending_bytes;
-	size_t num_bytes;
+	size_t remaining;
 	int ret;
 
 	s = crypto_acomp_lock_stream_bh(&zstd_streams);
@@ -121,60 +135,102 @@ static int zstd_compress(struct acomp_req *req)
 		goto out;
 	}
 
-	do {
-		dcur = acomp_walk_next_dst(&walk);
-		if (!dcur) {
-			ret = -ENOSPC;
+	ret = zstd_acomp_next_dst(&walk, &outbuf);
+	if (ret)
+		goto out;
+	dst_mapped = true;
+
+	for (;;) {
+		scur = acomp_walk_next_src(&walk);
+		src_mapped = scur;
+		if (outbuf.size == req->dlen && scur == req->slen) {
+			ret = zstd_compress_one(req, ctx, walk.src.virt.addr,
+						walk.dst.virt.addr, &total_out);
+			acomp_walk_done_src(&walk, scur);
+			src_mapped = false;
+			acomp_walk_done_dst(&walk, ret ? outbuf.size : total_out);
+			dst_mapped = false;
 			goto out;
 		}
 
-		outbuf.pos = 0;
-		outbuf.dst = (u8 *)walk.dst.virt.addr;
-		outbuf.size = dcur;
+		if (!scur)
+			break;
+
+		inbuf.pos = 0;
+		inbuf.src = walk.src.virt.addr;
+		inbuf.size = scur;
 
 		do {
-			scur = acomp_walk_next_src(&walk);
-			if (dcur == req->dlen && scur == req->slen) {
-				ret = zstd_compress_one(req, ctx, walk.src.virt.addr,
-							walk.dst.virt.addr, &total_out);
-				acomp_walk_done_src(&walk, scur);
-				acomp_walk_done_dst(&walk, dcur);
-				goto out;
+			remaining = zstd_compress_stream(ctx->cctx, &outbuf, &inbuf);
+			if (zstd_is_error(remaining)) {
+				ret = -EIO;
+				goto out_free_walk;
 			}
 
-			if (scur) {
-				inbuf.pos = 0;
-				inbuf.src = walk.src.virt.addr;
-				inbuf.size = scur;
-			} else {
-				data_available = false;
-				break;
+			if (outbuf.pos != outbuf.size) {
+				if (inbuf.pos == inbuf.size)
+					break;
+				continue;
 			}
 
-			num_bytes = zstd_compress_stream(ctx->cctx, &outbuf, &inbuf);
-			if (ZSTD_isError(num_bytes)) {
-				ret = -EIO;
-				goto out;
+			total_out += outbuf.pos;
+			acomp_walk_done_dst(&walk, outbuf.pos);
+			dst_mapped = false;
+
+			ret = zstd_acomp_next_dst(&walk, &outbuf);
+			if (ret)
+				goto out_free_walk;
+			dst_mapped = true;
+		} while (inbuf.pos != inbuf.size);
+
+		acomp_walk_done_src(&walk, inbuf.pos);
+		src_mapped = false;
+	}
+
+	for (;;) {
+		remaining = zstd_end_stream(ctx->cctx, &outbuf);
+		if (zstd_is_error(remaining)) {
+			ret = -EIO;
+			goto out_free_walk;
+		}
+
+		if (outbuf.pos == outbuf.size) {
+			total_out += outbuf.pos;
+			acomp_walk_done_dst(&walk, outbuf.pos);
+			dst_mapped = false;
+
+			if (!remaining) {
+				outbuf.pos = 0;
+				break;
 			}
 
-			pending_bytes = zstd_flush_stream(ctx->cctx, &outbuf);
-			if (ZSTD_isError(pending_bytes)) {
-				ret = -EIO;
+			ret = zstd_acomp_next_dst(&walk, &outbuf);
+			if (ret)
 				goto out;
-			}
-			acomp_walk_done_src(&walk, inbuf.pos);
-		} while (dcur != outbuf.pos);
+			dst_mapped = true;
+
+			continue;
+		}
+
+		if (!remaining)
+			break;
+	}
 
+	if (outbuf.pos) {
 		total_out += outbuf.pos;
-		acomp_walk_done_dst(&walk, dcur);
-	} while (data_available);
+		acomp_walk_done_dst(&walk, outbuf.pos);
+		dst_mapped = false;
+	}
 
-	pos = outbuf.pos;
-	num_bytes = zstd_end_stream(ctx->cctx, &outbuf);
-	if (ZSTD_isError(num_bytes))
-		ret = -EIO;
-	else
-		total_out += (outbuf.pos - pos);
+out_free_walk:
+	if (src_mapped)
+		acomp_walk_done_src(&walk, inbuf.pos);
+	if (dst_mapped)
+		/*
+		 * The streaming helpers can fail after dirtying part of the
+		 * current destination chunk while leaving outbuf.pos stale.
+		 */
+		acomp_walk_done_dst(&walk, ret ? outbuf.size : outbuf.pos);
 
 out:
 	if (ret)
@@ -209,12 +265,14 @@ static int zstd_decompress(struct acomp_req *req)
 {
 	struct crypto_acomp_stream *s;
 	unsigned int total_out = 0;
-	unsigned int scur, dcur;
+	unsigned int scur;
+	bool src_mapped = false;
+	bool dst_mapped = false;
 	zstd_out_buffer outbuf;
 	struct acomp_walk walk;
 	zstd_in_buffer inbuf;
 	struct zstd_ctx *ctx;
-	size_t pending_bytes;
+	size_t remaining = 1;
 	int ret;
 
 	s = crypto_acomp_lock_stream_bh(&zstd_streams);
@@ -230,48 +288,128 @@ static int zstd_decompress(struct acomp_req *req)
 		goto out;
 	}
 
-	do {
+	ret = zstd_acomp_next_dst(&walk, &outbuf);
+	if (ret)
+		goto out;
+	dst_mapped = true;
+
+	for (;;) {
 		scur = acomp_walk_next_src(&walk);
-		if (scur) {
-			inbuf.pos = 0;
-			inbuf.size = scur;
-			inbuf.src = walk.src.virt.addr;
-		} else {
+		src_mapped = scur;
+		if (!scur)
 			break;
+
+		inbuf.pos = 0;
+		inbuf.size = scur;
+		inbuf.src = walk.src.virt.addr;
+
+		if (!dst_mapped) {
+			ret = zstd_acomp_next_dst(&walk, &outbuf);
+			if (ret)
+				goto out_free_walk;
+			dst_mapped = true;
+		}
+
+		if (outbuf.size == req->dlen && scur == req->slen) {
+			ret = zstd_decompress_one(req, ctx, walk.src.virt.addr,
+						  walk.dst.virt.addr, &total_out);
+			acomp_walk_done_src(&walk, scur);
+			src_mapped = false;
+			acomp_walk_done_dst(&walk, ret ? outbuf.size : total_out);
+			dst_mapped = false;
+			goto out;
 		}
 
 		do {
-			dcur = acomp_walk_next_dst(&walk);
-			if (dcur == req->dlen && scur == req->slen) {
-				ret = zstd_decompress_one(req, ctx, walk.src.virt.addr,
-							  walk.dst.virt.addr, &total_out);
-				acomp_walk_done_dst(&walk, dcur);
-				acomp_walk_done_src(&walk, scur);
-				goto out;
+			remaining = zstd_decompress_stream(ctx->dctx, &outbuf,
+							   &inbuf);
+			if (zstd_is_error(remaining)) {
+				ret = -EIO;
+				goto out_free_walk;
 			}
 
-			if (!dcur) {
-				ret = -ENOSPC;
-				goto out;
+			if (outbuf.pos != outbuf.size) {
+				if (inbuf.pos == inbuf.size)
+					break;
+				continue;
 			}
 
-			outbuf.pos = 0;
-			outbuf.dst = (u8 *)walk.dst.virt.addr;
-			outbuf.size = dcur;
+			total_out += outbuf.pos;
+			acomp_walk_done_dst(&walk, outbuf.pos);
+			dst_mapped = false;
 
-			pending_bytes = zstd_decompress_stream(ctx->dctx, &outbuf, &inbuf);
-			if (ZSTD_isError(pending_bytes)) {
-				ret = -EIO;
-				goto out;
+			if (!remaining) {
+				outbuf.pos = 0;
+				break;
 			}
 
-			total_out += outbuf.pos;
+			ret = zstd_acomp_next_dst(&walk, &outbuf);
+			if (ret)
+				goto out_free_walk;
+			dst_mapped = true;
+		} while (inbuf.pos != inbuf.size);
 
-			acomp_walk_done_dst(&walk, outbuf.pos);
-		} while (inbuf.pos != scur);
+		acomp_walk_done_src(&walk, inbuf.pos);
+		src_mapped = false;
+	}
+
+	inbuf.pos = 0;
+	inbuf.size = 0;
+	inbuf.src = NULL;
+
+	/* Drain any buffered output after the source walk is exhausted. */
+	while (remaining) {
+		size_t pos = outbuf.pos;
+
+		remaining = zstd_decompress_stream(ctx->dctx, &outbuf, &inbuf);
+		if (zstd_is_error(remaining)) {
+			ret = -EIO;
+			goto out_free_walk;
+		}
+
+		if (outbuf.pos == pos) {
+			ret = -EINVAL;
+			goto out_free_walk;
+		}
+
+		if (outbuf.pos != outbuf.size) {
+			if (remaining) {
+				ret = -EINVAL;
+				goto out_free_walk;
+			}
+			break;
+		}
+
+		total_out += outbuf.pos;
+		acomp_walk_done_dst(&walk, outbuf.pos);
+		dst_mapped = false;
+
+		if (!remaining) {
+			outbuf.pos = 0;
+			break;
+		}
+
+		ret = zstd_acomp_next_dst(&walk, &outbuf);
+		if (ret)
+			goto out;
+		dst_mapped = true;
+	}
+
+	if (outbuf.pos) {
+		total_out += outbuf.pos;
+		acomp_walk_done_dst(&walk, outbuf.pos);
+		dst_mapped = false;
+	}
 
-		acomp_walk_done_src(&walk, scur);
-	} while (ret == 0);
+out_free_walk:
+	if (src_mapped)
+		acomp_walk_done_src(&walk, inbuf.pos);
+	if (dst_mapped)
+		/*
+		 * The streaming helpers can fail after dirtying part of the
+		 * current destination chunk while leaving outbuf.pos stale.
+		 */
+		acomp_walk_done_dst(&walk, ret ? outbuf.size : outbuf.pos);
 
 out:
 	if (ret)
base-commit: 1f318b96cc84d7c2ab792fcc0bfd42a7ca890681
-- 
2.34.1