kafl/tdx[v2]: implement atomic cache for virtio

When fuzzing TDX, the endianness conversions from the bounce buffer
results in unique values every time, which is impossible since the
buffer is copied and not modified until invalidated later. IE a read
at offset X within the buffer should always yield the same value but
currently results in a new random fuzz value.

To correct this, implement a cache for endianness conversions from the
DMA bounce buffer based on the original value used as in index into a
buffer of random fuzz data. Additionally, allow for index
granularities into this array based on data type. For example, if the
buffer is 512 bytes, it would support 64 unique u64s, 128 u32s and
256 u16s.

Implement this cache within the virtio_device struct so each
virtio_device gets its own cache and use an atomic. This avoids the need
to share a global cache and any associated locking.

Signed-off-by: William Roberts <william.c.roberts@intel.com>

arch/x86/include/asm/tdx.h

@@ -135,13 +135,25 @@ struct virtio_device;
 u64 tdx_fuzz(u64 var, uintptr_t addr, int size, enum tdx_fuzz_loc loc);
 bool tdx_fuzz_err(enum tdx_fuzz_loc loc);
 void tdx_fuzz_virtio_cache_init(struct virtio_device *vdev);
-u64 tdx_fuzz_virtio_cache_get_64(struct virtio_device *vdev, u64 orig_var);
+u16 tdx_fuzz_virtio_cache_get_u16(struct virtio_device *vdev, u16 orig_var);
+u32 tdx_fuzz_virtio_cache_get_u32(struct virtio_device *vdev, u32 orig_var);
+u64 tdx_fuzz_virtio_cache_get_u64(struct virtio_device *vdev, u64 orig_var);
 void tdx_fuzz_virtio_cache_refresh(struct device *dev);
 #else
 static inline u64 tdx_fuzz(u64 var, uintptr_t addr, int size, enum tdx_fuzz_loc loc) { return var; };
 static inline bool tdx_fuzz_err(enum tdx_fuzz_loc loc) { return false; }
 static inline void tdx_fuzz_virtio_cache_init(struct virtio_device *vdev) { }
-static inline u64 tdx_fuzz_virtio_cache_get_64(struct virtio_device *vdev, u64 orig_var)
+static inline u16 tdx_fuzz_virtio_cache_get_u16(struct virtio_device *vdev, u16 orig_var)
+{
+	return orig_var;
+}
+
+static inline u32 tdx_fuzz_virtio_cache_get_u32(struct virtio_device *vdev, u32 orig_var)
+{
+	return orig_var;
+}
+
+static inline u64 tdx_fuzz_virtio_cache_get_u64(struct virtio_device *vdev, u64 orig_var)
 {
 	return orig_var;
 }

arch/x86/kernel/kafl-agent.c

@@ -32,6 +32,7 @@
 #undef pr_fmt
 #define pr_fmt(fmt) "kAFL: " fmt
 
+#define ARRAY_LEN(X) (sizeof (X) / sizeof (*(X)))
 
 bool agent_initialized = false;
 bool fuzz_enabled = false;
