From 5fa4151e925c39c4d788918badec2ec7b4cc8938 Mon Sep 17 00:00:00 2001
From: Mark Johnston <markj@FreeBSD.org>
Date: Thu, 8 Feb 2024 11:01:55 -0500
Subject: [PATCH] arm64: Implement busdma bits for KMSAN

This works identically to amd64.  In particular, only the
bus_dma_bounce_impl busdma implementation handles KMSAN at the moment.

MFC after:	2 weeks
Sponsored by:	Klara, Inc.
Sponsored by:	Juniper Networks, Inc.
Differential Revision:	https://reviews.freebsd.org/D43157
---
 sys/arm64/arm64/busdma_bounce.c  | 23 +++++++++++++++++++++--
 sys/arm64/include/bus_dma.h      | 12 ++++++++++++
 sys/arm64/include/bus_dma_impl.h |  3 +++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/sys/arm64/arm64/busdma_bounce.c b/sys/arm64/arm64/busdma_bounce.c
index a117e1041658..57551a2edb47 100644
--- a/sys/arm64/arm64/busdma_bounce.c
+++ b/sys/arm64/arm64/busdma_bounce.c
@@ -40,9 +40,10 @@
 #include <sys/kernel.h>
 #include <sys/ktr.h>
 #include <sys/lock.h>
-#include <sys/proc.h>
 #include <sys/memdesc.h>
+#include <sys/msan.h>
 #include <sys/mutex.h>
+#include <sys/proc.h>
 #include <sys/sysctl.h>
 #include <sys/uio.h>
 
@@ -104,6 +105,9 @@ struct bus_dmamap {
 #define	DMAMAP_FROM_DMAMEM	(1 << 1)
 #define	DMAMAP_MBUF		(1 << 2)
 	int			sync_count;
+#ifdef KMSAN
+	struct memdesc	       kmsan_mem;
+#endif
 	struct sync_list	slist[];
 };
 
@@ -1149,8 +1153,20 @@ bounce_bus_dmamap_sync(bus_dma_tag_t dmat, bus_dmamap_t map,
 		 */
 		dsb(sy);
 	}
+
+	kmsan_bus_dmamap_sync(&map->kmsan_mem, op);
 }
 
+#ifdef KMSAN
+static void
+bounce_bus_dmamap_load_kmsan(bus_dmamap_t map, struct memdesc *mem)
+{
+	if (map == NULL)
+		return;
+	memcpy(&map->kmsan_mem, mem, sizeof(map->kmsan_mem));
+}
+#endif
+
 struct bus_dma_impl bus_dma_bounce_impl = {
 	.tag_create = bounce_bus_dma_tag_create,
 	.tag_destroy = bounce_bus_dma_tag_destroy,
@@ -1166,5 +1182,8 @@ struct bus_dma_impl bus_dma_bounce_impl = {
 	.map_waitok = bounce_bus_dmamap_waitok,
 	.map_complete = bounce_bus_dmamap_complete,
 	.map_unload = bounce_bus_dmamap_unload,
-	.map_sync = bounce_bus_dmamap_sync
+	.map_sync = bounce_bus_dmamap_sync,
+#ifdef KMSAN
+	.load_kmsan = bounce_bus_dmamap_load_kmsan,
+#endif
 };
diff --git a/sys/arm64/include/bus_dma.h b/sys/arm64/include/bus_dma.h
index 9e5a7124dcf5..fb494be0bdc4 100644
--- a/sys/arm64/include/bus_dma.h
+++ b/sys/arm64/include/bus_dma.h
@@ -149,4 +149,16 @@ _bus_dmamap_complete(bus_dma_tag_t dmat, bus_dmamap_t map,
 	return (tc->impl->map_complete(dmat, map, segs, nsegs, error));
 }
 
+#ifdef KMSAN
+static inline void
+_bus_dmamap_load_kmsan(bus_dma_tag_t dmat, bus_dmamap_t map,
+    struct memdesc *mem)
+{
+	struct bus_dma_tag_common *tc;
+
+	tc = (struct bus_dma_tag_common *)dmat;
+	return (tc->impl->load_kmsan(map, mem));
+}
+#endif
+
 #endif /* !_MACHINE_BUS_DMA_H_ */
diff --git a/sys/arm64/include/bus_dma_impl.h b/sys/arm64/include/bus_dma_impl.h
index 9e5741758ef5..1fb5e10c92b0 100644
--- a/sys/arm64/include/bus_dma_impl.h
+++ b/sys/arm64/include/bus_dma_impl.h
@@ -75,6 +75,9 @@ struct bus_dma_impl {
 	void (*map_unload)(bus_dma_tag_t dmat, bus_dmamap_t map);
 	void (*map_sync)(bus_dma_tag_t dmat, bus_dmamap_t map,
 	    bus_dmasync_op_t op);
+#ifdef KMSAN
+	void (*load_kmsan)(bus_dmamap_t map, struct memdesc *mem);
+#endif
 };
 
 int common_bus_dma_tag_create(struct bus_dma_tag_common *parent,