perf record: Return the written size from process_comp_header()

process_comp_header() is called from zstd_compress_stream_to_records()
twice per record: once with data_size == 0 to write the record header,
and once with the payload size to finalize it. It returns the increment
it was passed, and the loop separately decides whether a record still
fits by comparing the remaining 'dst_size' against the header size.

With the fit check split from the code that writes the record,
process_comp_header() cannot reject a record on its own, so any bytes it
writes into 'dst' have to be bounds-checked by the caller instead of
where they are produced.

Pass the space left in 'dst' to process_comp_header(), let it return the
number of bytes written or -1 when the header does not fit, and account
the compressed payload in the loop.

No functional change intended.

Signed-off-by: Dmitry Ilvokhin <d@ilvokhin.com>
Signed-off-by: Namhyung Kim <namhyung@kernel.org>

3 files changed, 31 insertions, 17 deletions

diff --git a/tools/perf/builtin-record.c b/tools/perf/builtin-record.c
index d1276382b77a..294bdd4b8d00 100644
--- a/tools/perf/builtin-record.c
+++ b/tools/perf/builtin-record.c
@@ -1592,16 +1592,25 @@ static void record__adjust_affinity(struct record *rec, struct mmap *map)
 	}
 }
 
-static size_t process_comp_header(void *record, size_t increment)
+/*
+ * Called once with data_size == 0 to start a record, then once with
+ * data_size == compressed payload size to finalize.
+ * Returns the bytes written, or -1 if it won't fit.
+ */
+static ssize_t process_comp_header(void *record, size_t dst_size,
+				   size_t data_size)
 {
 	struct perf_record_compressed2 *event = record;
 	size_t size = sizeof(*event);
 
-	if (increment) {
-		event->header.size += increment;
-		return increment;
+	if (data_size) {
+		event->header.size += data_size;
+		return 0;
 	}
 
+	if (size > dst_size)
+		return -1;
+
 	event->header.type = PERF_RECORD_COMPRESSED2;
 	event->header.size = size;
 
diff --git a/tools/perf/util/compress.h b/tools/perf/util/compress.h
index 6cfecfca16f2..ec6c38129e24 100644
--- a/tools/perf/util/compress.h
+++ b/tools/perf/util/compress.h
@@ -54,7 +54,8 @@ int zstd_fini(struct zstd_data *data);
 
 ssize_t zstd_compress_stream_to_records(struct zstd_data *data, void *dst, size_t dst_size,
 				       void *src, size_t src_size, size_t max_record_size,
-				       size_t process_header(void *record, size_t increment));
+				       ssize_t process_header(void *record, size_t dst_size,
+							      size_t data_size));
 
 size_t zstd_decompress_stream(struct zstd_data *data, void *src, size_t src_size,
 			      void *dst, size_t dst_size);
@@ -75,7 +76,8 @@ ssize_t zstd_compress_stream_to_records(struct zstd_data *data __maybe_unused,
 				       void *dst __maybe_unused, size_t dst_size __maybe_unused,
 				       void *src __maybe_unused, size_t src_size __maybe_unused,
 				       size_t max_record_size __maybe_unused,
-				       size_t process_header(void *record, size_t increment) __maybe_unused)
+				       ssize_t process_header(void *record, size_t dst_size,
+							      size_t data_size) __maybe_unused)
 {
 	return 0;
 }
diff --git a/tools/perf/util/zstd.c b/tools/perf/util/zstd.c
index 21a0eb58597c..d98014902f01 100644
--- a/tools/perf/util/zstd.c
+++ b/tools/perf/util/zstd.c
@@ -31,9 +31,11 @@ int zstd_fini(struct zstd_data *data)
 
 ssize_t zstd_compress_stream_to_records(struct zstd_data *data, void *dst, size_t dst_size,
 				       void *src, size_t src_size, size_t max_record_size,
-				       size_t process_header(void *record, size_t increment))
+				       ssize_t process_header(void *record, size_t dst_size,
+							      size_t data_size))
 {
-	size_t ret, size, compressed = 0;
+	size_t ret, compressed = 0;
+	ssize_t size;
 	ZSTD_inBuffer input = { src, src_size, 0 };
 	ZSTD_outBuffer output;
 	void *record;
@@ -55,12 +57,9 @@ ssize_t zstd_compress_stream_to_records(struct zstd_data *data, void *dst, size_
 
 	while (input.pos < input.size) {
 		record = dst;
-		/* process_header writes the event header into record */
-		if (dst_size < sizeof(struct perf_event_header))
-			goto reset;
-		size = process_header(record, 0);
+		size = process_header(record, dst_size, 0);
 		/* Output buffer full — cannot fit even the record header */
-		if (size > dst_size)
+		if (size < 0)
 			goto reset;
 		compressed += size;
 		dst += size;
@@ -74,17 +73,21 @@ ssize_t zstd_compress_stream_to_records(struct zstd_data *data, void *dst, size_
 				(long)src_size, ZSTD_getErrorName(ret));
 			goto reset;
 		}
-		size = output.pos;
+		compressed += output.pos;
+		dst += output.pos;
+		dst_size -= output.pos;
 		/*
 		 * No progress: ZSTD couldn't emit any bytes into the
 		 * remaining output buffer.  Calling process_header
-		 * with size=0 would re-trigger header initialization,
+		 * with output.pos=0 would re-trigger header initialization,
 		 * double-subtracting the header size from dst_size and
 		 * underflowing the unsigned counter.
 		 */
-		if (size == 0)
+		if (output.pos == 0)
+			goto reset;
+		size = process_header(record, dst_size, output.pos);
+		if (size < 0)
 			goto reset;
-		size = process_header(record, size);
 		compressed += size;
 		dst += size;
 		dst_size -= size;