3 files changed, 141 insertions, 1 deletions

diff --git a/tools/perf/builtin-inject.c b/tools/perf/builtin-inject.c
index 5b50a4abf95f..71b6eafe4c19 100644
--- a/tools/perf/builtin-inject.c
+++ b/tools/perf/builtin-inject.c
@@ -27,6 +27,8 @@
 #include "util/namespaces.h"
 #include "util/util.h"
 
+#include <internal/lib.h>
+
 #include <linux/err.h>
 #include <subcmd/parse-options.h>
 #include <uapi/linux/mman.h> /* To get things like MAP_HUGETLB even on older libc headers */
@@ -55,6 +57,7 @@ struct perf_inject {
 	struct list_head	samples;
 	struct itrace_synth_opts itrace_synth_opts;
 	char			event_copy[PERF_SAMPLE_MAX_SIZE];
+	struct perf_file_section secs[HEADER_FEAT_BITS];
 };
 
 struct event_entry {
@@ -763,6 +766,120 @@ static int parse_vm_time_correlation(const struct option *opt, const char *str,
 	return inject->itrace_synth_opts.vm_tm_corr_args ? 0 : -ENOMEM;
 }
 
+static int save_section_info_cb(struct perf_file_section *section,
+				struct perf_header *ph __maybe_unused,
+				int feat, int fd __maybe_unused, void *data)
+{
+	struct perf_inject *inject = data;
+
+	inject->secs[feat] = *section;
+	return 0;
+}
+
+static int save_section_info(struct perf_inject *inject)
+{
+	struct perf_header *header = &inject->session->header;
+	int fd = perf_data__fd(inject->session->data);
+
+	return perf_header__process_sections(header, fd, inject, save_section_info_cb);
+}
+
+static bool keep_feat(int feat)
+{
+	switch (feat) {
+	/* Keep original information that describes the machine or software */
+	case HEADER_TRACING_DATA:
+	case HEADER_HOSTNAME:
+	case HEADER_OSRELEASE:
+	case HEADER_VERSION:
+	case HEADER_ARCH:
+	case HEADER_NRCPUS:
+	case HEADER_CPUDESC:
+	case HEADER_CPUID:
+	case HEADER_TOTAL_MEM:
+	case HEADER_CPU_TOPOLOGY:
+	case HEADER_NUMA_TOPOLOGY:
+	case HEADER_PMU_MAPPINGS:
+	case HEADER_CACHE:
+	case HEADER_MEM_TOPOLOGY:
+	case HEADER_CLOCKID:
+	case HEADER_BPF_PROG_INFO:
+	case HEADER_BPF_BTF:
+	case HEADER_CPU_PMU_CAPS:
+	case HEADER_CLOCK_DATA:
+	case HEADER_HYBRID_TOPOLOGY:
+	case HEADER_HYBRID_CPU_PMU_CAPS:
+		return true;
+	/* Information that can be updated */
+	case HEADER_BUILD_ID:
+	case HEADER_CMDLINE:
+	case HEADER_EVENT_DESC:
+	case HEADER_BRANCH_STACK:
+	case HEADER_GROUP_DESC:
+	case HEADER_AUXTRACE:
+	case HEADER_STAT:
+	case HEADER_SAMPLE_TIME:
+	case HEADER_DIR_FORMAT:
+	case HEADER_COMPRESSED:
+	default:
+		return false;
+	};
+}
+
+static int read_file(int fd, u64 offs, void *buf, size_t sz)
+{
+	ssize_t ret = preadn(fd, buf, sz, offs);
+
+	if (ret < 0)
+		return -errno;
+	if ((size_t)ret != sz)
+		return -EINVAL;
+	return 0;
+}
+
+static int feat_copy(struct perf_inject *inject, int feat, struct feat_writer *fw)
+{
+	int fd = perf_data__fd(inject->session->data);
+	u64 offs = inject->secs[feat].offset;
+	size_t sz = inject->secs[feat].size;
+	void *buf = malloc(sz);
+	int ret;
+
+	if (!buf)
+		return -ENOMEM;
+
+	ret = read_file(fd, offs, buf, sz);
+	if (ret)
+		goto out_free;
+
+	ret = fw->write(fw, buf, sz);
+out_free:
+	free(buf);
+	return ret;
+}
+
+struct inject_fc {
+	struct feat_copier fc;
+	struct perf_inject *inject;
+};
+
+static int feat_copy_cb(struct feat_copier *fc, int feat, struct feat_writer *fw)
+{
+	struct inject_fc *inj_fc = container_of(fc, struct inject_fc, fc);
+	struct perf_inject *inject = inj_fc->inject;
+	int ret;
+
+	if (!inject->secs[feat].offset ||
+	    !keep_feat(feat))
+		return 0;
+
+	ret = feat_copy(inject, feat, fw);
+	if (ret < 0)
+		return ret;
+
+	return 1; /* Feature section copied */
+}
+
 static int output_fd(struct perf_inject *inject)
 {
 	return inject->in_place_update ? -1 : perf_data__fd(&inject->output);
@@ -848,6 +965,11 @@ static int __cmd_inject(struct perf_inject *inject)
 		return ret;
 
 	if (!inject->is_pipe && !inject->in_place_update) {
+		struct inject_fc inj_fc = {
+			.fc.copy = feat_copy_cb,
+			.inject = inject,
+		};
+
 		if (inject->build_ids)
 			perf_header__set_feat(&session->header,
 					      HEADER_BUILD_ID);
@@ -872,7 +994,7 @@ static int __cmd_inject(struct perf_inject *inject)
 		}
 		session->header.data_offset = output_data_offset;
 		session->header.data_size = inject->bytes_written;
-		perf_session__write_header(session, session->evlist, fd, true);
+		perf_session__inject_header(session, session->evlist, fd, &inj_fc.fc);
 	}
 
 	return ret;
@@ -1037,6 +1159,11 @@ int cmd_inject(int argc, const char **argv)
 	if (zstd_init(&(inject.session->zstd_data), 0) < 0)
 		pr_warning("Decompression initialization failed.\n");
 
+	/* Save original section info before feature bits change */
+	ret = save_section_info(&inject);
+	if (ret)
+		goto out_delete;
+
 	if (!data.is_pipe && inject.output.is_pipe) {
 		ret = perf_header__write_pipe(perf_data__fd(&inject.output));
 		if (ret < 0) {
diff --git a/tools/perf/util/header.c b/tools/perf/util/header.c
index b0c57a130d1e..53332da100e8 100644
--- a/tools/perf/util/header.c
+++ b/tools/perf/util/header.c
@@ -3686,6 +3686,14 @@ int perf_session__write_header(struct perf_session *session,
 	return perf_session__do_write_header(session, evlist, fd, at_exit, NULL);
 }
 
+int perf_session__inject_header(struct perf_session *session,
+				struct evlist *evlist,
+				int fd,
+				struct feat_copier *fc)
+{
+	return perf_session__do_write_header(session, evlist, fd, true, fc);
+}
+
 static int perf_header__getbuffer64(struct perf_header *header,
 				    int fd, void *buf, size_t size)
 {
diff --git a/tools/perf/util/header.h b/tools/perf/util/header.h
index e76ab02d5541..08563c1f1bff 100644
--- a/tools/perf/util/header.h
+++ b/tools/perf/util/header.h
@@ -131,6 +131,11 @@ struct feat_copier {
 	int (*copy)(struct feat_copier *fc, int feat, struct feat_writer *fw);
 };
 
+int perf_session__inject_header(struct perf_session *session,
+				struct evlist *evlist,
+				int fd,
+				struct feat_copier *fc);
+
 void perf_header__set_feat(struct perf_header *header, int feat);
 void perf_header__clear_feat(struct perf_header *header, int feat);
 bool perf_header__has_feat(const struct perf_header *header, int feat);