@@ -497,7 +498,7 @@ size_t kafl_fuzz_buffer(void* fuzz_buf, const void *orig_buf,
 
 	num_fuzzed = _kafl_fuzz_buffer(fuzz_buf, num_bytes);
 
-	if (agent_flags.dump_observed) {
+	if (agent_flags.dump_observed && orig_buf) {
 		// record input seen/used on this execution
 		// with exit_at_eof=0, this should produce good seeds?
 		if (ob_pos + num_bytes > ob_num) {
@@ -535,22 +536,29 @@ EXPORT_SYMBOL(tdx_fuzz);
 
 void tdx_fuzz_virtio_cache_init(struct virtio_device *vdev)
 {
-	u64 data;
+	size_t num_of_bytes;
 
 	pr_debug("virtio fuzz cache: updating.\n");
-	data = tdx_fuzz(0, (uintptr_t)&data, sizeof(data), TDX_FUZZ_VIRTIO);
-	atomic64_set(&vdev->tdx.fuzz_data, data);
+
+	/* Original buffer context here doesn't make sense since we don't know where the reads will come from */
+	kafl_fuzz_buffer(vdev->tdx.fuzz_data, NULL, (uintptr_t)vdev->tdx.fuzz_data, sizeof(vdev->tdx.fuzz_data), TDX_FUZZ_VIRTIO);
 }
 EXPORT_SYMBOL(tdx_fuzz_virtio_cache_init);
 
-u64 tdx_fuzz_virtio_cache_get_64(struct virtio_device *vdev, u64 orig_var)
-{
-	/* orig_var needed for signature when fuzzing is disabled */
-	(void)orig_var;
-	pr_debug("virtio fuzz cache: get u64.\n");
-	return atomic64_read(&vdev->tdx.fuzz_data);
-}
-EXPORT_SYMBOL(tdx_fuzz_virtio_cache_get_64);
+#define xstr(s) str(s)
+#define str(s) #s
+#define VIRTIO_CACHE_TO_OFFSET(fuzz_data, orig_var) (orig_var % ((sizeof(fuzz_data[0])/sizeof(orig_var)) * ARRAY_LEN(fuzz_data)))
+#define VIRTIO_CACHE_GET(fuzz_data, orig_var) ((typeof(orig_var) *)fuzz_data)[VIRTIO_CACHE_TO_OFFSET(fuzz_data, orig_var)];
+#define DEFINE_VIRTIO_CACHE_GET_FN(dtype) \
+	dtype tdx_fuzz_virtio_cache_get_##dtype(struct virtio_device *vdev, dtype orig_var) \
+	{ \
+		return VIRTIO_CACHE_GET(vdev->tdx.fuzz_data, orig_var); \
+	} \
+	EXPORT_SYMBOL(tdx_fuzz_virtio_cache_get_##dtype)
+
+DEFINE_VIRTIO_CACHE_GET_FN(u16);
+DEFINE_VIRTIO_CACHE_GET_FN(u32);
+DEFINE_VIRTIO_CACHE_GET_FN(u64);
 
 void tdx_fuzz_virtio_cache_refresh(struct device *dev)
 {

include/linux/virtio.h

@@ -119,7 +119,7 @@ struct virtio_device {
 	u64 features;
 #ifdef CONFIG_TDX_FUZZ_KAFL_VIRTIO
 	struct {
-		atomic64_t fuzz_data;
+		u64 fuzz_data[256];
 	} tdx;
 #endif
 	void *priv;

include/linux/virtio_config.h

@@ -279,7 +279,7 @@ static inline u16 virtio16_to_cpu(struct virtio_device *vdev, __virtio16 val)
 {
 	u16 ret = __virtio16_to_cpu(virtio_is_little_endian(vdev), val);
 
-	return tdx_fuzz_virtio_cache_get_64(vdev, ret);
+	return tdx_fuzz_virtio_cache_get_u16(vdev, ret);
 }
 
 static inline __virtio16 cpu_to_virtio16(struct virtio_device *vdev, u16 val)
@@ -291,7 +291,7 @@ static inline u32 virtio32_to_cpu(struct virtio_device *vdev, __virtio32 val)
 {
 	u32 ret = __virtio32_to_cpu(virtio_is_little_endian(vdev), val);
 
-	return tdx_fuzz_virtio_cache_get_64(vdev, ret);
+	return tdx_fuzz_virtio_cache_get_u32(vdev, ret);
 }
 
 static inline __virtio32 cpu_to_virtio32(struct virtio_device *vdev, u32 val)
@@ -303,7 +303,7 @@ static inline u64 virtio64_to_cpu(struct virtio_device *vdev, __virtio64 val)
 {
 	u64 ret = __virtio64_to_cpu(virtio_is_little_endian(vdev), val);
 
-	return tdx_fuzz_virtio_cache_get_64(vdev, ret);
+	return tdx_fuzz_virtio_cache_get_u64(vdev, ret);
 }
 
 static inline __virtio64 cpu_to_virtio64(struct virtio_device *vdev, u64 val)